Skip to content

Commit

Permalink
Fix OBAlign(includeH=False, symmetry=False) can't take keywords (#3403
Browse files Browse the repository at this point in the history
)

* run lint.yml with python 3.9

* OBAlign(includeH=False, symmetry=False) -> OBAlign(False, False)

* remove unused np.abs(d, d) in pbc_coord_intersection
  • Loading branch information
janosh committed Oct 12, 2023
1 parent b9a4ef4 commit 6e7e8f0
Show file tree
Hide file tree
Showing 6 changed files with 65 additions and 69 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/lint.yml
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ jobs:
- name: Set up Python
uses: actions/setup-python@v4
with:
python-version: "3.11"
python-version: "3.9"

- name: Install dependencies
run: |
Expand Down
40 changes: 20 additions & 20 deletions pymatgen/analysis/molecule_matcher.py
Original file line number Diff line number Diff line change
Expand Up @@ -143,7 +143,7 @@ def uniform_labels(self, mol1, mol2):
label2_list = tuple(tuple(p[1] + 1 for p in x) for x in sorted_isomorph)

vmol1 = ob_mol1
aligner = openbabel.OBAlign(True, False) # meaning includeH=True, symmetry=False # noqa: FBT003
aligner = openbabel.OBAlign(True, False) # includeH=True, symmetry=False # noqa: FBT003
aligner.SetRefMol(vmol1)
least_rmsd = float("Inf")
best_label2 = None
Expand Down Expand Up @@ -355,44 +355,44 @@ def _align_heavy_atoms(mol1, mol2, vmol1, vmol2, ilabel1, ilabel2, eq_atoms):
Return:
corrected inchi labels of heavy atoms of the second molecule
"""
nvirtual = vmol1.NumAtoms()
nheavy = len(ilabel1)
n_virtual = vmol1.NumAtoms()
n_heavy = len(ilabel1)

for i in ilabel2: # add all heavy atoms
for idx in ilabel2: # add all heavy atoms
a1 = vmol1.NewAtom()
a1.SetAtomicNum(1)
a1.SetVector(0.0, 0.0, 0.0) # useless, just to pair with vmol2
oa2 = mol2.GetAtom(i)
oa2 = mol2.GetAtom(idx)
a2 = vmol2.NewAtom()
a2.SetAtomicNum(1)
# align using the virtual atoms, these atoms are not
# used to align, but match by positions
a2.SetVector(oa2.GetVector())

aligner = openbabel.OBAlign(includeH=False, symmetry=False)
aligner = openbabel.OBAlign(False, False) # includeH=False, symmetry=False # noqa: FBT003
aligner.SetRefMol(vmol1)
aligner.SetTargetMol(vmol2)
aligner.Align()
aligner.UpdateCoords(vmol2)

canon_mol1 = openbabel.OBMol()
for i in ilabel1:
oa1 = mol1.GetAtom(i)
for idx in ilabel1:
oa1 = mol1.GetAtom(idx)
a1 = canon_mol1.NewAtom()
a1.SetAtomicNum(oa1.GetAtomicNum())
a1.SetVector(oa1.GetVector())

aligned_mol2 = openbabel.OBMol()
for i in range(nvirtual + 1, nvirtual + nheavy + 1):
oa2 = vmol2.GetAtom(i)
for idx in range(n_virtual + 1, n_virtual + n_heavy + 1):
oa2 = vmol2.GetAtom(idx)
a2 = aligned_mol2.NewAtom()
a2.SetAtomicNum(oa2.GetAtomicNum())
a2.SetVector(oa2.GetVector())

canon_label2 = list(range(1, nheavy + 1))
canon_label2 = list(range(1, n_heavy + 1))
for symm in eq_atoms:
for i in symm:
canon_label2[i - 1] = -1
for idx in symm:
canon_label2[idx - 1] = -1
for symm in eq_atoms:
candidates1 = list(symm)
candidates2 = list(symm)
Expand All @@ -409,7 +409,7 @@ def _align_heavy_atoms(mol1, mol2, vmol1, vmol2, ilabel1, ilabel2, eq_atoms):
canon_label2[c2 - 1] = canon_idx
candidates1.remove(canon_idx)

canon_inchi_orig_map2 = list(zip(canon_label2, list(range(1, nheavy + 1)), ilabel2))
canon_inchi_orig_map2 = list(zip(canon_label2, list(range(1, n_heavy + 1)), ilabel2))
canon_inchi_orig_map2.sort(key=lambda m: m[0])
return tuple(x[2] for x in canon_inchi_orig_map2)

Expand Down Expand Up @@ -448,7 +448,7 @@ def _align_hydrogen_atoms(mol1, mol2, heavy_indices1, heavy_indices2):
a2.SetAtomicNum(oa2.GetAtomicNum())
a2.SetVector(oa2.GetVector())

aligner = openbabel.OBAlign(includeH=False, symmetry=False)
aligner = openbabel.OBAlign(False, False) # includeH=False, symmetry=False # noqa: FBT003
aligner.SetRefMol(cmol1)
aligner.SetTargetMol(cmol2)
aligner.Align()
Expand Down Expand Up @@ -631,23 +631,23 @@ def _calc_rms(mol1, mol2, clabel1, clabel2):
Returns:
The RMSD.
"""
obmol1 = BabelMolAdaptor(mol1).openbabel_mol
obmol2 = BabelMolAdaptor(mol2).openbabel_mol
ob_mol1 = BabelMolAdaptor(mol1).openbabel_mol
ob_mol2 = BabelMolAdaptor(mol2).openbabel_mol

