Define set_rdkit_ring_dihedrals and make test

arc/plotter.py

@@ -1014,12 +1014,15 @@
     """
     num_comb = None
     torsions = list(torsion_angles.keys()) if torsions_sampling_points is None \
-        else list(torsions_sampling_points.keys())
+        else list(torsions_sampling_points.keys())    
+    torsions = [torsion for torsion in torsions if not (isinstance(torsion, tuple) and all(isinstance(sub_t, tuple) for sub_t in torsion))]
+
     ticks = [0, 60, 120, 180, 240, 300, 360]
     sampling_points = dict()
     if torsions_sampling_points is not None:
         for tor, points in torsions_sampling_points.items():
-            sampling_points[tor] = [point if point <= 360 else point - 360 for point in points]
+            if not isinstance(points[0],list):
+                sampling_points[tor] = [point if point <= 360 else point - 360 for point in points]
     if not torsions:
         return
     if len(torsions) == 1:
@@ -1048,6 +1051,8 @@
         fig.dpi = 120
         num_comb = 1
         for i, torsion in enumerate(torsions):
+            if tuple(torsion) not in torsion_angles:
+                continue
             axs[i].plot(np.array(torsion_angles[tuple(torsion)]),
                         np.zeros_like(np.arange(len(torsion_angles[tuple(torsion)]))), 'g.')
             if wells_dict is not None:
@@ -1121,6 +1126,76 @@
     return num_comb
 
 
+def plot_ring_torsion_angles(conformers, plot_path=None, tolerance=15):
+    """
+    Plot the torsion angles of the generated conformers for each ring,
+    considering torsion similarity within a given tolerance.
+
+    Args:
+        conformers (list): A list of conformers, each containing a 'puckering' key with ring torsion angles.
+        plot_path (str, optional): The path for saving the plot.
+        tolerance (float, optional): The angular tolerance to consider two torsion angle sets as similar.
+    """
+    if 'puckering' not in conformers[0]:
+        return
+
+    # Dictionary to store unique angle sets for each ring
+    ring_angle_data = {}
+
+    # Process each conformer
+    for conformer in conformers:
+        rings = conformer['puckering']  # Retrieve the puckering angles for rings
+        for torsions, angle_set in rings.items():
+            rounded_angle_set = tuple(round(angle) for angle in angle_set)  # Round angles
+            if torsions not in ring_angle_data:
+                ring_angle_data[torsions] = []
+
+            # Check for similarity within the current ring
+            is_similar = False
+            for i, (existing_set, count) in enumerate(ring_angle_data[torsions]):
+                if all(abs(a1 - a2) <= tolerance for a1, a2 in zip(rounded_angle_set, existing_set)):
+                    # If similar, increment count
+                    ring_angle_data[torsions][i] = (existing_set, count + 1)
+                    is_similar = True
+                    break
+            if not is_similar:
+                # Add unique angle set with a count
+                ring_angle_data[torsions].append((rounded_angle_set, 1))
+
+
+    # Plot data for each ring
+    for ring, angle_counts in ring_angle_data.items():
+        # Extract and sort data
+        angles, counts = zip(*angle_counts)
+        angles_counts_sorted = sorted(zip(angles, counts), key=lambda x: x[1], reverse=True)
+        angles_sorted, counts_sorted = zip(*angles_counts_sorted)
+
+        # Create bar plot for this ring
+        fig, ax = plt.subplots(figsize=(10, 5))
+        x = np.arange(len(angles_sorted))  # Label positions
+        ax.bar(x, counts_sorted, color='blue')
+        ax.set_xlabel('Rounded Angle Sets (Nearest Integer)')
+        ax.set_ylabel('Frequency')
+        ax.set_title(f'Frequency of Different Angle Sets for Ring {ring}')
+        ax.set_xticks(x)
+        ax.set_xticklabels([f'{angle}' for angle in angles_sorted], rotation=45, ha='right')
+
+        # Save or display the plot
+        if plot_path is not None:
+            ring_plot_path = os.path.join(plot_path, f'conformer_ring_torsions_{ring}.png')
+            if not os.path.isdir(plot_path):
+                os.makedirs(plot_path)
+            try:
+                plt.savefig(ring_plot_path, bbox_inches='tight')
+            except FileNotFoundError:
+                pass
+        if is_notebook():
+            plt.show()
+        plt.close(fig)
+
+    return ring_angle_data
+
+
 def plot_1d_rotor_scan(angles: Optional[Union[list, tuple, np.array]] = None,
                        energies: Optional[Union[list, tuple, np.array]] = None,
                        results: Optional[dict] = None,

arc/species/conformers.py

@@ -269,9 +269,12 @@
         mol_list=mol_list, label=label, xyzs=xyzs, torsion_num=len(torsions), charge=charge, multiplicity=multiplicity,
         num_confs=num_confs_to_generate, force_field=force_field)
 
+    rings, rings_indices = determine_rings(mol_list)
     lowest_confs = list()
     if len(conformers):
         conformers = determine_dihedrals(conformers, torsions)
+        if len(rings):
+            conformers = determine_puckering(conformers, rings_indices)
 
         new_conformers, symmetries = deduce_new_conformers(
             label, conformers, torsions, tops, mol_list, smeared_scan_res, plot_path=plot_path,
@@ -345,18 +348,25 @@
         symmetries[tuple(torsion)] = symmetry
     logger.debug(f'Identified {len([s for s in symmetries.values() if s > 1])} symmetric wells for {label}')
 
-    torsions_sampling_points, wells_dict = dict(), dict()
+    torsions_sampling_points, wells_dict, ring_sampling_points = dict(), dict(), dict()
     for tor, tor_angles in torsion_angles.items():
         torsions_sampling_points[tor], wells_dict[tor] = \
             determine_torsion_sampling_points(label, tor_angles, smeared_scan_res=smeared_scan_res,
                                               symmetry=symmetries[tor])
-
     if plot_path is not None:
         arc.plotter.plot_torsion_angles(torsion_angles, torsions_sampling_points, wells_dict=wells_dict,
                                         plot_path=plot_path)
+        ring_angle_data = arc.plotter.plot_ring_torsion_angles(conformers=conformers, plot_path=plot_path)
+    # Process the ring angle data
+    if ring_angle_data:
+        for ring, angle_counts in ring_angle_data.items():
+            angles = [list(angle) for angle, _ in angle_counts]
+            ring_sampling_points[tuple(ring)] = angles
+
+    combined_sampling_points = {**torsions_sampling_points, **ring_sampling_points}
 
     hypothetical_num_comb = 1
-    for points in torsions_sampling_points.values():
+    for points in combined_sampling_points.values():
         hypothetical_num_comb *= len(points)
     number_of_chiral_centers = get_number_of_chiral_centers(label, mol, conformer=conformers[0],
                                                             just_get_the_number=True)
@@ -367,10 +377,10 @@
         hypothetical_num_comb_str = str(hypothetical_num_comb)
     logger.info(f'\nHypothetical number of conformer combinations for {label}: {hypothetical_num_comb_str}')
 
-    # split torsions_sampling_points into two lists, use combinations only for those with multiple sampling points
+    # split combined_sampling_points into two lists, use combinations only for those with multiple sampling points
     single_tors, multiple_tors, single_sampling_point, multiple_sampling_points = list(), list(), list(), list()
     multiple_sampling_points_dict = dict()  # used for plotting an energy "scan"
-    for tor, points in torsions_sampling_points.items():
+    for tor, points in combined_sampling_points.items():
         if len(points) == 1:
             single_tors.append(tor)
             single_sampling_point.append((points[0]))
@@ -533,7 +543,8 @@
                                  'FF energy': round(energy, 3),
                                  'source': f'Changing dihedrals on most stable conformer, iteration {i}',
                                  'torsion': tor,
-                                 'dihedral': round(dihedral, 2),
+                                 'dihedral': round(dihedral, 2) if isinstance(dihedral, float)
+                                             else [round(angle, 2) for angle in dihedral],
                                  'dmat': dmat1,
                                  'fl_distance': fl_distance1}
                     newest_conformers_dict[tor].append(conformer)
@@ -549,7 +560,7 @@
                          'FF energy': None,
                          'source': f'Changing dihedrals on most stable conformer, iteration {i}, but FF energy is None',
                          'torsion': tor,
-                         'dihedral': round(dihedral, 2),
+                         'dihedral':  round(dihedral, 2) if isinstance(dihedral, float) else [round(angle, 2) for angle in dihedral],
                          'dmat': dmat1,
                          'fl_distance': fl_distance1})
         new_conformers.extend(newest_conformer_list)
@@ -748,12 +759,33 @@
         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_:
@@ -887,6 +919,49 @@
     return conformers
 
 
+def determine_rings(mol_list):
+    """
+    Determine the rings in the molecule.
+
+    Args:
+        mol_list (list): Localized structures (Molecule objects) by which all rotors will be determined.
+
+    Returns:
+        Tuple[list, list]:
+            - A list of ring atoms.
+            - A list of ring atom indices.
+    """
+    rings, rings_indexes = list(), list()
+    for mol in mol_list:
+        rings = mol.get_deterministic_sssr()
+        rings_indexes = [[mol.atoms.index(atom) for atom in ring] for ring in rings]
+    return rings, rings_indexes
+
+
+def determine_puckering(conformers, rings_indices):
+    """
+    For each conformer in `conformers` determine the respective puckering angles.
+
+    Args:
+        conformers (list): Entries are conformer dictionaries.
+        rings_indices (list): Entries are lists of ring atom indices.
+
+    Returns:
+        list: Entries are conformer dictionaries.
+    """
+    for conformer in conformers:
+        if isinstance(conformer['xyz'], str):
+            xyz = converter.str_to_xyz(conformer['xyz'])
+        else:
+            xyz = conformer['xyz']
+        if 'puckering' not in conformer or not conformer['puckering']:
+            conformer['puckering'] = dict()
+            for i, ring in enumerate(rings_indices):
+                theta = vectors.calculate_ring_dihedral_angles(coords=xyz['coords'], ring=ring, index=0)
+                conformer['puckering'][tuple((ring[i%len(ring)], ring[(i + 1)%len(ring)], ring[(i + 2)%len(ring)], ring[(i + 3)%len(ring)]) for i in range(len(ring)))] = theta    
+    return conformers
+
+
 def determine_torsion_sampling_points(label, torsion_angles, smeared_scan_res=None, symmetry=1):
     """
     Determine how many points to consider in each well of a torsion for conformer combinations.

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 @@
     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):
     """