Skip to content

Commit

Permalink
Merge pull request #66 from xiaoruiDong/pep8
Browse files Browse the repository at this point in the history
PEP-8 Formatting
  • Loading branch information
xiaoruiDong committed Sep 13, 2023
2 parents 71c1e5a + a67822c commit 39a9380
Show file tree
Hide file tree
Showing 44 changed files with 509 additions and 454 deletions.
2 changes: 1 addition & 1 deletion rdmc/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
#-*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

from rdmc.conf import *
from rdmc.forcefield import *
Expand Down
67 changes: 34 additions & 33 deletions rdmc/conf.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
#-*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

"""
This module provides class and methods for dealing with RDKit Conformer.
Expand Down Expand Up @@ -88,7 +88,7 @@ def FromRDKitMol(cls,

def GetBondLength(self,
atomIds: Sequence[int],
) -> float:
) -> float:
"""
Get the bond length between atoms in Angstrom.
Expand Down Expand Up @@ -271,7 +271,7 @@ def SetPositions(self,
def SetBondLength(self,
atomIds: Sequence[int],
value: Union[int, float],
) -> float:
) -> float:
"""
Set the bond length between atoms in Angstrom.
Expand Down Expand Up @@ -311,7 +311,7 @@ def SetAngleDeg(self,
# RDKit doesn't allow change bonds for non-bonding atoms
# A workaround may be form a bond and change the distance
edit_conf_by_add_bonds(self, 'SetAngleDeg', atomIds, value)
except:
except BaseException:
# RDKit doesn't allow change bonds for atoms in a ring
# A workaround hasn't been proposed
raise NotImplementedError(f'Approach for modifying the bond angle of {atomIds} is not available.')
Expand All @@ -335,7 +335,7 @@ def SetAngleRad(self,
# RDKit doesn't allow change bonds for non-bonding atoms
# A workaround may be form a bond and change the distance
edit_conf_by_add_bonds(self, 'SetAngleRad', atomIds, value)
except:
except BaseException:
# RDKit doesn't allow change bonds for atoms in a ring
# A workaround hasn't been proposed
raise NotImplementedError(f'Approach for modifying the bond angle of {atomIds} is not available.')
Expand All @@ -358,7 +358,7 @@ def SetTorsionDeg(self,
# RDKit doesn't allow change bonds for non-bonding atoms
# A workaround may be form a bond and change the distance
edit_conf_by_add_bonds(self, 'SetDihedralDeg', torsion, degree)
except:
except BaseException:
# RDKit doesn't allow change bonds for atoms in a ring
# A workaround hasn't been proposed
raise NotImplementedError(f'Approach for modifying the dihedral of {torsion} is not available.')
Expand Down Expand Up @@ -448,7 +448,7 @@ def split_by_energies(self,
Args:
decimals (int, optional): clustering children based on the number of digit
after the dot of the energy values.
For kcal/mol and J/mol, 0 or 1 is recommended; for
For kcal/mol and J/mol, 0 or 1 is recommended; for
hartree, 3 or 4 is recommended. Defaults to ``1``.
as_dict (bool, optional): If ``True``, return a dict object whose keys are
energy values and values are divided ConformerCluster
Expand Down Expand Up @@ -488,7 +488,7 @@ def merge(self,

# Update the children and energies for the current cluster
self.children, self.energies = (np.concatenate([c.children for c in new_cluster_list]),
np.concatenate([c.energies for c in new_cluster_list]))
np.concatenate([c.energies for c in new_cluster_list]))

def _update_energy_and_head(self):
"""
Expand All @@ -513,7 +513,7 @@ def __init__(self,
def get_torsional_angles(self,
confid: int,
adjust_periodicity: bool = True,
) -> 'np.array':
) -> 'np.array':
"""
Get torsional angles for a given conformers.
Expand All @@ -533,9 +533,9 @@ def get_torsional_angles(self,
return np.array(tor_angle)

def get_tor_matrix(self,
confs: Union['np.array',list],
confs: Union['np.array', list],
adjust_periodicity: bool = True,
) -> np.ndarray:
) -> np.ndarray:
"""
Get the torsional matrix consists of torsional angles for the input list of conformer indexes.
Expand All @@ -554,8 +554,8 @@ def get_tor_matrix(self,
return np.array(tor_matrix)

def _get_results_from_cluster_idxs(self,
confs: Union[list,'np.array'],
c_idxs: Union[list,'np.array'],
confs: Union[list, 'np.array'],
c_idxs: Union[list, 'np.array'],
as_dict: bool,
as_list_idx: bool,):
"""
Expand Down Expand Up @@ -593,7 +593,7 @@ def hierarchy_cluster(self,
adjust_periodicity: bool = False,
as_dict: bool = True,
as_list_idx: bool = False,
):
):
"""
The implementation of an hierarchy clustering method based on scipy.
It is basically defining clusters based on points within a hypercube defined by threshold.
Expand Down Expand Up @@ -624,9 +624,9 @@ def hierarchy_cluster(self,
# The output is an array of the same size as 'confs', but with numbers indicating the belonging
# cluster indexes. E.g., [1,1,1,2,2,2]
c_idxs = hcluster.fclusterdata(tor_matrix,
threshold,
criterion=criterion,
method=method)
threshold,
criterion=criterion,
method=method)

return self._get_results_from_cluster_idxs(confs=confs, c_idxs=c_idxs,
as_dict=as_dict, as_list_idx=as_list_idx)
Expand Down Expand Up @@ -688,7 +688,7 @@ def filter_by_iter_hcluster(self,
method=method,
as_dict=False,
as_list_idx=True,
adjust_periodicity=i%2)
adjust_periodicity=i % 2)
clusters = self._get_clusters_from_grouped_idxs(clusters, grouped_conf_idx)
cur_num_clusters = len(clusters)
if cur_num_clusters == last_num_clusters:
Expand Down Expand Up @@ -733,9 +733,11 @@ def check_dihed_angle_diff(self,

if mask:
# Allowing mask some dimensions
masked_tor = lambda tor: np.ma.masked_array(tor, mask)
def masked_tor(tor):
return np.ma.masked_array(tor, mask)
else:
masked_tor = lambda tor: tor
def masked_tor(tor):
return tor

# Create an array to store cluster indexes
# Initializing all elements to -1
Expand Down Expand Up @@ -801,16 +803,16 @@ def filter_by_dihed_angles(self,
adjust_periodicity=True)
clusters = [ConformerCluster(clusters.children[c],
clusters.energies[c]) for c in c_tmp]
max_iter-=1
max_iter -= 1

for i in range(max_iter):
confs = [cl.head for cl in clusters]
grouped_conf_idx = self.check_dihed_angle_diff(confs,
threshold=threshold,
mask=mask,
as_dict=False,
as_list_idx=True,
adjust_periodicity=i%2)
threshold=threshold,
mask=mask,
as_dict=False,
as_list_idx=True,
adjust_periodicity=i % 2)

