Change ring dihedral angles using RDKit

Author: JintaoWu98

arc/species/conformers.py

@@ -751,12 +751,33 @@ def change_dihedrals_and_force_field_it(label, mol, xyz, torsions, new_dihedrals
         new_dihedrals = [new_dihedrals]
 
     for dihedrals in new_dihedrals:
+        skip_to_next_dihedral = False  # Initialize a flag
         xyz_dihedrals = xyz
         for torsion, dihedral in zip(torsions, dihedrals):
             conf, rd_mol = converter.rdkit_conf_from_mol(mol, xyz_dihedrals)
             if conf is not None:
-                torsion_0_indexed = [tor - 1 for tor in torsion]
-                xyz_dihedrals = converter.set_rdkit_dihedrals(conf, rd_mol, torsion_0_indexed, deg_abs=dihedral)
+                if not isinstance(dihedral, list):
+                    torsion_0_indexed = [tor - 1 for tor in torsion]
+                    xyz_dihedrals = converter.set_rdkit_dihedrals(conf, rd_mol, torsion_0_indexed, deg_abs=dihedral)
+                elif isinstance(dihedral, list):
+                    try:
+                        torsion_0_indexed = [[tor - 0 for tor in sub_torsion] for sub_torsion in torsion]
+                        ring_length = len(torsion_0_indexed)
+                        head = torsion_0_indexed[0][0]
+                        for torsion in torsion_0_indexed:
+                            if head == torsion[-1]:
+                                tail = torsion[-2]
+                                break
+                        xyz_dihedrals = converter.set_rdkit_ring_dihedrals(
+                            conf, rd_mol, head, tail, torsion_0_indexed[0:ring_length - 3], dihedral[0:ring_length - 3]
+                        )
+                    except Exception as e:  # Catch exceptions specifically from set_rdkit_ring_dihedrals
+                        print(f"Skipping ring dihedral due to an error: {e}")
+                        skip_to_next_dihedral = True  # Set the flag to skip this dihedral set
+                        break  # Exit the inner loop for this dihedral set
+        if skip_to_next_dihedral:
+            continue  # Skip the rest of this `dihedrals` iteration
+
         xyz_, energy = get_force_field_energies(label, mol=mol, xyz=xyz_dihedrals, optimize=True,
                                                 force_field=force_field, suppress_warning=True)
         if energy and xyz_:

arc/species/converter.py

@@ -12,7 +12,7 @@
 from openbabel import pybel
 from rdkit import Chem
 from rdkit.Chem import rdMolTransforms as rdMT
-from rdkit.Chem import SDWriter
+from rdkit.Chem import SDWriter, AllChem
 from rdkit.Chem.rdchem import AtomValenceException
 
 from arkane.common import get_element_mass, mass_by_symbol, symbol_by_number
@@ -1851,6 +1851,104 @@ def set_rdkit_dihedrals(conf, rd_mol, torsion, deg_increment=None, deg_abs=None)
     new_xyz = xyz_from_data(coords=coords, symbols=symbols)
     return new_xyz
 
+def set_rdkit_ring_dihedrals(conf_original, rd_mol, ring_head, ring_tail, torsions, dihedrals):
+    """
+    A helper function for setting dihedral angles in a ring using RDKit.
+
+    Args:
+        rd_mol: The respective RDKit molecule.
+        ring_head: The first atom index of the ring(0-indexed).
+        ring_tail: The last atom index of the ring(0-indexed).
+        torsions: A list of torsions, each corresponding to a dihedral.
+        dihedrals: A list of dihedral angles in degrees, each corresponding to a torsion.
+
+    Example of a 6-membered ring:
+        ring_head = 0
+        ring_tail = 5
+        torsions = [(0, 1, 2, 3), (1, 2, 3, 4), (2, 3, 4, 5)]
+        dihedrals = [30, 300, 30]
+    
+    Returns:
+        dict: The xyz with the new dihedral, ordered according to the map.
+    """
+
+    # Create a copy of the molecule to modify
+    rd_mol_mod = Chem.Mol(rd_mol)
+
+    # Apply the original coordinates to the conformer
+    conf_mod = rd_mol_mod.GetConformer()
+    for i in range(rd_mol_mod.GetNumAtoms()):
+        pos = conf_original.GetAtomPosition(i)
+        conf_mod.SetAtomPosition(i, pos)
+    # Remove hydrogens
+    rd_mol_noH = Chem.RemoveHs(rd_mol_mod)
+    Chem.SanitizeMol(rd_mol_noH)
+    
+    # Map positions from conf_mod to conf_noH
+    conf_noH = Chem.Conformer(rd_mol_noH.GetNumAtoms())
+    atom_map = {}  # Map heavy atom indices between rd_mol_mod and rd_mol_noH
+    idx_noH = 0
+    for idx in range(rd_mol_mod.GetNumAtoms()):
+        atom = rd_mol_mod.GetAtomWithIdx(idx)
+        if atom.GetAtomicNum() != 1:  # Not hydrogen
+            pos = conf_mod.GetAtomPosition(idx)
+            conf_noH.SetAtomPosition(idx_noH, pos)
+            atom_map[idx] = idx_noH
+            idx_noH += 1
+    rd_mol_noH.AddConformer(conf_noH)
+
+    # Remove the bond to open the ring
+    rd_mol_noH = Chem.RWMol(rd_mol_noH)
+    rd_mol_noH.RemoveBond(atom_map[ring_head], atom_map[ring_tail])
+    Chem.SanitizeMol(rd_mol_noH)
+
+    # Set the specified dihedral angles
+    conf_noH = rd_mol_noH.GetConformer()
+    for torsion, dihedral in zip(torsions, dihedrals):
+        torsion_noH = [atom_map[atom_idx] for atom_idx in torsion]
+        rdMT.SetDihedralDeg(conf_noH, *torsion_noH, dihedral)
+    # Re-add the bond to close the ring
+    rd_mol_noH.AddBond(
+        atom_map[ring_head], atom_map[ring_tail], rd_mol.GetBondBetweenAtoms(ring_head, ring_tail).GetBondType()
+    )
+    Chem.SanitizeMol(rd_mol_noH)
+    # Optimize the molecule
+    uff_ff = AllChem.UFFGetMoleculeForceField(rd_mol_noH)
+    if uff_ff is None:
+        raise ValueError("UFF force field could not be generated for the molecule.")
+
+    # Add torsion constraints to keep dihedral angles fixed
+    force_constant = 1000.0  # A high force constant to strongly enforce the constraint
+    for torsion, dihedral in zip(torsions, dihedrals):
+        torsion_noH = [atom_map[atom_idx] for atom_idx in torsion]
+        i, j, k, l = torsion_noH
+        uff_ff.UFFAddTorsionConstraint(
+            i, j, k, l, False, dihedral, dihedral, force_constant
+        )
+
+    # Optimize the molecule
+    uff_ff.Minimize()
+
+    # Retrieve the optimized conformer
+    conf_noH = rd_mol_noH.GetConformer()
+    # Add hydrogens back to the optimized molecule
+    rd_mol_opt_H = Chem.AddHs(rd_mol_noH)
+    # Generate new conformer with hydrogens
+    AllChem.EmbedMolecule(rd_mol_opt_H, coordMap={atom.GetIdx(): conf_noH.GetAtomPosition(atom.GetIdx()) for atom in rd_mol_noH.GetAtoms()})
+    AllChem.UFFOptimizeMolecule(rd_mol_opt_H)
+    # Extract updated coordinates
+    conf_opt_H = rd_mol_opt_H.GetConformer()
+    coords = []
+    symbols = []
+    for i, atom in enumerate(rd_mol_opt_H.GetAtoms()):
+        pos = conf_opt_H.GetAtomPosition(i)
+        coords.append([pos.x, pos.y, pos.z])
+        symbols.append(atom.GetSymbol())
+
+    # Create the new xyz dictionary
+    new_xyz = xyz_from_data(coords=coords, symbols=symbols)
+    return new_xyz
+
 
 def check_isomorphism(mol1, mol2, filter_structures=True, convert_to_single_bonds=False):
     """

arc/species/converter_test.py

@@ -4066,6 +4066,43 @@ def test_set_rdkit_dihedrals(self):
 H       2.16336803    0.09985803    0.03295192"""
         self.assertEqual(converter.xyz_to_str(new_xyz4), expected_xyz4)
 
+    def test_set_rdkit_ring_dihedrals(self):
+        """Test setting the dihedral angles of an RDKit ring molecule"""
+        xyz_original = """C       1.17528959    0.88689342   -0.09425887
+C      -0.23165323    1.40815606   -0.37444021
+C      -1.28915380    0.60789983    0.38119602
+C      -1.17528947   -0.88689346    0.09425817
+C       0.23165279   -1.40815571    0.37444068
+C       1.28915350   -0.60789979   -0.38119592
+H       1.90063595    1.43610053   -0.70501194
+H       1.43190556    1.07695419    0.95523181
+H      -0.29672067    2.46483469   -0.09164586
+H      -0.43309289    1.35229514   -1.45133707
+H      -2.28822258    0.96189799    0.10312701
+H      -1.17664390    0.78164432    1.45848873
+H      -1.43190253   -1.07695588   -0.95523291
+H      -1.90063264   -1.43610606    0.70501042
+H       0.29671416   -2.46483479    0.09164748
+H       0.43309139   -1.35229454    1.45133785
+H       1.17664469   -0.78164459   -1.45848883
+H       2.28822409   -0.96189136   -0.10312655"""
+        xyz_original = converter.str_to_xyz(xyz_original)
+        
+        ring_head = 0
+        ring_tail = 5
+        torsions = [(0, 1, 2, 3), (1, 2, 3, 4), (2, 3, 4, 5)]
+        dihedrals = [29.167577928701704, 299.8936870462789, 29.167577208303104]
+            
+        s_mol, b_mol = converter.molecules_from_xyz(xyz_original)
+        mol = b_mol if b_mol is not None else s_mol
+        _, rd_mol = converter.rdkit_conf_from_mol(mol, xyz_original)
+            
+        xyz_final = converter.set_rdkit_ring_dihedrals(rd_mol, ring_head, ring_tail, torsions, dihedrals)
+        
+        self.assertAlmostEqual(calculate_dihedral_angle(xyz_final,[0,1,2,3]), 29.167577928701704, 2)
+        self.assertAlmostEqual(calculate_dihedral_angle(xyz_final,[1,2,3,4]), 299.8936870462789,  2)
+        self.assertAlmostEqual(calculate_dihedral_angle(xyz_final,[2,3,4,5]), 29.167577208303104, 2)
+
     def test_get_center_of_mass(self):
         """Test calculating the center of mass for coordinates"""
         xyz = """O       1.28706525    0.52121353    0.04219198