cmol1 = openbabel.OBMol()
for i in clabel1:
oa1 = obmol1.GetAtom(i)
oa1 = ob_mol1.GetAtom(i)
a1 = cmol1.NewAtom()
a1.SetAtomicNum(oa1.GetAtomicNum())
a1.SetVector(oa1.GetVector())
cmol2 = openbabel.OBMol()
for i in clabel2:
oa2 = obmol2.GetAtom(i)
oa2 = ob_mol2.GetAtom(i)
a2 = cmol2.NewAtom()
a2.SetAtomicNum(oa2.GetAtomicNum())
a2.SetVector(oa2.GetVector())

aligner = openbabel.OBAlign(True, False) # meaning includeH=True, symmetry=False # noqa: FBT003
aligner = openbabel.OBAlign(True, False) # includeH=True, symmetry=False # noqa: FBT003
aligner.SetRefMol(cmol1)
aligner.SetTargetMol(cmol2)
aligner.Align()
Expand Down
11 changes: 5 additions & 6 deletions pymatgen/core/structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -2300,7 +2300,7 @@ def site_label(site):
# min_vecs are approximate periodicities of the cell. The exact
# periodicities from the supercell matrices are checked against these
# first
min_fcoords = min(grouped_fcoords, key=lambda x: len(x))
min_fcoords = min(grouped_fcoords, key=len)
min_vecs = min_fcoords - min_fcoords[0]

# fractional tolerance in the supercell
Expand All @@ -2311,18 +2311,17 @@ def pbc_coord_intersection(fc1, fc2, tol):
"""Returns the fractional coords in fc1 that have coordinates
within tolerance to some coordinate in fc2.
"""
d = fc1[:, None, :] - fc2[None, :, :]
d -= np.round(d)
np.abs(d, d)
return fc1[np.any(np.all(d < tol, axis=-1), axis=-1)]
dist = fc1[:, None, :] - fc2[None, :, :]
dist -= np.round(dist)
return fc1[np.any(np.all(dist < tol, axis=-1), axis=-1)]

# here we reduce the number of min_vecs by enforcing that every
# vector in min_vecs approximately maps each site onto a similar site.
# The subsequent processing is O(fu^3 * min_vecs) = O(n^4) if we do no
# reduction.
# This reduction is O(n^3) so usually is an improvement. Using double
# the tolerance because both vectors are approximate
for group in sorted(grouped_fcoords, key=lambda x: len(x)):
for group in sorted(grouped_fcoords, key=len):
for frac_coords in group:
min_vecs = pbc_coord_intersection(min_vecs, group - frac_coords, super_ftol_2)

Expand Down
29 changes: 13 additions & 16 deletions pymatgen/core/surface.py
Original file line number Diff line number Diff line change
Expand Up @@ -225,32 +225,29 @@ def get_tasker2_slabs(self, tol: float = 0.01, same_species_only=True):
spga = SpacegroupAnalyzer(self)
symm_structure = spga.get_symmetrized_structure()

def equi_index(site):
for i, equi_sites in enumerate(symm_structure.equivalent_sites):
def equi_index(site) -> int:
for idx, equi_sites in enumerate(symm_structure.equivalent_sites):
if site in equi_sites:
return i
return idx
raise ValueError("Cannot determine equi index!")

for surface_site, shift in [
(sorted_csites[0], slab_ratio),
(sorted_csites[-1], -slab_ratio),
]:
tomove = []
for surface_site, shift in [(sorted_csites[0], slab_ratio), (sorted_csites[-1], -slab_ratio)]:
to_move = []
fixed = []
for site in sites:
if abs(site.c - surface_site.c) < tol and (
(not same_species_only) or site.species == surface_site.species
):
tomove.append(site)
to_move.append(site)
else:
fixed.append(site)

