def _get_interatomic_distances(self, coords, cell, pbc): """ Return the interatomic distances matrix and its associated coordinates differences matrices. Returns: dists: a `float32` array of shape `[N, N]` where N is the number of atoms as the interatomic distances matrix. """ if not self.is_periodic: dists = pairwise_distances(coords) else: atoms = Atoms( symbols=self._species, positions=coords, cell=cell, pbc=pbc ) dists = atoms.get_all_distances(mic=True) del atoms # Manually set the distances between ghost atoms and real atoms to inf. if self._num_ghosts > 0: dists[:, -self._num_ghosts:] = np.inf dists[-self._num_ghosts:, :] = np.inf return dists
def check_correct(methods, func_map, pos, lbox): # get ase distance table (reference distance table) natom = len(pos) axes = lbox * np.eye(3) s1 = Atoms('H%d' % natom, cell=axes, positions=pos, pbc=[1, 1, 1]) #print [cmd for cmd in dir(s1) if 'dist' in cmd] #assert 0 # check only unique pair distances idx = np.triu_indices(natom, k=1) adt = s1.get_all_distances(mic=True) for method in methods: func = func_map[method] mydt = func(pos, lbox) print('%s correct:' % method, np.allclose(mydt[idx], adt[idx]))
a = Atoms('HOC', positions=[(1, 1, 1), (3, 1, 1), (6, 1, 1)]) a.set_cell((9, 2, 2)) a.set_pbc((True, False, False)) # Calculate indiviually with mic=True assert a.get_distance(0, 1, mic=True) == 2 assert a.get_distance(1, 2, mic=True) == 3 assert a.get_distance(0, 2, mic=True) == 4 # Calculate indiviually with mic=False assert a.get_distance(0, 1, mic=False) == 2 assert a.get_distance(1, 2, mic=False) == 3 assert a.get_distance(0, 2, mic=False) == 5 # Calculate in groups with mic=True assert (a.get_distances(0, [1, 2], mic=True) == [2, 4]).all() # Calculate in groups with mic=False assert (a.get_distances(0, [1, 2], mic=False) == [2, 5]).all() # Calculate all with mic=True assert (a.get_all_distances(mic=True) == [[0, 2, 4], [2, 0, 3], [4, 3, 0]]).all() # Calculate all with mic=False assert (a.get_all_distances(mic=False) == [[0, 2, 5], [2, 0, 3], [5, 3, 0]]).all()
class Conformer(): """ A class for generating and editing 3D conformers of molecules """ def __init__(self, smiles=None, rmg_molecule=None, index=0): self.energy = None self.index = index if (smiles or rmg_molecule): if smiles and rmg_molecule: assert rmg_molecule.isIsomorphic(RMGMolecule( SMILES=smiles)), "SMILES string did not match RMG Molecule object" self.smiles = smiles self.rmg_molecule = rmg_molecule elif rmg_molecule: self.rmg_molecule = rmg_molecule self.smiles = rmg_molecule.toSMILES() else: self.smiles = smiles self.rmg_molecule = RMGMolecule(SMILES=smiles) self.rmg_molecule.updateMultiplicity() self.get_molecules() self.get_geometries() self._symmetry_number = None else: self.smiles = None self.rmg_molecule = None self.rdkit_molecule = None self.ase_molecule = None self.bonds = [] self.angles = [] self.torsions = [] self.cistrans = [] self.chiral_centers = [] self._symmetry_number = None def __repr__(self): return '<Conformer "{}">'.format(self.smiles) def copy(self): copy_conf = Conformer() copy_conf.smiles = self.smiles copy_conf.rmg_molecule = self.rmg_molecule.copy() copy_conf.rdkit_molecule = self.rdkit_molecule.__copy__() copy_conf.ase_molecule = self.ase_molecule.copy() copy_conf.get_geometries() copy_conf.energy = self.energy return copy_conf @property def symmetry_number(self): if not self._symmetry_number: self._symmetry_number = self.calculate_symmetry_number() return self._symmetry_number def get_rdkit_mol(self): """ A method for creating an rdkit geometry from an rmg mol """ assert self.rmg_molecule, "Cannot create an RDKit geometry without an RMG molecule object" RDMol = self.rmg_molecule.toRDKitMol(removeHs=False) rdkit.Chem.AllChem.EmbedMolecule(RDMol) self.rdkit_molecule = RDMol mol_list = AllChem.MolToMolBlock(self.rdkit_molecule).split('\n') for i, atom in enumerate(self.rmg_molecule.atoms): j = i + 4 coords = mol_list[j].split()[:3] for k, coord in enumerate(coords): coords[k] = float(coord) atom.coords = np.array(coords) return self.rdkit_molecule def get_ase_mol(self): """ A method for creating an ase atoms object from an rdkit mol """ if not self.rdkit_molecule: self.get_rdkit_mol() mol_list = AllChem.MolToMolBlock(self.rdkit_molecule).split('\n') ase_atoms = [] for i, line in enumerate(mol_list): if i > 3: try: atom0, atom1, bond, rest = line atom0 = int(atom0) atom0 = int(atom1) bond = float(bond) except ValueError: try: x, y, z, symbol = line.split()[0:4] x = float(x) y = float(y) z = float(z) ase_atoms.append( Atom(symbol=symbol, position=(x, y, z))) except BaseException: continue self.ase_molecule = Atoms(ase_atoms) return self.ase_molecule def get_molecules(self): if not self.rmg_molecule: self.rmg_molecule = RMGMolecule(SMILES=self.smiles) self.rdkit_molecule = self.get_rdkit_mol() self.ase_molecule = self.get_ase_mol() self.get_geometries() return self.rdkit_molecule, self.ase_molecule def view(self): """ A method designed to create a 3D figure of the AutoTST_Molecule with py3Dmol from the rdkit_molecule """ mb = Chem.MolToMolBlock(self.rdkit_molecule) p = py3Dmol.view(width=600, height=600) p.addModel(mb, "sdf") p.setStyle({'stick': {}}) p.setBackgroundColor('0xeeeeee') p.zoomTo() return p.show() def get_bonds(self): """ A method for identifying all of the bonds in a conformer """ bond_list = [] for bond in self.rdkit_molecule.GetBonds(): bond_list.append((bond.GetBeginAtomIdx(), bond.GetEndAtomIdx())) bonds = [] for index, indices in enumerate(bond_list): i, j = indices length = self.ase_molecule.get_distance(i, j) center = False if ((self.rmg_molecule.atoms[i].label) and ( self.rmg_molecule.atoms[j].label)): center = True bond = Bond(index=index, atom_indices=indices, length=length, reaction_center=center) mask = self.get_mask(bond) bond.mask = mask bonds.append(bond) self.bonds = bonds return self.bonds def get_angles(self): """ A method for identifying all of the angles in a conformer """ angle_list = [] for atom1 in self.rdkit_molecule.GetAtoms(): for atom2 in atom1.GetNeighbors(): for atom3 in atom2.GetNeighbors(): if atom1.GetIdx() == atom3.GetIdx(): continue to_add = (atom1.GetIdx(), atom2.GetIdx(), atom3.GetIdx()) if (to_add in angle_list) or ( tuple(reversed(to_add)) in angle_list): continue angle_list.append(to_add) angles = [] for index, indices in enumerate(angle_list): i, j, k = indices degree = self.ase_molecule.get_angle(i, j, k) ang = Angle(index=index, atom_indices=indices, degree=degree, mask=[]) mask = self.get_mask(ang) reaction_center = False angles.append(Angle(index=index, atom_indices=indices, degree=degree, mask=mask, reaction_center=reaction_center)) self.angles = angles return self.angles def get_torsions(self): """ A method for identifying all of the torsions in a conformer """ torsion_list = [] for bond1 in self.rdkit_molecule.GetBonds(): atom1 = bond1.GetBeginAtom() atom2 = bond1.GetEndAtom() if atom1.IsInRing() or atom2.IsInRing(): # Making sure that bond1 we're looking at are not in a ring continue bond_list1 = list(atom1.GetBonds()) bond_list2 = list(atom2.GetBonds()) if not len(bond_list1) > 1 and not len(bond_list2) > 1: # Making sure that there are more than one bond attached to # the atoms we're looking at continue # Getting the 0th and 3rd atom and insuring that atoms # attached to the 1st and 2nd atom are not terminal hydrogens # We also make sure that all of the atoms are properly bound # together # If the above are satisfied, we append a tuple of the torsion our # torsion_list got_atom0 = False got_atom3 = False for bond0 in bond_list1: atomX = bond0.GetOtherAtom(atom1) # if atomX.GetAtomicNum() == 1 and len(atomX.GetBonds()) == 1: # This means that we have a terminal hydrogen, skip this # NOTE: for H_abstraction TSs, a non teminal H should exist # continue if atomX.GetIdx() != atom2.GetIdx(): got_atom0 = True atom0 = atomX for bond2 in bond_list2: atomY = bond2.GetOtherAtom(atom2) # if atomY.GetAtomicNum() == 1 and len(atomY.GetBonds()) == 1: # This means that we have a terminal hydrogen, skip this # continue if atomY.GetIdx() != atom1.GetIdx(): got_atom3 = True atom3 = atomY if not (got_atom0 and got_atom3): # Making sure atom0 and atom3 were not found continue # Looking to make sure that all of the atoms are properly bonded to # eached if ( "SINGLE" in str( self.rdkit_molecule.GetBondBetweenAtoms( atom1.GetIdx(), atom2.GetIdx()).GetBondType()) and self.rdkit_molecule.GetBondBetweenAtoms( atom0.GetIdx(), atom1.GetIdx()) and self.rdkit_molecule.GetBondBetweenAtoms( atom1.GetIdx(), atom2.GetIdx()) and self.rdkit_molecule.GetBondBetweenAtoms( atom2.GetIdx(), atom3.GetIdx())): torsion_tup = (atom0.GetIdx(), atom1.GetIdx(), atom2.GetIdx(), atom3.GetIdx()) already_in_list = False for torsion_entry in torsion_list: a, b, c, d = torsion_entry e, f, g, h = torsion_tup if (b, c) == (f, g) or (b, c) == (g, f): already_in_list = True if not already_in_list: torsion_list.append(torsion_tup) torsions = [] for index, indices in enumerate(torsion_list): i, j, k, l = indices dihedral = self.ase_molecule.get_dihedral(i, j, k, l) tor = Torsion(index=index, atom_indices=indices, dihedral=dihedral, mask=[]) mask = self.get_mask(tor) reaction_center = False torsions.append(Torsion(index=index, atom_indices=indices, dihedral=dihedral, mask=mask, reaction_center=reaction_center)) self.torsions = torsions return self.torsions def get_cistrans(self): """ A method for identifying all possible cistrans bonds in a molecule """ torsion_list = [] cistrans_list = [] for bond1 in self.rdkit_molecule.GetBonds(): atom1 = bond1.GetBeginAtom() atom2 = bond1.GetEndAtom() if atom1.IsInRing() or atom2.IsInRing(): # Making sure that bond1 we're looking at are not in a ring continue bond_list1 = list(atom1.GetBonds()) bond_list2 = list(atom2.GetBonds()) if not len(bond_list1) > 1 and not len(bond_list2) > 1: # Making sure that there are more than one bond attached to # the atoms we're looking at continue # Getting the 0th and 3rd atom and insuring that atoms # attached to the 1st and 2nd atom are not terminal hydrogens # We also make sure that all of the atoms are properly bound # together # If the above are satisfied, we append a tuple of the torsion our # torsion_list got_atom0 = False got_atom3 = False for bond0 in bond_list1: atomX = bond0.GetOtherAtom(atom1) # if atomX.GetAtomicNum() == 1 and len(atomX.GetBonds()) == 1: # This means that we have a terminal hydrogen, skip this # NOTE: for H_abstraction TSs, a non teminal H should exist # continue if atomX.GetIdx() != atom2.GetIdx(): got_atom0 = True atom0 = atomX for bond2 in bond_list2: atomY = bond2.GetOtherAtom(atom2) # if atomY.GetAtomicNum() == 1 and len(atomY.GetBonds()) == 1: # This means that we have a terminal hydrogen, skip this # continue if atomY.GetIdx() != atom1.GetIdx(): got_atom3 = True atom3 = atomY if not (got_atom0 and got_atom3): # Making sure atom0 and atom3 were not found continue # Looking to make sure that all of the atoms are properly bonded to # eached if ( "DOUBLE" in str( self.rdkit_molecule.GetBondBetweenAtoms( atom1.GetIdx(), atom2.GetIdx()).GetBondType()) and self.rdkit_molecule.GetBondBetweenAtoms( atom0.GetIdx(), atom1.GetIdx()) and self.rdkit_molecule.GetBondBetweenAtoms( atom1.GetIdx(), atom2.GetIdx()) and self.rdkit_molecule.GetBondBetweenAtoms( atom2.GetIdx(), atom3.GetIdx())): torsion_tup = (atom0.GetIdx(), atom1.GetIdx(), atom2.GetIdx(), atom3.GetIdx()) already_in_list = False for torsion_entry in torsion_list: a, b, c, d = torsion_entry e, f, g, h = torsion_tup if (b, c) == (f, g) or (b, c) == (g, f): already_in_list = True if not already_in_list: cistrans_list.append(torsion_tup) cistrans = [] for ct_index, indices in enumerate(cistrans_list): i, j, k, l = indices b0 = self.rdkit_molecule.GetBondBetweenAtoms(i, j) b1 = self.rdkit_molecule.GetBondBetweenAtoms(j, k) b2 = self.rdkit_molecule.GetBondBetweenAtoms(k, l) b0.SetBondDir(Chem.BondDir.ENDUPRIGHT) b2.SetBondDir(Chem.BondDir.ENDDOWNRIGHT) Chem.AssignStereochemistry(self.rdkit_molecule, force=True) if "STEREOZ" in str(b1.GetStereo()): if round(self.ase_molecule.get_dihedral(i, j, k, l), -1) == 0: atom = self.rdkit_molecule.GetAtomWithIdx(k) bonds = atom.GetBonds() for bond in bonds: indexes = [ bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()] if not ((sorted([j, k]) == sorted(indexes)) or ( sorted([k, l]) == sorted(indexes))): break for index in indexes: if not (index in indices): l = index break indices = [i, j, k, l] stero = "Z" else: if round( self.ase_molecule.get_dihedral( i, j, k, l), -1) == 180: atom = self.rdkit_molecule.GetAtomWithIdx(k) bonds = atom.GetBonds() for bond in bonds: indexes = [ bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()] if not ((sorted([j, k]) == sorted(indexes)) or ( sorted([k, l]) == sorted(indexes))): break for index in indexes: if not (index in indices): l = index break indices = [i, j, k, l] stero = "E" dihedral = self.ase_molecule.get_dihedral(i, j, k, l) tor = CisTrans(index=ct_index, atom_indices=indices, dihedral=dihedral, mask=[], stero=stero) mask = self.get_mask(tor) reaction_center = False cistrans.append(CisTrans(index=ct_index, atom_indices=indices, dihedral=dihedral, mask=mask, stero=stero ) ) self.cistrans = cistrans return self.cistrans def get_mask(self, geometry): """ Getting the right hand mask for a geometry object: - self: an AutoTST Conformer object - geometry: a Bond, Angle, Dihedral, or Torsion object """ rdkit_atoms = self.rdkit_molecule.GetAtoms() if (isinstance(geometry, autotst.geometry.Torsion) or isinstance(geometry, autotst.geometry.CisTrans)): L1, L0, R0, R1 = geometry.atom_indices # trying to get the left hand side of this torsion LHS_atoms_index = [L0, L1] RHS_atoms_index = [R0, R1] elif isinstance(geometry, autotst.geometry.Angle): a1, a2, a3 = geometry.atom_indices LHS_atoms_index = [a2, a1] RHS_atoms_index = [a2, a3] elif isinstance(geometry, autotst.geometry.Bond): a1, a2 = geometry.atom_indices LHS_atoms_index = [a1] RHS_atoms_index = [a2] complete_RHS = False i = 0 atom_index = RHS_atoms_index[0] while complete_RHS is False: try: RHS_atom = rdkit_atoms[atom_index] for neighbor in RHS_atom.GetNeighbors(): if (neighbor.GetIdx() in RHS_atoms_index) or ( neighbor.GetIdx() in LHS_atoms_index): continue else: RHS_atoms_index.append(neighbor.GetIdx()) i += 1 atom_index = RHS_atoms_index[i] except IndexError: complete_RHS = True mask = [index in RHS_atoms_index for index in range( len(self.ase_molecule))] return mask def get_chiral_centers(self): """ A method to identify """ centers = rdkit.Chem.FindMolChiralCenters( self.rdkit_molecule, includeUnassigned=True) chiral_centers = [] for index, center in enumerate(centers): atom_index, chirality = center chiral_centers.append( ChiralCenter( index=index, atom_index=atom_index, chirality=chirality)) self.chiral_centers = chiral_centers return self.chiral_centers def get_geometries(self): """ A helper method to obtain all geometry things """ self.bonds = self.get_bonds() self.angles = self.get_angles() self.torsions = self.get_torsions() self.cistrans = self.get_cistrans() self.chiral_centers = self.get_chiral_centers() return ( self.bonds, self.angles, self.torsions, self.cistrans, self.chiral_centers) def update_coords(self): """ A function that creates distance matricies for the RMG, ASE, and RDKit molecules and finds which (if any) are different. If one is different, this will update the coordinates of the other two with the different one. If all three are different, nothing will happen. If all are the same, nothing will happen. """ rdkit_dm = rdkit.Chem.rdmolops.Get3DDistanceMatrix(self.rdkit_molecule) ase_dm = self.ase_molecule.get_all_distances() l = len(self.rmg_molecule.atoms) rmg_dm = np.zeros((l, l)) for i, atom_i in enumerate(self.rmg_molecule.atoms): for j, atom_j in enumerate(self.rmg_molecule.atoms): rmg_dm[i][j] = np.linalg.norm(atom_i.coords - atom_j.coords) d1 = round(abs(rdkit_dm - ase_dm).max(), 3) d2 = round(abs(rdkit_dm - rmg_dm).max(), 3) d3 = round(abs(ase_dm - rmg_dm).max(), 3) if np.all(np.array([d1, d2, d3]) > 0): return False, None if np.any(np.array([d1, d2, d3]) > 0): if d1 == 0: diff = "rmg" self.update_coords_from("rmg") elif d2 == 0: diff = "ase" self.update_coords_from("ase") else: diff = "rdkit" self.update_coords_from("rdkit") return True, diff else: return True, None def update_coords_from(self, mol_type="ase"): """ A method to update the coordinates of the RMG, RDKit, and ASE objects with a chosen object. """ possible_mol_types = ["ase", "rmg", "rdkit"] assert (mol_type.lower() in possible_mol_types), "Please specifiy a valid mol type. Valid types are {}".format( possible_mol_types) if mol_type.lower() == "rmg": conf = self.rdkit_molecule.GetConformers()[0] ase_atoms = [] for i, atom in enumerate(self.rmg_molecule.atoms): x, y, z = atom.coords symbol = atom.symbol conf.SetAtomPosition(i, [x, y, z]) ase_atoms.append(Atom(symbol=symbol, position=(x, y, z))) self.ase_molecule = Atoms(ase_atoms) # self.calculate_symmetry_number() elif mol_type.lower() == "ase": conf = self.rdkit_molecule.GetConformers()[0] for i, position in enumerate(self.ase_molecule.get_positions()): self.rmg_molecule.atoms[i].coords = position conf.SetAtomPosition(i, position) # self.calculate_symmetry_number() elif mol_type.lower() == "rdkit": mol_list = AllChem.MolToMolBlock(self.rdkit_molecule).split('\n') for i, atom in enumerate(self.rmg_molecule.atoms): j = i + 4 coords = mol_list[j].split()[:3] for k, coord in enumerate(coords): coords[k] = float(coord) atom.coords = np.array(coords) self.get_ase_mol() # self.calculate_symmetry_number() def set_bond_length(self, bond_index, length): """ This is a method to set bond lengths Variabels: - bond_index (int): the index of the bond you want to edit - length (float, int): the distance you want to set the bond (in angstroms) """ assert isinstance(length, (float, int)) matched = False for bond in self.bonds: if bond.index == bond_index: matched = True break if not matched: logging.info("Angle index provided is out of range. Nothing was changed.") return self i, j = bond.atom_indices self.ase_molecule.set_distance( a0=i, a1=j, distance=length, mask=bond.mask, fix=0 ) bond.length = length self.update_coords_from(mol_type="ase") return self def set_angle(self, angle_index, angle): """ A method that will set the angle of an Angle object accordingly """ assert isinstance( angle, (int, float)), "Plese provide a float or an int for the angle" matched = False for a in self.angles: if a.index == angle_index: matched = True break if not matched: logging.info("Angle index provided is out of range. Nothing was changed.") return self i, j, k = a.atom_indices self.ase_molecule.set_angle( a1=i, a2=j, a3=k, angle=angle, mask=a.mask ) a.degree = angle self.update_coords_from(mol_type="ase") return self def set_torsion(self, torsion_index, dihedral): """ A method that will set the diehdral angle of a Torsion object accordingly. """ assert isinstance( dihedral, (int, float)), "Plese provide a float or an int for the diehdral angle" matched = False for torsion in self.torsions: if torsion.index == torsion_index: matched = True break if not matched: logging.info("Torsion index provided is out of range. Nothing was changed.") return self i, j, k, l = torsion.atom_indices self.ase_molecule.set_dihedral( a1=i, a2=j, a3=k, a4=l, angle=dihedral, mask=torsion.mask ) torsion.dihedral = dihedral self.update_coords_from(mol_type="ase") return self def set_cistrans(self, cistrans_index, stero="E"): """ A module that will set a corresponding cistrans bond to the proper E/Z config """ assert stero.upper() in [ "E", "Z"], "Please specify a valid stero direction." matched = False for cistrans in self.cistrans: if cistrans.index == cistrans_index: matched = True break if not matched: logging.info("CisTrans index provided is out of range. Nothing was changed.") return self if cistrans.stero == stero.upper(): self.update_coords_from("ase") return self else: cistrans.stero = stero.upper() i, j, k, l = cistrans.atom_indices self.ase_molecule.rotate_dihedral( a1=i, a2=j, a3=k, a4=l, angle=float(180), mask=cistrans.mask ) cistrans.stero = stero.upper() self.update_coords_from(mol_type="ase") return self def set_chirality(self, chiral_center_index, stero="R"): """ A module that can set the orientation of a chiral center. """ assert stero.upper() in ["R", "S"], "Specify a valid stero orientation" centers_dict = { 'R': Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CW, 'S': Chem.rdchem.ChiralType.CHI_TETRAHEDRAL_CCW } assert isinstance(chiral_center_index, int), "Please provide an integer for the index" rdmol = self.rdkit_molecule.__copy__() match = False for chiral_center in self.chiral_centers: if chiral_center.index == chiral_center_index: match = True break if not match: logging.info("ChiralCenter index provided is out of range. Nothing was changed") return self rdmol.GetAtomWithIdx(chiral_center.atom_index).SetChiralTag( centers_dict[stero.upper()]) rdkit.Chem.rdDistGeom.EmbedMolecule(rdmol) old_torsions = self.torsions[:] + self.cistrans[:] self.rdkit_molecule = rdmol self.update_coords_from(mol_type="rdkit") # Now resetting dihedral angles in case if they changed. for torsion in old_torsions: i, j, k, l = torsion.atom_indices self.ase_molecule.set_dihedral( a1=i, a2=j, a3=k, a4=l, mask=torsion.mask, angle=torsion.dihedral, ) self.update_coords_from(mol_type="ase") return self def calculate_symmetry_number(self): from rmgpy.qm.symmetry import PointGroupCalculator from rmgpy.qm.qmdata import QMData atom_numbers = self.ase_molecule.get_atomic_numbers() coordinates = self.ase_molecule.get_positions() qmdata = QMData( groundStateDegeneracy=1, # Only needed to check if valid QMData numberOfAtoms=len(atom_numbers), atomicNumbers=atom_numbers, atomCoords=(coordinates, str('angstrom')), energy=(0.0, str('kcal/mol')) # Only needed to avoid error ) settings = type(str(''), (), dict(symmetryPath=str( 'symmetry'), scratchDirectory="."))() # Creates anonymous class pgc = PointGroupCalculator(settings, self.smiles, qmdata) pg = pgc.calculate() #os.remove("{}.symm".format(self.smiles)) if pg is not None: symmetry_number = pg.symmetryNumber else: symmetry_number = 1 return symmetry_number
assert a.get_distance(1, 2, mic=True) == 3 assert a.get_distance(0, 2, mic=True) == 4 # Calculate indiviually with mic=False assert a.get_distance(0, 1, mic=False) == 2 assert a.get_distance(1, 2, mic=False) == 3 assert a.get_distance(0, 2, mic=False) == 5 # Calculate in groups with mic=True assert (a.get_distances(0, [1, 2], mic=True) == [2, 4]).all() # Calculate in groups with mic=False assert (a.get_distances(0, [1, 2], mic=False) == [2, 5]).all() # Calculate all with mic=True assert (a.get_all_distances(mic=True) == [[0, 2, 4], [2, 0, 3], [4, 3, 0]]).all() # Calculate all with mic=False assert (a.get_all_distances(mic=False) == [[0, 2, 5], [2, 0, 3], [5, 3, 0]]).all() # Scale Distance old = a.get_distance(0, 1) a.set_distance(0, 1, 0.9, add=True, factor=True) new = a.get_distance(0, 1) diff = new - 0.9 * old assert abs(diff) < 10e-6 # Change Distance old = a.get_distance(0, 1)
cy = float(cy) cz = float(cz) c_mag = math.sqrt(cx * cx + cy * cy + cz * cz) coords = [] oh = [] hh = [] zncl = [] znzn = [] for j in range(7, (n_atoms + 7)): x, y, z = lines[j + (n_atoms + 7) * n].split() coords.append([a_mag * float(x), b_mag * float(y), c_mag * float(z)]) atoms = Atoms(elements, positions=coords, cell=[[ax, ay, az], [bx, by, bz], [cx, cy, cz]]) distances = np.asarray(atoms.get_all_distances(mic=True)) for i in range(n_atoms): for j in range(i, n_atoms): #because distance matrix is symmetric if ((elements[i] == 'H' and elements[j] == 'O') or (elements[i] == 'O' and elements[j] == 'H')): oh.append(distances[i][j]) if ((elements[i] == 'H' and elements[j] == 'H') or (elements[i] == 'H' and elements[j] == 'H')): hh.append(distances[i][j]) if ((elements[i] == 'Zn' and elements[j] == 'Cl') or (elements[i] == 'Cl' and elements[j] == 'Zn')): zncl.append(distances[i][j]) if ((elements[i] == 'Zn' and elements[j] == 'Zn') or (elements[i] == 'Zn' and elements[j] == 'Zn')): znzn.append(distances[i][j])
def test_cell_list(self): """Tests that the cell list implementation returns identical results with the naive calculation. """ # Cubic system: different cutoffs, smaller and larger than the system. a = 5.64 n_copies = 3 system = bulk("NaCl", crystalstructure="rocksalt", a=a, cubic=True) system *= (n_copies, n_copies, n_copies) pos = system.get_positions() all_distances_naive = system.get_all_distances() for cutoff in np.array([0.5, 1, 1.5, 2])*a*n_copies: cell_list = CellList(pos, cutoff) for idx in range(len(system)): result = cell_list.get_neighbours_for_index(idx) indices = result.indices distances = result.distances sort_order = np.argsort(indices) indices = np.array(indices)[sort_order] distances = np.array(distances)[sort_order] indices_naive = np.where(np.linalg.norm(pos-pos[idx], axis=1) <= cutoff)[0] indices_naive = indices_naive[indices_naive != idx] distances_naive = all_distances_naive[idx, indices_naive] self.assertTrue(np.array_equal(indices, indices_naive)) self.assertTrue(np.allclose(distances, distances_naive, atol=1e-16, rtol=1e-16)) # Triclinic finite system: cell > cutoff system = Atoms( cell=[ [0.0, 2.0, 2.0], [2.0, 0.0, 2.0], [2.0, 2.0, 0.0] ], positions=[ [0, 0, 0], [0.95, 0, 0], [0.95*(1+math.cos(76/180*math.pi)), 0.95*math.sin(76/180*math.pi), 0.0] ], symbols=["H", "O", "H"], ) system *= (3, 3, 3) pos = system.get_positions() all_distances_naive = system.get_all_distances() for cutoff in np.arange(1, 5): cell_list = CellList(pos, cutoff) for idx in range(len(system)): result = cell_list.get_neighbours_for_index(idx) indices = result.indices distances = result.distances sort_order = np.argsort(indices) indices = np.array(indices)[sort_order] distances = np.array(distances)[sort_order] indices_naive = np.where(np.linalg.norm(pos-pos[idx], axis=1) <= cutoff)[0] indices_naive = indices_naive[indices_naive != idx] distances_naive = all_distances_naive[idx, indices_naive] self.assertTrue(np.array_equal(indices, indices_naive)) self.assertTrue(np.allclose(distances, distances_naive, atol=1e-16, rtol=1e-16)) # System smaller than cutoff system = Atoms( positions=[[0, 0, 0]], symbols=["H"], ) pos = system.get_positions() cutoff = 5 cell_list = CellList(pos, cutoff) result = cell_list.get_neighbours_for_position(4, 3, 0) indices = result.indices distances = result.distances self.assertEqual(indices, [0]) self.assertAlmostEqual(distances[0], 5.0, places=7) result = cell_list.get_neighbours_for_position(4, 3, 0.001) indices = result.indices distances = result.distances self.assertEqual(indices, []) self.assertEqual(distances, []) # Position way outside bins system = Atoms( positions=[[0, 0, 0], [1, 1, 1]], symbols=["H", "H"], ) pos = system.get_positions() cutoff = 0.2 cell_list = CellList(pos, cutoff) result = cell_list.get_neighbours_for_position(500, 500, 500) indices = result.indices distances = result.distances self.assertEqual(indices, []) self.assertEqual(distances, []) result = cell_list.get_neighbours_for_position(-500, -500, -500) indices = result.indices distances = result.distances self.assertEqual(indices, []) self.assertEqual(distances, []) # Position at the brink of bins system = Atoms( positions=[[0, 0, 0], [1, 1, 1]], symbols=["H", "H"], ) pos = system.get_positions() cutoff = 0.2 cell_list = CellList(pos, cutoff) result = cell_list.get_neighbours_for_position(-0.2, 0, 0) indices = result.indices distances = result.distances self.assertEqual(indices, [0]) self.assertAlmostEqual(distances[0], 0.2, places=7) result = cell_list.get_neighbours_for_position(1, 1.2, 1) indices = result.indices distances = result.distances self.assertEqual(indices, [1]) self.assertAlmostEqual(distances[0], 0.2, places=7)
def test_potential_model(): """A simple example to test training and using a potential""" from ase import Atoms from ase.calculators.lj import LennardJones from pinn.io import load_numpy, sparse_batch from pinn.models import potential_model def three_body_sample(atoms, a, r): x = a * np.pi / 180 pos = [[0, 0, 0], [0, 2, 0], [0, r * np.cos(x), r * np.sin(x)]] atoms.set_positions(pos) return atoms tmp = tempfile.mkdtemp(prefix='pinn_test') atoms = Atoms('H3', calculator=LennardJones()) na, nr = 50, 50 arange = np.linspace(30, 180, na) rrange = np.linspace(1, 3, nr) # Truth agrid, rgrid = np.meshgrid(arange, rrange) egrid = np.zeros([na, nr]) for i in range(na): for j in range(nr): atoms = three_body_sample(atoms, arange[i], rrange[j]) egrid[i, j] = atoms.get_potential_energy() # Samples nsample = 50 asample, rsample = [], [] distsample = [] data = {'e_data': [], 'f_data': [], 'elems': [], 'coord': []} for i in range(nsample): a, r = np.random.choice(arange), np.random.choice(rrange) atoms = three_body_sample(atoms, a, r) dist = atoms.get_all_distances() dist = dist[np.nonzero(dist)] data['e_data'].append(atoms.get_potential_energy()) data['f_data'].append(atoms.get_forces()) data['coord'].append(atoms.get_positions()) data['elems'].append(atoms.numbers) asample.append(a) rsample.append(r) distsample.append(dist) data = {k: np.array(v) for k, v in data.items()} dataset = lambda: load_numpy(data) train = lambda: dataset()['train'].shuffle(100).repeat().apply( sparse_batch(100)) test = lambda: dataset()['test'].repeat().apply(sparse_batch(100)) params = { 'model_dir': tmp, 'network': 'pinet', 'network_params': { 'ii_nodes': [8, 8], 'pi_nodes': [8, 8], 'pp_nodes': [8, 8], 'en_nodes': [8, 8], 'rc': 3.0, 'atom_types': [1] }, 'model_params': { 'use_force': True } } model = potential_model(params) train_spec = tf.estimator.TrainSpec(input_fn=train, max_steps=200) eval_spec = tf.estimator.EvalSpec(input_fn=test, steps=10) tf.estimator.train_and_evaluate(model, train_spec, eval_spec) rmtree(tmp, ignore_errors=True)
def test_atoms_distance(): # Setup a chain of H,O,C # H-O Dist = 2 # O-C Dist = 3 # C-H Dist = 5 with mic=False # C-H Dist = 4 with mic=True a = Atoms('HOC', positions=[(1, 1, 1), (3, 1, 1), (6, 1, 1)]) a.set_cell((9, 2, 2)) a.set_pbc((True, False, False)) # Calculate indiviually with mic=True assert a.get_distance(0, 1, mic=True) == 2 assert a.get_distance(1, 2, mic=True) == 3 assert a.get_distance(0, 2, mic=True) == 4 # Calculate indiviually with mic=False assert a.get_distance(0, 1, mic=False) == 2 assert a.get_distance(1, 2, mic=False) == 3 assert a.get_distance(0, 2, mic=False) == 5 # Calculate in groups with mic=True assert (a.get_distances(0, [1, 2], mic=True) == [2, 4]).all() # Calculate in groups with mic=False assert (a.get_distances(0, [1, 2], mic=False) == [2, 5]).all() # Calculate all with mic=True assert (a.get_all_distances(mic=True) == [[0, 2, 4], [2, 0, 3], [4, 3, 0]]).all() # Calculate all with mic=False assert (a.get_all_distances(mic=False) == [[0, 2, 5], [2, 0, 3], [5, 3, 0]]).all() # Scale Distance old = a.get_distance(0, 1) a.set_distance(0, 1, 0.9, add=True, factor=True) new = a.get_distance(0, 1) diff = new - 0.9 * old assert abs(diff) < 10e-6 # Change Distance old = a.get_distance(0, 1) a.set_distance(0, 1, 0.9, add=True) new = a.get_distance(0, 1) diff = new - old - 0.9 assert abs(diff) < 10e-6