def get_conformer_rmsd(mol: RDKitMol) -> np.ndarray:
"""
Calculate conformer-conformer RMSD.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
rmsd: np.ndarray
A conformer-conformer RMSD value. The shape is `(NumConformers, NumConformers)`
"""
try:
from rdkit.Chem import AllChem
except ModuleNotFoundError:
raise ValueError("This function requires RDKit to be installed.")
rmsd = np.zeros((mol.GetNumConformers(), mol.GetNumConformers()),
dtype=float)
for i, ref_conf in enumerate(mol.GetConformers()):
for j, fit_conf in enumerate(mol.GetConformers()):
if i >= j:
continue
rmsd[i, j] = AllChem.GetBestRMS(mol, mol, ref_conf.GetId(),
fit_conf.GetId())
rmsd[j, i] = rmsd[i, j]
return rmsd
def mol_to_graph(mol: RDKitMol):
"""Convert RDKit Mol to NetworkX graph
Convert mol into a graph representation atoms are nodes, and bonds
are vertices stored as graph
Parameters
----------
mol: RDKit Mol
The molecule to convert into a graph.
Returns
-------
graph: networkx.Graph
Contains atoms indices as nodes, edges as bonds.
Note
----
This function requires NetworkX to be installed.
"""
try:
import networkx as nx
except ModuleNotFoundError:
raise ValueError("This function requires NetworkX to be installed.")
G = nx.Graph()
num_atoms = mol.GetNumAtoms()
G.add_nodes_from(range(num_atoms))
for i in range(mol.GetNumBonds()):
from_idx = mol.GetBonds()[i].GetBeginAtomIdx()
to_idx = mol.GetBonds()[i].GetEndAtomIdx()
G.add_edge(from_idx, to_idx)
return G
def generate_conformers(self, mol: RDKitMol) -> RDKitMol:
"""
Generate conformers for a molecule.
This function returns a copy of the original molecule with embedded
conformers.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
mol: rdkit.Chem.rdchem.Mol
A new RDKit Mol object containing the chosen conformers, sorted by
increasing energy.
"""
# initial embedding
mol = self.embed_molecule(mol)
if not mol.GetNumConformers():
msg = 'No conformers generated for molecule'
if mol.HasProp('_Name'):
name = mol.GetProp('_Name')
msg += ' "{}".'.format(name)
else:
msg += '.'
raise RuntimeError(msg)
# minimization and pruning
self.minimize_conformers(mol)
mol = self.prune_conformers(mol)
return mol
def prune_conformers(self, mol: RDKitMol) -> RDKitMol:
"""
Prune conformers from a molecule using an RMSD threshold, starting
with the lowest energy conformer.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
new_mol: rdkit.Chem.rdchem.Mol
A new rdkit.Chem.rdchem.Mol containing the chosen conformers, sorted by
increasing energy.
"""
try:
from rdkit import Chem
except ModuleNotFoundError:
raise ValueError("This function requires RDKit to be installed.")
if self.rmsd_threshold < 0 or mol.GetNumConformers() <= 1:
return mol
energies = self.get_conformer_energies(mol)
rmsd = self.get_conformer_rmsd(mol)
sort = np.argsort(energies) # sort by increasing energy
keep: List[float] = [] # always keep lowest-energy conformer
discard = []
for i in sort:
# always keep lowest-energy conformer
if len(keep) == 0:
keep.append(i)
continue
# discard conformers after max_conformers is reached
if len(keep) >= self.max_conformers:
discard.append(i)
continue
# get RMSD to selected conformers
this_rmsd = rmsd[i][np.asarray(keep, dtype=int)]
# discard conformers within the RMSD threshold
if np.all(this_rmsd >= self.rmsd_threshold):
keep.append(i)
else:
discard.append(i)
# create a new molecule to hold the chosen conformers
# this ensures proper conformer IDs and energy-based ordering
new_mol = Chem.Mol(mol)
new_mol.RemoveAllConformers()
conf_ids = [conf.GetId() for conf in mol.GetConformers()]
for i in keep:
conf = mol.GetConformer(conf_ids[i])
new_mol.AddConformer(conf, assignId=True)
return new_mol
def _featurize(self, mol: RDKitMol) -> GraphData:
"""Calculate molecule graph features from RDKit mol object.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphData
A molecule graph with some features.
"""
from rdkit import Chem
from rdkit.Chem import AllChem
# construct atom and bond features
try:
mol.GetAtomWithIdx(0).GetProp('_GasteigerCharge')
except:
# If partial charges were not computed
AllChem.ComputeGasteigerCharges(mol)
h_bond_infos = construct_hydrogen_bonding_info(mol)
sssr = Chem.GetSymmSSSR(mol)
# construct atom (node) feature
atom_features = np.array(
[
_construct_atom_feature(atom, h_bond_infos, sssr)
for atom in mol.GetAtoms()
],
dtype=np.float,
)
# construct edge (bond) information
src, dest, bond_features = [], [], []
for bond in mol.GetBonds():
# add edge list considering a directed graph
start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
src += [start, end]
dest += [end, start]
bond_features += 2 * [_construct_bond_feature(bond)]
if self.add_self_edges:
num_atoms = mol.GetNumAtoms()
src += [i for i in range(num_atoms)]
dest += [i for i in range(num_atoms)]
# add dummy edge features
bond_fea_length = len(bond_features[0])
bond_features += num_atoms * [[0 for _ in range(bond_fea_length)]]
return GraphData(node_features=atom_features,
edge_index=np.array([src, dest], dtype=np.int),
edge_features=np.array(bond_features, dtype=np.float))
def _featurize(self, mol: RDKitMol) -> GraphData:
"""Calculate molecule graph features from RDKit mol object.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphData
A molecule graph with some features.
"""
if self.use_partial_charge:
try:
mol.GetAtomWithIdx(0).GetProp('_GasteigerCharge')
except:
# If partial charges were not computed
try:
from rdkit.Chem import AllChem
AllChem.ComputeGasteigerCharges(mol)
except ModuleNotFoundError:
raise ImportError(
"This class requires RDKit to be installed.")
# construct atom (node) feature
h_bond_infos = construct_hydrogen_bonding_info(mol)
atom_features = np.asarray(
[
_construct_atom_feature(atom, h_bond_infos, self.use_chirality,
self.use_partial_charge)
for atom in mol.GetAtoms()
],
dtype=float,
)
# construct edge (bond) index
src, dest = [], []
for bond in mol.GetBonds():
# add edge list considering a directed graph
start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
src += [start, end]
dest += [end, start]
# construct edge (bond) feature
bond_features = None # deafult None
if self.use_edges:
features = []
for bond in mol.GetBonds():
features += 2 * [_construct_bond_feature(bond)]
bond_features = np.asarray(features, dtype=float)
return GraphData(node_features=atom_features,
edge_index=np.asarray([src, dest], dtype=int),
edge_features=bond_features)
def _featurize(self, datapoint: RDKitMol, **kwargs) -> np.ndarray:
"""Calculate atomic coordinates.
Parameters
----------
datapoint: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
np.ndarray
A numpy array of atomic coordinates. The shape is `(n_atoms, 3)`.
"""
try:
from rdkit import Chem
from rdkit.Chem import AllChem
except ModuleNotFoundError:
raise ImportError("This class requires RDKit to be installed.")
if 'mol' in kwargs:
datapoint = kwargs.get("mol")
raise DeprecationWarning(
'Mol is being phased out as a parameter, please pass "datapoint" instead.'
)
# Check whether num_confs >=1 or not
num_confs = len(datapoint.GetConformers())
if num_confs == 0:
datapoint = Chem.AddHs(datapoint)
AllChem.EmbedMolecule(datapoint, AllChem.ETKDG())
datapoint = Chem.RemoveHs(datapoint)
N = datapoint.GetNumAtoms()
coords = np.zeros((N, 3))
# RDKit stores atomic coordinates in Angstrom. Atomic unit of length is the
# bohr (1 bohr = 0.529177 Angstrom). Converting units makes gradient calculation
# consistent with most QM software packages.
if self.use_bohr:
coords_list = [
datapoint.GetConformer(0).GetAtomPosition(i).__idiv__(
0.52917721092) for i in range(N)
]
else:
coords_list = [
datapoint.GetConformer(0).GetAtomPosition(i) for i in range(N)
]
for atom in range(N):
coords[atom, 0] = coords_list[atom].x
coords[atom, 1] = coords_list[atom].y
coords[atom, 2] = coords_list[atom].z
return coords
def _featurize(self, datapoint: RDKitMol,
**kwargs) -> Optional[GraphMatrix]:
"""
Calculate adjacency matrix and nodes features for RDKitMol.
It strips any chirality and charges
Parameters
----------
datapoint: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphMatrix
A molecule graph with some features.
"""
try:
from rdkit import Chem
except ModuleNotFoundError:
raise ImportError("This method requires RDKit to be installed.")
if 'mol' in kwargs:
datapoint = kwargs.get("mol")
raise DeprecationWarning(
'Mol is being phased out as a parameter, please pass "datapoint" instead.'
)
if self.kekulize:
Chem.Kekulize(datapoint)
A = np.zeros(shape=(self.max_atom_count, self.max_atom_count),
dtype=np.float32)
bonds = datapoint.GetBonds()
begin, end = [b.GetBeginAtomIdx()
for b in bonds], [b.GetEndAtomIdx() for b in bonds]
bond_type = [self.bond_encoder[b.GetBondType()] for b in bonds]
A[begin, end] = bond_type
A[end, begin] = bond_type
degree = np.sum(A[:datapoint.GetNumAtoms(), :datapoint.GetNumAtoms()],
axis=-1)
X = np.array(
[
self.atom_encoder[atom.GetAtomicNum()]
for atom in datapoint.GetAtoms()
] + [0] * (self.max_atom_count - datapoint.GetNumAtoms()),
dtype=np.int32,
)
graph = GraphMatrix(A, X)
return graph if (degree > 0).all() else None
def _featurize(self, mol: RDKitMol) -> np.ndarray:
"""Calculate atomic coordinates.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
np.ndarray
A numpy array of atomic coordinates. The shape is `(n_atoms, 3)`.
"""
try:
from rdkit import Chem
from rdkit.Chem import AllChem
except ModuleNotFoundError:
raise ImportError("This class requires RDKit to be installed.")
# Check whether num_confs >=1 or not
num_confs = len(mol.GetConformers())
if num_confs == 0:
mol = Chem.AddHs(mol)
AllChem.EmbedMolecule(mol, AllChem.ETKDG())
mol = Chem.RemoveHs(mol)
N = mol.GetNumAtoms()
coords = np.zeros((N, 3))
# RDKit stores atomic coordinates in Angstrom. Atomic unit of length is the
# bohr (1 bohr = 0.529177 Angstrom). Converting units makes gradient calculation
# consistent with most QM software packages.
if self.use_bohr:
coords_list = [
mol.GetConformer(0).GetAtomPosition(i).__idiv__(0.52917721092)
for i in range(N)
]
else:
coords_list = [
mol.GetConformer(0).GetAtomPosition(i) for i in range(N)
]
for atom in range(N):
coords[atom, 0] = coords_list[atom].x
coords[atom, 1] = coords_list[atom].y
coords[atom, 2] = coords_list[atom].z
return coords
def _create_component_map(mol: RDKitMol,
components: List[List[int]]) -> Dict[int, int]:
"""Creates a map from atom ids to disconnected component id
For each atom in `mol`, maps it to the id of the component in the
molecule. The intent is that this is used on a molecule whose
rotatable bonds have been removed. `components` is a list of the
connected components after this surgery.
Parameters
----------
mol: RDKit Mol
molecule to find disconnected compontents in
components: List[List[int]]
List of connected components
Returns
-------
comp_map: Dict[int, int]
Maps atom ids to component ides
"""
comp_map = {}
for i in range(mol.GetNumAtoms()):
for j in range(len(components)):
if i in components[j]:
comp_map[i] = j
break
return comp_map
def _featurize(self, datapoint: RDKitMol, **kwargs) -> GraphData:
"""Calculate molecule graph features from RDKit mol object.
Parameters
----------
datapoint: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphData
A molecule graph with some features.
"""
if 'mol' in kwargs:
datapoint = kwargs.get("mol")
raise DeprecationWarning(
'Mol is being phased out as a parameter, please pass "datapoint" instead.'
)
node_features = np.asarray(
[self._pagtn_atom_featurizer(atom) for atom in datapoint.GetAtoms()],
dtype=np.float)
edge_index, edge_features = self._pagtn_edge_featurizer(datapoint)
graph = GraphData(node_features, edge_index, edge_features)
return graph
def convert_protein_to_pdbqt(mol: RDKitMol, outfile: str) -> None:
"""Convert a protein PDB file into a pdbqt file.
Writes the extra PDBQT terms directly to `outfile`.
Parameters
----------
mol: RDKit Mol
Protein molecule
outfile: str
filename which already has a valid pdb representation of mol
"""
lines = [x.strip() for x in open(outfile).readlines()]
out_lines = []
for line in lines:
if "ROOT" in line or "ENDROOT" in line or "TORSDOF" in line:
out_lines.append("%s\n" % line)
continue
if not line.startswith("ATOM"):
continue
line = line[:66]
atom_index = int(line[6:11])
atom = mol.GetAtoms()[atom_index - 1]
line = "%s +0.000 %s\n" % (line, atom.GetSymbol().ljust(2))
out_lines.append(line)
with open(outfile, 'w') as fout:
for line in out_lines:
fout.write(line)
def _featurize(self, datapoint: RDKitMol, **kwargs) -> np.ndarray:
"""Calculate symmetry function.
Parameters
----------
datapoint: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
np.ndarray
A numpy array of symmetry function. The shape is `(max_atoms, 4)`.
"""
if 'mol' in kwargs:
datapoint = kwargs.get("mol")
raise DeprecationWarning(
'Mol is being phased out as a parameter, please pass "datapoint" instead.'
)
coordinates = self.coordfeat._featurize(datapoint)
atom_numbers = np.array(
[atom.GetAtomicNum() for atom in datapoint.GetAtoms()])
atom_numbers = np.expand_dims(atom_numbers, axis=1)
assert atom_numbers.shape[0] == coordinates.shape[0]
features = np.concatenate([atom_numbers, coordinates], axis=1)
return pad_array(features, (self.max_atoms, 4))
def _featurize(self, mol: RDKitMol) -> Optional[GraphMatrix]:
"""
Calculate adjacency matrix and nodes features for RDKitMol.
It strips any chirality and charges
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphMatrix
A molecule graph with some features.
"""
try:
from rdkit import Chem
except ModuleNotFoundError:
raise ImportError("This method requires RDKit to be installed.")
if self.kekulize:
Chem.Kekulize(mol)
A = np.zeros(shape=(self.max_atom_count, self.max_atom_count),
dtype=np.float32)
bonds = mol.GetBonds()
begin, end = [b.GetBeginAtomIdx()
for b in bonds], [b.GetEndAtomIdx() for b in bonds]
bond_type = [self.bond_encoder[b.GetBondType()] for b in bonds]
A[begin, end] = bond_type
A[end, begin] = bond_type
degree = np.sum(A[:mol.GetNumAtoms(), :mol.GetNumAtoms()], axis=-1)
X = np.array(
[
self.atom_encoder[atom.GetAtomicNum()]
for atom in mol.GetAtoms()
] + [0] * (self.max_atom_count - mol.GetNumAtoms()),
dtype=np.int32,
)
graph = GraphMatrix(A, X)
return graph if (degree > 0).all() else None
def max_pair_distance_pairs(mol: RDKitMol,
max_pair_distance: Optional[int]) -> np.ndarray:
"""Helper method which finds atom pairs within max_pair_distance graph distance.
This helper method is used to find atoms which are within max_pair_distance
graph_distance of one another. This is done by using the fact that the
powers of an adjacency matrix encode path connectivity information. In
particular, if `adj` is the adjacency matrix, then `adj**k` has a nonzero
value at `(i, j)` if and only if there exists a path of graph distance `k`
between `i` and `j`. To find all atoms within `max_pair_distance` of each
other, we can compute the adjacency matrix powers `[adj, adj**2,
...,adj**max_pair_distance]` and find pairs which are nonzero in any of
these matrices. Since adjacency matrices and their powers are positive
numbers, this is simply the nonzero elements of `adj + adj**2 + ... +
adj**max_pair_distance`.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit molecules
max_pair_distance: Optional[int], (default None)
This value can be a positive integer or None. This
parameter determines the maximum graph distance at which pair
features are computed. For example, if `max_pair_distance==2`,
then pair features are computed only for atoms at most graph
distance 2 apart. If `max_pair_distance` is `None`, all pairs are
considered (effectively infinite `max_pair_distance`)
Returns
-------
np.ndarray
Of shape `(2, num_pairs)` where `num_pairs` is the total number of pairs
within `max_pair_distance` of one another.
"""
from rdkit import Chem
from rdkit.Chem import rdmolops
N = len(mol.GetAtoms())
if (max_pair_distance is None or max_pair_distance >= N):
max_distance = N
elif max_pair_distance is not None and max_pair_distance <= 0:
raise ValueError(
"max_pair_distance must either be a positive integer or None")
elif max_pair_distance is not None:
max_distance = max_pair_distance
adj = rdmolops.GetAdjacencyMatrix(mol)
# Handle edge case of self-pairs (i, i)
sum_adj = np.eye(N)
for i in range(max_distance):
# Increment by 1 since we don't want 0-indexing
power = i + 1
sum_adj += np.linalg.matrix_power(adj, power)
nonzero_locs = np.where(sum_adj != 0)
num_pairs = len(nonzero_locs[0])
# This creates a matrix of shape (2, num_pairs)
pair_edges = np.reshape(np.array(list(zip(nonzero_locs))), (2, num_pairs))
return pair_edges
def coulomb_matrix(self, mol: RDKitMol) -> np.ndarray:
"""
Generate Coulomb matrices for each conformer of the given molecule.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
np.ndarray
The coulomb matrices of the given molecule
"""
try:
from rdkit import Chem
from rdkit.Chem import AllChem
except ModuleNotFoundError:
raise ImportError("This class requires RDKit to be installed.")
# Check whether num_confs >=1 or not
num_confs = len(mol.GetConformers())
if num_confs == 0:
mol = Chem.AddHs(mol)
AllChem.EmbedMolecule(mol, AllChem.ETKDG())
if self.remove_hydrogens:
mol = Chem.RemoveHs(mol)
n_atoms = mol.GetNumAtoms()
z = [atom.GetAtomicNum() for atom in mol.GetAtoms()]
rval = []
for conf in mol.GetConformers():
d = self.get_interatomic_distances(conf)
m = np.outer(z, z) / d
m[range(n_atoms), range(n_atoms)] = 0.5 * np.array(z)**2.4
if self.randomize:
for random_m in self.randomize_coulomb_matrix(m):
random_m = pad_array(random_m, self.max_atoms)
rval.append(random_m)
else:
m = pad_array(m, self.max_atoms)
rval.append(m)
rval = np.asarray(rval)
return rval
def _edge_features(self, mol: RDKitMol, path_atoms: Tuple[int, ...],
ring_info) -> np.ndarray:
"""Computes the edge features for a given pair of nodes.
Parameters
----------
mol : : RDKitMol
RDKit molecule instance.
path_atoms: tuple
Shortest path between the given pair of nodes.
ring_info: list
Different rings that contain the pair of atoms
"""
features = []
path_bonds = []
path_length = len(path_atoms)
for path_idx in range(path_length - 1):
bond = mol.GetBondBetweenAtoms(path_atoms[path_idx],
path_atoms[path_idx + 1])
if bond is None:
import warnings
warnings.warn('Valid idx of bonds must be passed')
path_bonds.append(bond)
for path_idx in range(self.max_length):
if path_idx < len(path_bonds):
bond_type = get_bond_type_one_hot(path_bonds[path_idx])
conjugacy = get_bond_is_conjugated_one_hot(
path_bonds[path_idx])
ring_attach = get_bond_is_in_same_ring_one_hot(
path_bonds[path_idx])
features.append(
np.concatenate([bond_type, conjugacy, ring_attach]))
else:
features.append(np.zeros(6))
if path_length + 1 > self.max_length:
path_length = self.max_length + 1
position_feature = np.zeros(self.max_length + 2)
position_feature[path_length] = 1
features.append(position_feature)
if ring_info:
rfeat = [
one_hot_encode(r, allowable_set=self.RING_TYPES)
for r in ring_info
]
# The 1.0 float value represents True Boolean
rfeat = [1.0] + np.any(rfeat, axis=0).tolist()
features.append(rfeat)
else:
# This will return a boolean vector with all entries False
features.append(
[0.0] +
one_hot_encode(ring_info, allowable_set=self.RING_TYPES))
return np.concatenate(features, axis=0)
def minimize_conformers(self, mol: RDKitMol) -> None:
"""
Minimize molecule conformers.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object with embedded conformers.
"""
for conf in mol.GetConformers():
ff = self.get_molecule_force_field(mol, conf_id=conf.GetId())
ff.Minimize()
def _featurize(self, mol: RDKitMol) -> GraphMatrix:
"""Calculate adjacency matrix and nodes features for RDKitMol.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphMatrix
A molecule graph with some features.
"""
if self.kekulize:
Chem.Kekulize(mol)
A = np.zeros(shape=(self.max_atom_count, self.max_atom_count),
dtype=np.float32)
bonds = mol.GetBonds()
begin, end = [b.GetBeginAtomIdx()
for b in bonds], [b.GetEndAtomIdx() for b in bonds]
bond_type = [self.bond_encoder[b.GetBondType()] for b in bonds]
A[begin, end] = bond_type
A[end, begin] = bond_type
degree = np.sum(A[:mol.GetNumAtoms(), :mol.GetNumAtoms()], axis=-1)
X = np.array(
[
self.atom_encoder[atom.GetAtomicNum()]
for atom in mol.GetAtoms()
] + [0] * (self.max_atom_count - mol.GetNumAtoms()),
dtype=np.int32,
)
graph = GraphMatrix(A, X)
return graph if (degree > 0).all() else None
def compute_all_ecfp(mol: RDKitMol,
indices: Optional[Set[int]] = None,
degree: int = 2) -> Dict[int, str]:
"""Obtain molecular fragment for all atoms emanating outward to given degree.
For each fragment, compute SMILES string (for now) and hash to
an int. Return a dictionary mapping atom index to hashed
SMILES.
Parameters
----------
mol: rdkit Molecule
Molecule to compute ecfp fragments on
indices: Optional[Set[int]]
List of atom indices for molecule. Default is all indices. If
specified will only compute fragments for specified atoms.
degree: int
Graph degree to use when computing ECFP fingerprints
Returns
----------
dict
Dictionary mapping atom index to hashed smiles.
"""
ecfp_dict = {}
from rdkit import Chem
for i in range(mol.GetNumAtoms()):
if indices is not None and i not in indices:
continue
env = Chem.FindAtomEnvironmentOfRadiusN(mol, degree, i, useHs=True)
submol = Chem.PathToSubmol(mol, env)
smile = Chem.MolToSmiles(submol)
ecfp_dict[i] = "%s,%s" % (mol.GetAtoms()[i].GetAtomicNum(), smile)
return ecfp_dict
def construct_node_features_matrix(self, mol: RDKitMol) -> np.ndarray:
"""
This function constructs a matrix of atom features for all atoms in a given molecule using the atom_features function.
Parameters
----------
mol: RDKitMol
RDKit Mol object.
Returns
----------
Atom_features: ndarray
Numpy array containing atom features.
"""
return np.array([self.atom_features(atom) for atom in mol.GetAtoms()])
def _featurize(self, mol: RDKitMol) -> np.ndarray:
"""Calculate symmetry function.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
np.ndarray
A numpy array of symmetry function. The shape is `(max_atoms, 4)`.
"""
coordinates = self.coordfeat._featurize(mol)
atom_numbers = np.array([atom.GetAtomicNum() for atom in mol.GetAtoms()])
atom_numbers = np.expand_dims(atom_numbers, axis=1)
assert atom_numbers.shape[0] == coordinates.shape[0]
features = np.concatenate([atom_numbers, coordinates], axis=1)
return pad_array(features, (self.max_atoms, 4))
def _featurize(self, mol: RDKitMol) -> GraphData:
"""Calculate molecule graph features from RDKit mol object.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphData
A molecule graph with some features.
"""
node_features = np.asarray(
[self._pagtn_atom_featurizer(atom) for atom in mol.GetAtoms()],
dtype=np.float)
edge_index, edge_features = self._pagtn_edge_featurizer(mol)
graph = GraphData(node_features, edge_index, edge_features)
return graph
def get_conformer_energies(self, mol: RDKitMol) -> np.ndarray:
"""
Calculate conformer energies.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object with embedded conformers.
Returns
-------
energies : np.ndarray
Minimized conformer energies.
"""
energies = []
for conf in mol.GetConformers():
ff = self.get_molecule_force_field(mol, conf_id=conf.GetId())
energy = ff.CalcEnergy()
energies.append(energy)
return np.asarray(energies, dtype=float)
def _pagtn_edge_featurizer(self,
mol: RDKitMol) -> Tuple[np.ndarray, np.ndarray]:
"""Calculate bond features from RDKit mol object.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
np.ndarray
Source and Destination node indexes of each bond.
np.ndarray
numpy vector of bond features.
"""
n_atoms = mol.GetNumAtoms()
# To get the shortest paths between two nodes.
paths_dict = compute_all_pairs_shortest_path(mol)
# To get info if two nodes belong to the same ring.
rings_dict = compute_pairwise_ring_info(mol)
# Featurizer
feats = []
src = []
dest = []
for i in range(n_atoms):
for j in range(n_atoms):
src.append(i)
dest.append(j)
if (i, j) not in paths_dict:
feats.append(np.zeros(7 * self.max_length + 7))
continue
ring_info = rings_dict.get(self.ordered_pair(i, j), [])
feats.append(
self._edge_features(mol, paths_dict[(i, j)], ring_info))
return np.array([src, dest], dtype=np.int), np.array(feats,
dtype=np.float)
def get_rotatable_bonds(mol: RDKitMol) -> List[Tuple[int, int]]:
"""
https://github.com/rdkit/rdkit/blob/f4529c910e546af590c56eba01f96e9015c269a6/Code/GraphMol/Descriptors/Lipinski.cpp#L107
Taken from rdkit source to find which bonds are rotatable store
rotatable bonds in (from_atom, to_atom)
Parameters
----------
mol: RDKit Mol
Ligand molecule
Returns
-------
rotatable_bonds: List[List[int, int]]
List of rotatable bonds in molecule
Note
----
This function requires RDKit to be installed.
"""
try:
from rdkit import Chem
from rdkit.Chem import rdmolops
except ModuleNotFoundError:
raise ValueError("This function requires RDKit to be installed.")
pattern = Chem.MolFromSmarts(
"[!$(*#*)&!D1&!$(C(F)(F)F)&!$(C(Cl)(Cl)Cl)&!$(C(Br)(Br)Br)&!$(C([CH3])("
"[CH3])[CH3])&!$([CD3](=[N,O,S])-!@[#7,O,S!D1])&!$([#7,O,S!D1]-!@[CD3]="
"[N,O,S])&!$([CD3](=[N+])-!@[#7!D1])&!$([#7!D1]-!@[CD3]=[N+])]-!@[!$(*#"
"*)&!D1&!$(C(F)(F)F)&!$(C(Cl)(Cl)Cl)&!$(C(Br)(Br)Br)&!$(C([CH3])([CH3])"
"[CH3])]")
rdmolops.FastFindRings(mol)
rotatable_bonds = mol.GetSubstructMatches(pattern)
return rotatable_bonds
def _featurize(self, mol: RDKitMol) -> np.ndarray:
"""Featurizes a single SMILE into an image.
Parameters
----------
mol: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
np.ndarray
A 3D array of image, the shape is `(img_size, img_size, 1)`.
If the length of SMILES is longer than `max_len`, this value is an empty array.
"""
from rdkit import Chem
from rdkit.Chem import AllChem
smile = Chem.MolToSmiles(mol)
if len(smile) > self.max_len:
return np.array([])
cmol = Chem.Mol(mol.ToBinary())
cmol.ComputeGasteigerCharges()
AllChem.Compute2DCoords(cmol)
atom_coords = cmol.GetConformer(0).GetPositions()
if self.img_spec == "std":
# Setup image
img = np.zeros((self.img_size, self.img_size, 1))
# Compute bond properties
bond_props = np.array(
[[2.0, bond.GetBeginAtomIdx(),
bond.GetEndAtomIdx()] for bond in mol.GetBonds()])
# Compute atom properties
atom_props = np.array([[atom.GetAtomicNum()]
for atom in cmol.GetAtoms()])
bond_props = bond_props.astype(np.float32)
atom_props = atom_props.astype(np.float32)
else:
# Setup image
img = np.zeros((self.img_size, self.img_size, 4))
# Compute bond properties
bond_props = np.array([[
bond.GetBondTypeAsDouble(),
bond.GetBeginAtomIdx(),
bond.GetEndAtomIdx()
] for bond in mol.GetBonds()])
# Compute atom properties
atom_props = np.array([[
atom.GetAtomicNum(),
atom.GetProp("_GasteigerCharge"),
atom.GetExplicitValence(),
atom.GetHybridization().real,
] for atom in cmol.GetAtoms()])
bond_props = bond_props.astype(np.float32)
atom_props = atom_props.astype(np.float32)
partial_charges = atom_props[:, 1]
if np.any(np.isnan(partial_charges)):
return np.array([])
frac = np.linspace(0, 1, int(1 / self.res * 2))
# Reshape done for proper broadcast
frac = frac.reshape(-1, 1, 1)
bond_begin_idxs = bond_props[:, 1].astype(int)
bond_end_idxs = bond_props[:, 2].astype(int)
# Reshapes, and axes manipulations to facilitate vector processing.
begin_coords = atom_coords[bond_begin_idxs]
begin_coords = np.expand_dims(begin_coords.T, axis=0)
end_coords = atom_coords[bond_end_idxs]
end_coords = np.expand_dims(end_coords.T, axis=0)
# Draw a line between the two atoms.
# The coordinates of this line, are indicated in line_coords
line_coords = frac * begin_coords + (1 - frac) * end_coords
# Turn the line coordinates into image positions
bond_line_idxs = np.ceil(
(line_coords[:, 0] + self.embed) / self.res).astype(int)
bond_line_idys = np.ceil(
(line_coords[:, 1] + self.embed) / self.res).astype(int)
# Set the bond line coordinates to the bond property used.
img[bond_line_idxs, bond_line_idys, 0] = bond_props[:, 0]
# Turn atomic coordinates into image positions
atom_idxs = np.round(
(atom_coords[:, 0] + self.embed) / self.res).astype(int)
atom_idys = np.round(
(atom_coords[:, 1] + self.embed) / self.res).astype(int)
# Set the atom positions in image to different atomic properties in channels
img[atom_idxs, atom_idys, :] = atom_props
return img
def pair_features(mol: RDKitMol,
bond_features_map: dict,
bond_adj_list: List,
bt_len: int = 6,
graph_distance: bool = True,
max_pair_distance: Optional[int] = None) -> np.ndarray:
"""Helper method used to compute atom pair feature vectors.
Many different featurization methods compute atom pair features
such as WeaveFeaturizer. Note that atom pair features could be
for pairs of atoms which aren't necessarily bonded to one
another.
Parameters
----------
mol: RDKit Mol
Molecule to compute features on.
bond_features_map: dict
Dictionary that maps pairs of atom ids (say `(2, 3)` for a bond between
atoms 2 and 3) to the features for the bond between them.
bond_adj_list: list of lists
`bond_adj_list[i]` is a list of the atom indices that atom `i` shares a
bond with . This list is symmetrical so if `j in bond_adj_list[i]` then `i
in bond_adj_list[j]`.
bt_len: int, optional (default 6)
The number of different bond types to consider.
graph_distance: bool, optional (default True)
If true, use graph distance between molecules. Else use euclidean
distance. The specified `mol` must have a conformer. Atomic
positions will be retrieved by calling `mol.getConformer(0)`.
max_pair_distance: Optional[int], (default None)
This value can be a positive integer or None. This
parameter determines the maximum graph distance at which pair
features are computed. For example, if `max_pair_distance==2`,
then pair features are computed only for atoms at most graph
distance 2 apart. If `max_pair_distance` is `None`, all pairs are
considered (effectively infinite `max_pair_distance`)
Note
----
This method requires RDKit to be installed.
Returns
-------
features: np.ndarray
Of shape `(N_edges, bt_len + max_distance + 1)`. This is the array
of pairwise features for all atom pairs, where N_edges is the
number of edges within max_pair_distance of one another in this
molecules.
pair_edges: np.ndarray
Of shape `(2, num_pairs)` where `num_pairs` is the total number of
pairs within `max_pair_distance` of one another.
"""
if graph_distance:
max_distance = 7
else:
max_distance = 1
N = mol.GetNumAtoms()
pair_edges = max_pair_distance_pairs(mol, max_pair_distance)
num_pairs = pair_edges.shape[1]
N_edges = pair_edges.shape[1]
features = np.zeros((N_edges, bt_len + max_distance + 1))
# Get mapping
mapping = {}
for n in range(N_edges):
a1, a2 = pair_edges[:, n]
mapping[(int(a1), int(a2))] = n
num_atoms = mol.GetNumAtoms()
rings = mol.GetRingInfo().AtomRings()
for a1 in range(num_atoms):
for a2 in bond_adj_list[a1]:
# first `bt_len` features are bond features(if applicable)
if (int(a1), int(a2)) not in mapping:
raise ValueError(
"Malformed molecule with bonds not in specified graph distance.")
else:
n = mapping[(int(a1), int(a2))]
features[n, :bt_len] = np.asarray(
bond_features_map[tuple(sorted((a1, a2)))], dtype=float)
for ring in rings:
if a1 in ring:
for a2 in ring:
if (int(a1), int(a2)) not in mapping:
# For ring pairs outside max pairs distance continue
continue
else:
n = mapping[(int(a1), int(a2))]
# `bt_len`-th feature is if the pair of atoms are in the same ring
if a2 == a1:
features[n, bt_len] = 0
else:
features[n, bt_len] = 1
# graph distance between two atoms
if graph_distance:
# distance is a matrix of 1-hot encoded distances for all atoms
distance = find_distance(
a1, num_atoms, bond_adj_list, max_distance=max_distance)
for a2 in range(num_atoms):
if (int(a1), int(a2)) not in mapping:
# For ring pairs outside max pairs distance continue
continue
else:
n = mapping[(int(a1), int(a2))]
features[n, bt_len + 1:] = distance[a2]
# Euclidean distance between atoms
if not graph_distance:
coords = np.zeros((N, 3))
for atom in range(N):
pos = mol.GetConformer(0).GetAtomPosition(atom)
coords[atom, :] = [pos.x, pos.y, pos.z]
features[:, :, -1] = np.sqrt(np.sum(np.square(
np.stack([coords] * N, axis=1) - \
np.stack([coords] * N, axis=0)), axis=2))
return features, pair_edges
def _featurize(self, datapoint: RDKitMol, **kwargs) -> np.ndarray:
"""Featurizes a single SMILE into an image.
Parameters
----------
datapoint: rdkit.Chem.rdchem.Mol
RDKit Mol object
Returns
-------
np.ndarray
A 3D array of image, the shape is `(img_size, img_size, 1)`.
If the length of SMILES is longer than `max_len`, this value is an empty array.
"""
try:
from rdkit import Chem
from rdkit.Chem import AllChem
except ModuleNotFoundError:
raise ImportError("This class requires RDKit to be installed.")
if 'mol' in kwargs:
datapoint = kwargs.get("mol")
raise DeprecationWarning(
'Mol is being phased out as a parameter, please pass "datapoint" instead.'
)
smile = Chem.MolToSmiles(datapoint)
if len(smile) > self.max_len:
return np.array([])
cmol = Chem.Mol(datapoint.ToBinary())
cmol.ComputeGasteigerCharges()
AllChem.Compute2DCoords(cmol)
atom_coords = cmol.GetConformer(0).GetPositions()
if self.img_spec == "std":
# Setup image
img = np.zeros((self.img_size, self.img_size, 1))
# Compute bond properties
bond_props = np.array(
[[2.0, bond.GetBeginAtomIdx(),
bond.GetEndAtomIdx()] for bond in datapoint.GetBonds()])
# Compute atom properties
atom_props = np.array([[atom.GetAtomicNum()]
for atom in cmol.GetAtoms()])
bond_props = bond_props.astype(np.float32)
atom_props = atom_props.astype(np.float32)
else:
# Setup image
img = np.zeros((self.img_size, self.img_size, 4))
# Compute bond properties
bond_props = np.array([[
bond.GetBondTypeAsDouble(),
bond.GetBeginAtomIdx(),
bond.GetEndAtomIdx()
] for bond in datapoint.GetBonds()])
# Compute atom properties
atom_props = np.array([[
atom.GetAtomicNum(),
atom.GetProp("_GasteigerCharge"),
atom.GetExplicitValence(),
atom.GetHybridization().real,
] for atom in cmol.GetAtoms()])
bond_props = bond_props.astype(np.float32)
atom_props = atom_props.astype(np.float32)
partial_charges = atom_props[:, 1]
if np.any(np.isnan(partial_charges)):
return np.array([])
frac = np.linspace(0, 1, int(1 / self.res * 2))
# Reshape done for proper broadcast
frac = frac.reshape(-1, 1, 1)
bond_begin_idxs = bond_props[:, 1].astype(int)
bond_end_idxs = bond_props[:, 2].astype(int)
# Reshapes, and axes manipulations to facilitate vector processing.
begin_coords = atom_coords[bond_begin_idxs]
begin_coords = np.expand_dims(begin_coords.T, axis=0)
end_coords = atom_coords[bond_end_idxs]
end_coords = np.expand_dims(end_coords.T, axis=0)
# Draw a line between the two atoms.
# The coordinates of this line, are indicated in line_coords
line_coords = frac * begin_coords + (1 - frac) * end_coords
# Turn the line coordinates into image positions
bond_line_idxs = np.ceil(
(line_coords[:, 0] + self.embed) / self.res).astype(int)
bond_line_idys = np.ceil(
(line_coords[:, 1] + self.embed) / self.res).astype(int)
# Turn atomic coordinates into image positions
atom_idxs = np.round(
(atom_coords[:, 0] + self.embed) / self.res).astype(int)
atom_idys = np.round(
(atom_coords[:, 1] + self.embed) / self.res).astype(int)
try:
# Set the bond line coordinates to the bond property used.
img[bond_line_idxs, bond_line_idys, 0] = bond_props[:, 0]
# Set the atom positions in image to different atomic properties in channels
img[atom_idxs, atom_idys, :] = atom_props
except IndexError:
# With fixed res and img_size some molecules (e.g. long chains) may not fit.
raise IndexError(
"The molecule does not fit into the image. Consider increasing img_size or res of the SmilesToImage featurizer."
)
return img
def _featurize(self, datapoint: RDKitMol, **kwargs) -> GraphData:
"""Calculate molecule graph features from RDKit mol object.
Parameters
----------
datapoint: rdkit.Chem.rdchem.Mol
RDKit mol object.
Returns
-------
graph: GraphData
A molecule graph with some features.
"""
assert datapoint.GetNumAtoms(
) > 1, "More than one atom should be present in the molecule for this featurizer to work."
if 'mol' in kwargs:
datapoint = kwargs.get("mol")
raise DeprecationWarning(
'Mol is being phased out as a parameter, please pass "datapoint" instead.'
)
if self.use_partial_charge:
try:
datapoint.GetAtomWithIdx(0).GetProp('_GasteigerCharge')
except:
# If partial charges were not computed
try:
from rdkit.Chem import AllChem
AllChem.ComputeGasteigerCharges(datapoint)
except ModuleNotFoundError:
raise ImportError(
"This class requires RDKit to be installed.")
# construct atom (node) feature
h_bond_infos = construct_hydrogen_bonding_info(datapoint)
atom_features = np.asarray(
[
_construct_atom_feature(atom, h_bond_infos, self.use_chirality,
self.use_partial_charge)
for atom in datapoint.GetAtoms()
],
dtype=float,
)
# construct edge (bond) index
src, dest = [], []
for bond in datapoint.GetBonds():
# add edge list considering a directed graph
start, end = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
src += [start, end]
dest += [end, start]
# construct edge (bond) feature
bond_features = None # deafult None
if self.use_edges:
features = []
for bond in datapoint.GetBonds():
features += 2 * [_construct_bond_feature(bond)]
bond_features = np.asarray(features, dtype=float)
return GraphData(node_features=atom_features,
edge_index=np.asarray([src, dest], dtype=int),
edge_features=bond_features)