graphinvent/GraphGenerator.py
def _graph_to_mol(node_features : torch.Tensor,
                         edge_features : torch.Tensor,
                         n_nodes : int) -> rdkit.Chem.Mol:
            """
            Converts input graph represenetation (node and edge features) into
            an `rdkit.Mol` object.

            Args:
            ----
                node_features (torch.Tensor) : Node features tensor.
                edge_features (torch.Tensor) : Edge features tensor.
                n_nodes (int)                : Number of nodes in the graph representation.

            Returns:
            -------
                molecule (rdkit.Chem.Mol) : Molecule object.
            """

            # create empty editable `rdkit.Chem.Mol` object
            molecule = rdkit.Chem.RWMol()
            node_to_idx = {}

            # add atoms to editable mol object
            for node_idx in range(n_nodes):
                atom_to_add = _features_to_atom(node_idx, node_features)
                molecule_idx = molecule.AddAtom(atom_to_add)
                node_to_idx[node_idx] = molecule_idx

            # add bonds to atoms in editable mol object; to not add the same bond twice
            # (which leads to an error), mask half of the edge features beyond diagonal
            n_max_nodes = constants.dim_nodes[0]
            edge_mask = torch.triu(
                torch.ones((n_max_nodes, n_max_nodes), device=constants.device),
                diagonal=1
            )
            edge_mask = edge_mask.view(n_max_nodes, n_max_nodes, 1)
            edges_idc = torch.nonzero(edge_features * edge_mask)

            for node_idx1, node_idx2, bond_idx in edges_idc:
                molecule.AddBond(
                    node_to_idx[node_idx1.item()],
                    node_to_idx[node_idx2.item()],
                    constants.int_to_bondtype[bond_idx.item()],
                )

            try:  # convert editable mol object to non-editable mol object
                molecule.GetMol()
            except AttributeError:  # will throw an error if molecule is `None`
                pass

            if constants.ignore_H and molecule:
                try:  # correct for ignored Hs
                    rdkit.Chem.SanitizeMol(molecule)
                except ValueError:  # throws exception if molecule is too ugly to correct
                    pass

            return molecule

            Converts input graph represenetation (node and edge features) into
            an `rdkit.Mol` object.

            Args:
            ----
                node_features (torch.Tensor) : Node features tensor.
                edge_features (torch.Tensor) : Edge features tensor.
                n_nodes (int)                : Number of nodes in the graph representation.

            Returns:
            -------
                molecule (rdkit.Chem.Mol) : Molecule object.
            
Converts input graph represenetation (node and edge features) into
an `rdkit.Mol` object.
Converts input graph represenetation ( node and edge features ) into an ` rdkit . Mol ` object .
def _graph_to_mol(node_features : torch.Tensor,
                         edge_features : torch.Tensor,
                         n_nodes : int) -> rdkit.Chem.Mol:
            molecule = rdkit.Chem.RWMol()
            node_to_idx = {}
            for node_idx in range(n_nodes):
                atom_to_add = _features_to_atom(node_idx, node_features)
                molecule_idx = molecule.AddAtom(atom_to_add)
                node_to_idx[node_idx] = molecule_idx
            n_max_nodes = constants.dim_nodes[0]
            edge_mask = torch.triu(
                torch.ones((n_max_nodes, n_max_nodes), device=constants.device),
                diagonal=1
            )
            edge_mask = edge_mask.view(n_max_nodes, n_max_nodes, 1)
            edges_idc = torch.nonzero(edge_features * edge_mask)
            for node_idx1, node_idx2, bond_idx in edges_idc:
                molecule.AddBond(
                    node_to_idx[node_idx1.item()],
                    node_to_idx[node_idx2.item()],
                    constants.int_to_bondtype[bond_idx.item()],
                )
            try:  
                molecule.GetMol()
            except AttributeError:  
                pass
            if constants.ignore_H and molecule:
                try:  
                    rdkit.Chem.SanitizeMol(molecule)
                except ValueError:  
                    pass
            return molecule
def _graph_to_mol ( node_features : torch . Tensor , edge_features : torch . Tensor , n_nodes : int ) -> rdkit . Chem . Mol : molecule = rdkit . Chem . RWMol ( ) node_to_idx = { } for node_idx in range ( n_nodes ) : atom_to_add = _features_to_atom ( node_idx , node_features ) molecule_idx = molecule . AddAtom ( atom_to_add ) node_to_idx [ node_idx ] = molecule_idx n_max_nodes = constants . dim_nodes [ 0 ] edge_mask = torch . triu ( torch . ones ( ( n_max_nodes , n_max_nodes ) , device = constants . device ) , diagonal = 1 ) edge_mask = edge_mask . view ( n_max_nodes , n_max_nodes , 1 ) edges_idc = torch . nonzero ( edge_features * edge_mask ) for node_idx1 , node_idx2 , bond_idx in edges_idc : molecule . AddBond ( node_to_idx [ node_idx1 . item ( ) ] , node_to_idx [ node_idx2 . item ( ) ] , constants . int_to_bondtype [ bond_idx . item ( ) ] , ) try : molecule . GetMol ( ) except AttributeError : pass if constants . ignore_H and molecule : try : rdkit . Chem . SanitizeMol ( molecule ) except ValueError : pass return molecule
Converts input graph represenetation (node and edge features) into
an `rdkit.Mol` object.
Converts input graph represenetation ( node and edge features ) into an ` rdkit . Mol ` object .
"""
            Converts input graph represenetation (node and edge features) into
            an `rdkit.Mol` object.

            Args:
            ----
                node_features (torch.Tensor) : Node features tensor.
                edge_features (torch.Tensor) : Edge features tensor.
                n_nodes (int)                : Number of nodes in the graph representation.

            Returns:
            -------
                molecule (rdkit.Chem.Mol) : Molecule object.
            """ # create empty editable `rdkit.Chem.Mol` object # add atoms to editable mol object # add bonds to atoms in editable mol object; to not add the same bond twice # (which leads to an error), mask half of the edge features beyond diagonal # convert editable mol object to non-editable mol object # will throw an error if molecule is `None` # correct for ignored Hs # throws exception if molecule is too ugly to correct
molecule (rdkit.Chem.Mol) : Molecule object.
molecule ( rdkit . Chem . Mol ) : Molecule object .
node_features
torch.Tensor
Node features tensor.
Node features tensor .
edge_features
torch.Tensor
Edge features tensor.
Edge features tensor .
n_nodes
int
Number of nodes in the graph representation.
Number of nodes in the graph representation .