clusters = self._get_clusters_from_grouped_idxs(clusters, grouped_conf_idx)

Expand Down Expand Up @@ -844,10 +846,9 @@ def reset_atom_maps(self, max_atom_maps=100000):
maxMatches=max_atom_maps)
self._atom_maps = [list(enumerate(match)) for match in matches]
if len(self._atom_maps) == 100000:
print('WARNING: The atom index mappings are not complete (possibly due to the'
'large size of the molecule or high symmetry). You may want to regenerate'
'atom mappings by `reset_atom_maps` with a number larger than 100000.')

print('WARNING: The atom index mappings are not complete (possibly due to the'
'large size of the molecule or high symmetry). You may want to regenerate'
'atom mappings by `reset_atom_maps` with a number larger than 100000.')

def pairwise_rmsd(self,
i: int,
Expand All @@ -862,7 +863,7 @@ def pairwise_rmsd(self,
Args:
i (int): Index of one of the conformer. Usually, used as the 'reference'.
j (int): Index of the other conformer.
reflect (bool, optional): Whether to reflect the j conformer to rule out
reflect (bool, optional): Whether to reflect the j conformer to rule out
mirror symmetry. Defaults to ``False``.
reorder (bool, optional): Whether to allow atom index order to change (based
on isomorphism check to rule out torsional symmetry
Expand Down Expand Up @@ -914,8 +915,8 @@ def edit_conf_by_add_bonds(conf, function_name, atoms, value):
tmp_atoms = sorted(atoms)
bonds_to_add = []
for i in range(len(tmp_atoms) - 1):
if not (tmp_atoms[i], tmp_atoms[i+1]) in all_bonds:
bonds_to_add.append([tmp_atoms[i], tmp_atoms[i+1]])
if not (tmp_atoms[i], tmp_atoms[i + 1]) in all_bonds:
bonds_to_add.append([tmp_atoms[i], tmp_atoms[i + 1]])
tmp_mol = tmp_mol.AddRedundantBonds(bonds_to_add)
tmp_mol.SetPositions(conf.GetPositions())
tmp_conf = tmp_mol.GetConformer()
Expand Down
20 changes: 13 additions & 7 deletions rdmc/conformer_generation/align.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
#-*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

"""
This module provides class and methods for reaction path analysis.
Expand Down Expand Up @@ -129,7 +129,7 @@ def from_reactants(cls,
raise ValueError(f'The conf_ids\'s length (currently {len(conf_ids)}) should be '
f'the same as the length of moles (currently {len(mols)}.')
coord_list = [mol.GetPostions(id=conf_id) for mol, conf_id in zip(mols, conf_ids)]
coords = np.concatenate(coord_list, axis=0)
coords = np.concatenate(coord_list, axis=0)
atom_maps, counter = [], 0
for mol in mols:
atom_maps.append(list(range(counter, counter + mol.GetNumAtoms())))
Expand Down Expand Up @@ -222,7 +222,7 @@ def initialize_align(self,
np.zeros(3))
elif len(self.atom_maps) == 2:
pos = [np.zeros(3), np.array([self.dist, 0., 0.])]
for i in [0, 1] :
for i in [0, 1]:
atom_map = self.atom_maps[i]
# Make the first fragment centered at (0, 0, 0)
# Make the second fragment centered at (R1 + R2 + dist)
Expand Down Expand Up @@ -269,8 +269,8 @@ def score_bimolecule(self,
# Square euclideans distance is used as score for each bond.
score2 = 0.
for bond in self.formed_bonds:
# Only bonds form between fragments will help decrease this score3
# Bonds formed inside a fragment will not change this score since molecules are rigid and the distances are fixed
# Only bonds form between fragments will help decrease this score3
# Bonds formed inside a fragment will not change this score since molecules are rigid and the distances are fixed
score2 += np.sum((coords[bond[0], :] - coords[bond[1], :]) ** 2)
score2 = score2 / len(self.formed_bonds) / dist_ref

Expand All @@ -281,7 +281,12 @@ def score_bimolecule(self,
react_atom_center = [get_centroid(coords[self.reacting_atoms[i], :]) for i in range(2)]
for i in [0, 1]:
if len(self.non_reacting_atoms[i]):
score3 += np.sum(1 / distance.cdist(coords[self.non_reacting_atoms[i],:], react_atom_center[i-1].reshape(1, -1), 'sqeuclidean') ** 2)
score3 += np.sum(
1 / distance.cdist(
coords[self.non_reacting_atoms[i], :],
react_atom_center[i - 1].reshape(1, -1),
'sqeuclidean') ** 2
)

return score1 + score2 + score3

Expand Down Expand Up @@ -351,7 +356,8 @@ def reset_pmol(r_mol: 'RDKitMol',
broken_bonds = get_broken_bonds(r_mol, p_mol)
r_conf = r_mol.GetConformer()
current_distances = [r_conf.GetBondLength(b) for b in broken_bonds]
[obff.add_distance_constraint(b, 1.5*d) for b, d in zip(broken_bonds, current_distances)]
[obff.add_distance_constraint(b, 1.5 * d)
for b, d in zip(broken_bonds, current_distances)]
obff.optimize(max_step=2000)

# second minimization without constraints
Expand Down
6 changes: 5 additions & 1 deletion rdmc/conformer_generation/embedders.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
#-*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

"""
Modules for providing initial guess geometries
Expand Down Expand Up @@ -32,6 +32,7 @@ class ConfGenEmbedder:
"""
Base class for conformer generation embedders.
"""

def __init__(self, track_stats=False):

self.iter = 0
Expand Down Expand Up @@ -141,6 +142,7 @@ class GeoMolEmbedder(ConfGenEmbedder):
temp_schedule (str, optional): Temperature schedule. Defaults to ``"linear"``.
track_stats (bool, optional): Whether to track the statistics of the conformer generation. Defaults to ``False``.
"""

def __init__(self,
trained_model_dir: str,
dataset: str = "drugs",
Expand Down Expand Up @@ -198,6 +200,7 @@ def embed_conformers(self,
conf.SetPositions(x.squeeze(axis=1))
return self.mol


class ETKDGEmbedder(ConfGenEmbedder):
"""
Embed conformers using ETKDG.
Expand All @@ -216,6 +219,7 @@ def embed_conformers(self, n_conformers: int):
self.mol.EmbedMultipleConfs(n_conformers)
return self.mol


class RandomEmbedder(ConfGenEmbedder):
"""
Embed conformers with coordinates of random numbers.
Expand Down
19 changes: 11 additions & 8 deletions rdmc/conformer_generation/generators.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
#!/usr/bin/env python3
#-*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

"""
Modules for conformer generation workflows
Expand Down Expand Up @@ -49,6 +49,7 @@ class StochasticConformerGenerator:
max_iters (int, optional}: Maximum number of iterations for which to run the module.
final_modules (list): List of instances of optimizer/pruner to run after initial cycles complete.
"""

def __init__(self,
smiles,
embedder: Optional['ConfGenEmbedder'] = None,
Expand Down Expand Up @@ -204,12 +205,12 @@ def set_config(self,
optimizer (ConfGenOptimizer, optional): Instance of a :obj:`ConfGenOptimizer <rdmc.conformer_generation.optimizers.ConfGenOptimizer>`.
Available options are :obj:`XTBOptimizer <rdmc.conformer_generation.optimizers.XTBOptimizer>`,
:obj:`GaussianOptimizer <rdmc.conformer_generation.optimizers.GaussianOptimizer>`, and
:obj:`MMFFOptimizer <rdmc.conformer_generation.optimizers.MMFFOptimizer>`. Defaults to
:obj:`MMFFOptimizer <rdmc.conformer_generation.optimizers.MMFFOptimizer>`. Defaults to
:obj:`XTBOptimizer <rdmc.conformer_generation.optimizers.XTBOptimizer>` with ``"gff"`` method.
pruner (ConfGenPruner, optional): Instance of a :obj:`ConfGenPruner <rdmc.conformer_generation.pruners.ConfGenPruner>`. Available options are
:obj:`CRESTPruner <rdmc.conformer_generation.pruners.CRESTPruner>` and
:obj:`TorsionPruner <rdmc.conformer_generation.pruners.TorsionPruner>`. By default,
``"loose"`` utilizes :obj:`TorsionPruner <rdmc.conformer_generation.pruners.TorsionPruner>` with
``"loose"`` utilizes :obj:`TorsionPruner <rdmc.conformer_generation.pruners.TorsionPruner>` with
``mean_chk_threshold=20`` and ``max_chk_threshold=30``, and ``"normal"`` utilizes
:obj:`CRESTPruner <rdmc.conformer_generation.pruners.CRESTPruner>`.
metric (SCGMetric, optional): The available option is `SCGMetric <rdmc.conformer_generation.metrics.SCGMetric>`.
Expand Down Expand Up @@ -244,6 +245,7 @@ def set_config(self,
self.min_iters = 5 if not min_iters else min_iters
self.max_iters = 100 if not max_iters else max_iters


class ConformerGenerator:
"""
A module for conformer generation. The workflow follows an embed -> optimize -> prune cycle with
Expand All @@ -266,14 +268,15 @@ class ConformerGenerator:
final_modules (list): List of instances of optimizer/pruner to run after initial cycles complete.
save_dir (str or Pathlike object, optional): The path to save the intermediate files and outputs generated during the generation.
"""

def __init__(self,
smiles: str,
multiplicity: Optional[int] = None,
optimizer: Optional['ConfGenOptimizer'] = None,
pruner: Optional['ConfGenPruner'] = None,
verifiers: Optional[Union['Verifier',List['Verifier']]] = None,
verifiers: Optional[Union['Verifier', List['Verifier']]] = None,
sampler: Optional['TorisonalSampler'] = None,
final_modules: Optional[Union['ConfGenOptimizer','Verifier']] = None,
final_modules: Optional[Union['ConfGenOptimizer', 'Verifier']] = None,
save_dir: Optional[str] = None,
) -> 'ConformerGenerator':
"""
Expand All @@ -299,7 +302,7 @@ def __init__(self,
self.logger = logging.getLogger(f"{self.__class__.__name__}")
self.smiles = smiles
if multiplicity:
self.multiplicity = multiplicity
self.multiplicity = multiplicity
else:
mol = RDKitMol.FromSmiles(smiles)
mul = mol.GetSpinMultiplicity()
Expand Down Expand Up @@ -375,7 +378,7 @@ def set_filter(self,
energy_dict = mol.energy
KeepIDs = mol.KeepIDs

sorted_index = [k for k, v in sorted(energy_dict.items(), key = lambda item: item[1])] # Order by energy
sorted_index = [k for k, v in sorted(energy_dict.items(), key=lambda item: item[1])] # Order by energy
filter_index = [k for k in sorted_index if KeepIDs[k]][:n_conformers]
for i in range(mol.GetNumConformers()):
if i not in filter_index:
Expand Down Expand Up @@ -428,7 +431,7 @@ def __call__(self,
self.logger.info("Running torsional sampling...")
energy_dict = opt_mol.energy
KeepIDs = opt_mol.KeepIDs
sorted_index = [k for k, v in sorted(energy_dict.items(), key = lambda item: item[1])] # Order by energy
sorted_index = [k for k, v in sorted(energy_dict.items(), key=lambda item: item[1])] # Order by energy
filter_index = [k for k in sorted_index if KeepIDs[k]][:n_sampling]
found_lower_energy_index = {i: False for i in range(opt_mol.GetNumConformers())}
for id in filter_index:
Expand Down
Loading

0 comments on commit 39a9380

Please sign in to comment.