# Sort and group the sites by the species and symmetry equivalence
tomove = sorted(tomove, key=lambda s: equi_index(s))
to_move = sorted(to_move, key=equi_index)

grouped = [list(sites) for k, sites in itertools.groupby(tomove, key=lambda s: equi_index(s))]
grouped = [list(sites) for k, sites in itertools.groupby(to_move, key=equi_index)]

if len(tomove) == 0 or any(len(g) % 2 != 0 for g in grouped):
if len(to_move) == 0 or any(len(g) % 2 != 0 for g in grouped):
warnings.warn(
"Odd number of sites to divide! Try changing "
"the tolerance to ensure even division of "
Expand All @@ -266,7 +263,7 @@ def equi_index(site):
species = [site.species for site in fixed]
fcoords = [site.frac_coords for site in fixed]

for s in tomove:
for s in to_move:
species.append(s.species)
for group in selection:
if s in group:
Expand Down Expand Up @@ -1171,11 +1168,11 @@ def repair_broken_bonds(self, slab, bonds):
# find its NNs with the corresponding
# species it should be coordinated with
neighbors = slab.get_neighbors(slab[i], blength, include_index=True)
tomove = [nn[2] for nn in neighbors if nn[0].species_string == element2]
tomove.append(i)
to_move = [nn[2] for nn in neighbors if nn[0].species_string == element2]
to_move.append(i)
# and then move those NNs along with the central
# atom back to the other side of the slab again
slab = self.move_to_other_side(slab, tomove)
slab = self.move_to_other_side(slab, to_move)

return slab

Expand Down
48 changes: 24 additions & 24 deletions pymatgen/io/babel.py
Original file line number Diff line number Diff line change
Expand Up @@ -74,11 +74,11 @@ def __init__(self, mol: Molecule | openbabel.OBMol | pybel.Molecule) -> None:
ob_mol.SetTotalCharge(int(mol.charge))
ob_mol.Center()
ob_mol.EndModify()
self._obmol = ob_mol
self._ob_mol = ob_mol
elif isinstance(mol, openbabel.OBMol):
self._obmol = mol
self._ob_mol = mol
elif isinstance(mol, pybel.Molecule):
self._obmol = mol.OBMol
self._ob_mol = mol.OBMol
else:
raise ValueError(f"Unsupported input type {type(mol)}, must be Molecule, openbabel.OBMol or pybel.Molecule")

Expand All @@ -87,15 +87,15 @@ def pymatgen_mol(self):
"""Returns pymatgen Molecule object."""
sp = []
coords = []
for atom in openbabel.OBMolAtomIter(self._obmol):
for atom in openbabel.OBMolAtomIter(self._ob_mol):
sp.append(atom.GetAtomicNum())
coords.append([atom.GetX(), atom.GetY(), atom.GetZ()])
return Molecule(sp, coords)

@property
def openbabel_mol(self):
"""Returns OpenBabel's OBMol."""
return self._obmol
return self._ob_mol

def localopt(self, forcefield="mmff94", steps=500):
"""
Expand All @@ -106,9 +106,9 @@ def localopt(self, forcefield="mmff94", steps=500):
'mmff94', 'mmff94s', and 'uff'.
steps: Default is 500.
"""
pybelmol = pybel.Molecule(self._obmol)
pybelmol = pybel.Molecule(self._ob_mol)
pybelmol.localopt(forcefield=forcefield, steps=steps)
self._obmol = pybelmol.OBMol
self._ob_mol = pybelmol.OBMol

def make3d(self, forcefield="mmff94", steps=50):
"""
Expand All @@ -129,13 +129,13 @@ def make3d(self, forcefield="mmff94", steps=50):
'mmff94', 'mmff94s', and 'uff'.
steps: Default is 50.
"""
pybelmol = pybel.Molecule(self._obmol)
pybelmol = pybel.Molecule(self._ob_mol)
pybelmol.make3D(forcefield=forcefield, steps=steps)
self._obmol = pybelmol.OBMol
self._ob_mol = pybelmol.OBMol

def add_hydrogen(self):
"""Add hydrogens (make all hydrogen explicit)."""
self._obmol.AddHydrogens()
self._ob_mol.AddHydrogens()

def remove_bond(self, idx1, idx2):
"""
Expand All @@ -145,11 +145,11 @@ def remove_bond(self, idx1, idx2):
idx1: The atom index of one of the atoms participating the in bond
idx2: The atom index of the other atom participating in the bond
"""
for obbond in openbabel.OBMolBondIter(self._obmol):
for obbond in openbabel.OBMolBondIter(self._ob_mol):
if (obbond.GetBeginAtomIdx() == idx1 and obbond.GetEndAtomIdx() == idx2) or (
obbond.GetBeginAtomIdx() == idx2 and obbond.GetEndAtomIdx() == idx1
):
self._obmol.DeleteBond(obbond)
self._ob_mol.DeleteBond(obbond)

def rotor_conformer(self, *rotor_args, algo="WeightedRotorSearch", forcefield="mmff94"):
"""
Expand All @@ -172,7 +172,7 @@ def rotor_conformer(self, *rotor_args, algo="WeightedRotorSearch", forcefield="m
forcefield (str): Default is mmff94. Options are 'gaff', 'ghemical',
'mmff94', 'mmff94s', and 'uff'.
"""
if self._obmol.GetDimension() != 3:
if self._ob_mol.GetDimension() != 3:
self.make3d()
else:
self.add_hydrogen()
Expand All @@ -199,7 +199,7 @@ def rotor_conformer(self, *rotor_args, algo="WeightedRotorSearch", forcefield="m
)
rotor_search = ff.WeightedRotorSearch
rotor_search(*rotor_args)
ff.GetConformers(self._obmol)
ff.GetConformers(self._ob_mol)

def gen3d_conformer(self):
"""
Expand All @@ -222,7 +222,7 @@ def gen3d_conformer(self):
between speed and finding the global energy minimum.
"""
gen3d = openbabel.OBOp.FindType("Gen3D")
gen3d.Do(self._obmol)
gen3d.Do(self._ob_mol)

def confab_conformers(
self,
Expand Down Expand Up @@ -254,7 +254,7 @@ def confab_conformers(
Returns:
list[Molecule]: Molecule objects for generated conformers.
"""
if self._obmol.GetDimension() != 3:
if self._ob_mol.GetDimension() != 3:
self.make3d()
else:
self.add_hydrogen()
Expand All @@ -268,31 +268,31 @@ def confab_conformers(
print(f"{len(freeze_atoms)} atoms will be freezed")
constraints = openbabel.OBFFConstraints()

for atom in openbabel.OBMolAtomIter(self._obmol):
for atom in openbabel.OBMolAtomIter(self._ob_mol):
atom_id = atom.GetIndex() + 1
if id in freeze_atoms:
constraints.AddAtomConstraint(atom_id)
ff.SetConstraints(constraints)

# Confab conformer generation
ff.DiverseConfGen(rmsd_cutoff, conf_cutoff, energy_cutoff, verbose)
ff.GetConformers(self._obmol)
ff.GetConformers(self._ob_mol)

# Number of conformers generated by Confab conformer generation
conformer_num = self._obmol.NumConformers()
conformer_num = self._ob_mol.NumConformers()

conformers = []
for i in range(conformer_num):
self._obmol.SetConformer(i)
conformer = copy.deepcopy(BabelMolAdaptor(self._obmol).pymatgen_mol)
self._ob_mol.SetConformer(i)
conformer = copy.deepcopy(BabelMolAdaptor(self._ob_mol).pymatgen_mol)
conformers.append(conformer)
self._obmol.SetConformer(0)
self._ob_mol.SetConformer(0)
return conformers

@property
def pybel_mol(self):
"""Returns Pybel's Molecule object."""
return pybel.Molecule(self._obmol)
return pybel.Molecule(self._ob_mol)

def write_file(self, filename, file_format="xyz"):
"""
Expand All @@ -302,7 +302,7 @@ def write_file(self, filename, file_format="xyz"):
filename: Filename of file to output
file_format: String specifying any OpenBabel supported formats.
"""
mol = pybel.Molecule(self._obmol)
mol = pybel.Molecule(self._ob_mol)
return mol.write(file_format, filename, overwrite=True)

@staticmethod
Expand Down
4 changes: 2 additions & 2 deletions pymatgen/io/lammps/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -336,7 +336,7 @@ def write_pdb(mol: Molecule, filename: str, name: str | None = None, num=None) -
name = name or f"ml{num}"

# bma = BabelMolAdaptor(mol)
pbm = pybel.Molecule(bma._obmol)
pbm = pybel.Molecule(bma._ob_mol)
for x in pbm.residues:
x.OBResidue.SetName(name)
x.OBResidue.SetNum(num)
Expand Down Expand Up @@ -414,7 +414,7 @@ def restore_site_properties(self, site_property: str = "ff_map", filename: str |

filename = filename or self.control_params["output"]
bma = BabelMolAdaptor.from_file(filename, "pdb")
pbm = pybel.Molecule(bma._obmol)
pbm = pybel.Molecule(bma._ob_mol)

assert len(pbm.residues) == sum(x["number"] for x in self.param_list)

Expand Down

0 comments on commit 6e7e8f0

Please sign in to comment.