From b016dce6f351494bad33cdb4e26500c328619224 Mon Sep 17 00:00:00 2001
From: Jintao <jintaowu@campus.technion.ac.il>
Date: Sun, 8 Dec 2024 10:59:29 +0800
Subject: [PATCH] Define ring and puckering info for confs

---
 arc/species/conformers.py | 46 +++++++++++++++++++++++++++++++++++++++
 arc/species/vectors.py    | 21 +++++++++++++++++-
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/arc/species/conformers.py b/arc/species/conformers.py
index bc077cb801..a3240288f9 100644
--- a/arc/species/conformers.py
+++ b/arc/species/conformers.py
@@ -269,9 +269,12 @@ def generate_conformers(mol_list: Union[List[Molecule], Molecule],
         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,
@@ -887,6 +890,49 @@ def determine_dihedrals(conformers, torsions):
     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.
diff --git a/arc/species/vectors.py b/arc/species/vectors.py
index 99fe05e5b7..71ddd07fb1 100644
--- a/arc/species/vectors.py
+++ b/arc/species/vectors.py
@@ -205,7 +205,7 @@ def calculate_dihedral_angle(coords: Union[list, tuple, dict],
     """
     if isinstance(coords, dict) and 'coords' in coords:
         coords = coords['coords']
-    if not isinstance(coords, (list, tuple)):
+    if not isinstance(coords, (list, tuple, np.ndarray)):
         raise TypeError(f'coords must be a list or a tuple, got\n{coords}\nwhich is a {type(coords)}')
     if index not in [0, 1]:
         raise VectorsError(f'index must be either 0 or 1, got {index}')
@@ -232,6 +232,25 @@ def calculate_dihedral_angle(coords: Union[list, tuple, dict],
     return get_dihedral(v1, v2, v3, units=units)
 
 
+def calculate_ring_dihedral_angles(coords: Union[list, tuple, dict], 
+                                   ring: list, 
+                                   index: int = 0
+                                   ) -> list:
+    if isinstance(coords, dict) and 'coords' in coords:
+        coords = coords['coords']
+    if not isinstance(coords, (list, tuple)):
+        raise TypeError(f'coords must be a list or a tuple, got {type(coords)}')
+
+    coords = np.array(coords, dtype=np.float32)
+    ring = [atom - index for atom in ring]  # Adjusting for zero-indexed
+    angles = []
+    for i in range(len(ring)):
+        angle_deg = calculate_dihedral_angle(coords, [ring[i%len(ring)], ring[(i + 1)%len(ring)], ring[(i + 2)%len(ring)], ring[(i + 3)%len(ring)]])
+        angles.append(angle_deg)
+
+    return angles
+
+
 def calculate_param(coords: Union[list, tuple, dict],
                     atoms: list,
                     index: int = 0,