MultiOptPy 1.20.7__py3-none-any.whl → 1.20.8__py3-none-any.whl

multioptpy/Potential/keep_dihedral_angle_potential.py

@@ -1,95 +1,294 @@
-
 from multioptpy.Parameters.parameter import UnitValueLib
 from multioptpy.Utils.calc_tools import torch_calc_dihedral_angle_from_vec
 import torch
+import math
 
 class StructKeepDihedralAnglePotential:
+    """
+    Computes the harmonic potential energy for a dihedral angle defined by four atoms.
+    
+    This class incorporates a robust singularity handling mechanism for cases where
+    three consecutive atoms become collinear (i.e., the bond angle approaches 0 or 180 degrees).
+    In such configurations, the dihedral angle is mathematically undefined, leading to 
+    numerical instabilities and exploding gradients.
+
+    To resolve this, a smooth switching function (smoothstep) is applied to the squared norm 
+    of the cross products of the bond vectors. This function smoothly attenuates the 
+    potential energy and forces to zero as the geometry enters the collinear region.
+
+    Energy Function:
+        E_total = E_harmonic * S(|n1|^2) * S(|n2|^2)
+    
+    Where:
+        E_harmonic = 0.5 * k * (phi - phi_0)^2
+        S(x) is a cubic Hermite interpolation function (smoothstep) mapping [0, 1].
+    """
+
     def __init__(self, **kwarg):
         self.config = kwarg
         UVL = UnitValueLib()
         self.hartree2kcalmol = UVL.hartree2kcalmol 
         self.bohr2angstroms = UVL.bohr2angstroms 
         self.hartree2kjmol = UVL.hartree2kjmol 
+        
+        # Thresholds for collinearity smoothing (Squared Norm of Cross Product).
+        # These values determine the range over which the potential is switched off.
+        # COLLINEAR_CUT_MIN: Below this value, the weighting factor is 0.0.
+        # COLLINEAR_CUT_MAX: Above this value, the weighting factor is 1.0.
+        # A squared norm of 1e-8 corresponds to a sine of approx 1e-4.
+        self.COLLINEAR_CUT_MIN = 1e-10
+        self.COLLINEAR_CUT_MAX = 1e-8
+        
         return
-    def calc_energy(self, geom_num_list, bias_pot_params=[]):
+
+    def _compute_switching(self, val):
         """
-        # required variables: self.config["keep_dihedral_angle_spring_const"],
-                              self.config["keep_dihedral_angle_atom_pairs"]
-                              self.config["keep_dihedral_angle_angle"]
-        bias_pot_params[0] : keep_dihedral_angle_spring_const
-        bias_pot_params[1] : keep_dihedral_angle_angle
+        Computes the smoothstep switching factor.
+
+        Args:
+            val (Tensor): The squared norm of the cross product vector.
+
+        Returns:
+            Tensor: A scalar factor between 0.0 and 1.0 using cubic interpolation.
+                    Returns 0.0 if val < min, 1.0 if val > max.
+        """
+        t = (val - self.COLLINEAR_CUT_MIN) / (self.COLLINEAR_CUT_MAX - self.COLLINEAR_CUT_MIN)
+        t = torch.clamp(t, 0.0, 1.0)
+        # Smoothstep function: 3t^2 - 2t^3
+        return t * t * (3.0 - 2.0 * t)
+
+    def calc_energy(self, geom_num_list, bias_pot_params=[]):
         """
-        a1 = geom_num_list[self.config["keep_dihedral_angle_atom_pairs"][1]-1] - geom_num_list[self.config["keep_dihedral_angle_atom_pairs"][0]-1]
-        a2 = geom_num_list[self.config["keep_dihedral_angle_atom_pairs"][2]-1] - geom_num_list[self.config["keep_dihedral_angle_atom_pairs"][1]-1]
-        a3 = geom_num_list[self.config["keep_dihedral_angle_atom_pairs"][3]-1] - geom_num_list[self.config["keep_dihedral_angle_atom_pairs"][2]-1]
+        Calculates the potential energy with collinearity smoothing.
 
-        angle = torch.abs(torch_calc_dihedral_angle_from_vec(a1, a2, a3))
+        Args:
+            geom_num_list (Tensor): Atomic coordinates tensor of shape (N_atoms, 3).
+            bias_pot_params (list, optional): Optional override for parameters [k, phi_0].
+
+        Returns:
+            Tensor: The calculated potential energy (scalar).
+        """
+        
+        # ========================================
+        # 1. Parameter Retrieval
+        # ========================================
         if len(bias_pot_params) == 0:
-            energy = 0.5 * self.config["keep_dihedral_angle_spring_const"] * (angle - torch.deg2rad(torch.tensor(self.config["keep_dihedral_angle_angle"]))) ** 2
+            k = self.config["keep_dihedral_angle_spring_const"]
+            phi_0_deg = torch.tensor(self.config["keep_dihedral_angle_angle"])
+            phi_0 = torch.deg2rad(phi_0_deg)
         else:
-            energy = 0.5 * bias_pot_params[0] * (angle - torch.deg2rad(bias_pot_params[1])) ** 2
-      
-        return energy #hartree    
+            k = bias_pot_params[0]
+            phi_0_deg = bias_pot_params[1]
+            if isinstance(phi_0_deg, torch.Tensor):
+                phi_0 = torch.deg2rad(phi_0_deg)
+            else:
+                phi_0 = torch.deg2rad(torch.tensor(phi_0_deg))
+
+        device = geom_num_list.device if isinstance(geom_num_list, torch.Tensor) else torch.device("cpu")
+        dtype = geom_num_list.dtype
+        
+        PI = torch.tensor(math.pi, device=device, dtype=dtype)
+        phi_0 = phi_0.to(device=device, dtype=dtype)
+
+        # ========================================
+        # 2. Vector Calculations
+        # ========================================
+        i1 = self.config["keep_dihedral_angle_atom_pairs"][0] - 1
+        i2 = self.config["keep_dihedral_angle_atom_pairs"][1] - 1
+        i3 = self.config["keep_dihedral_angle_atom_pairs"][2] - 1
+        i4 = self.config["keep_dihedral_angle_atom_pairs"][3] - 1
+        
+        b1 = geom_num_list[i2] - geom_num_list[i1]
+        b2 = geom_num_list[i3] - geom_num_list[i2]
+        b3 = geom_num_list[i4] - geom_num_list[i3]
+        
+        # Normal vectors to the planes defined by bond pairs
+        n1 = torch.linalg.cross(b1, b2)
+        n2 = torch.linalg.cross(b2, b3)
+        
+        # ========================================
+        # 3. Collinearity Guard (Singularity Smoothing)
+        # ========================================
+        n1_sq_norm = torch.sum(n1**2, dim=-1)
+        n2_sq_norm = torch.sum(n2**2, dim=-1)
+        
+        # Compute switching factors to dampen energy in collinear regions
+        switch_1 = self._compute_switching(n1_sq_norm)
+        switch_2 = self._compute_switching(n2_sq_norm)
+        
+        # Safe normalization for angle calculation.
+        # Clamping prevents NaN in the graph, while the switching factor ensures
+        # that the resulting gradients in the clamped region are zeroed out.
+        n1_norm = torch.clamp(torch.sqrt(n1_sq_norm), min=1e-12)
+        n2_norm = torch.clamp(torch.sqrt(n2_sq_norm), min=1e-12)
+        
+        n1_hat = n1 / (n1_norm.unsqueeze(-1))
+        n2_hat = n2 / (n2_norm.unsqueeze(-1))
+        
+        b2_norm = torch.clamp(torch.linalg.norm(b2), min=1e-12)
+        b2_hat = b2 / (b2_norm.unsqueeze(-1))
+        
+        # ========================================
+        # 4. Angle Calculation
+        # ========================================
+        x = torch.sum(n1_hat * n2_hat, dim=-1)
+        # (n1 x n2) is parallel to b2
+        m1 = torch.linalg.cross(n1_hat, n2_hat)
+        y = torch.sum(m1 * b2_hat, dim=-1)
+        
+        phi = torch.atan2(y, x)
+        
+        # ========================================
+        # 5. Energy Calculation with Smoothing
+        # ========================================
+        diff = phi - phi_0
+        # Wrap difference to [-pi, pi]
+        diff = diff - 2.0 * PI * torch.round(diff / (2.0 * PI))
+        
+        raw_energy = 0.5 * k * diff**2
+        
+        # Apply smoothing factors
+        energy = raw_energy * switch_1 * switch_2
+        
+        return energy
 
 
 class StructKeepDihedralAnglePotentialv2:
+    """
+    Computes the dihedral angle potential energy defined by the centroids of four atom fragments.
+    
+    This class is designed for coarse-grained constraints where the dihedral is defined
+    by the geometric centers of specified groups of atoms rather than single atoms.
+    It includes the same singularity smoothing mechanism as the standard atom-based potential
+    to handle cases where the fragment centroids become collinear.
+    """
+
     def __init__(self, **kwarg):
         self.config = kwarg
         UVL = UnitValueLib()
         self.hartree2kcalmol = UVL.hartree2kcalmol 
         self.bohr2angstroms = UVL.bohr2angstroms 
         self.hartree2kjmol = UVL.hartree2kjmol 
+        
+        # Thresholds for collinearity smoothing
+        self.COLLINEAR_CUT_MIN = 1e-10
+        self.COLLINEAR_CUT_MAX = 1e-8
+        
         return
-    def calc_energy(self, geom_num_list, bias_pot_params=[]):
+
+    def _compute_switching(self, val):
         """
-        # required variables: self.config["keep_dihedral_angle_v2_spring_const"], 
-                             self.config["keep_dihedral_angle_v2_angle"], 
-                             self.config["keep_dihedral_angle_v2_fragm1"],
-                             self.config["keep_dihedral_angle_v2_fragm2"],
-                             self.config["keep_dihedral_angle_v2_fragm3"],
-                             self.config["keep_dihedral_angle_v2_fragm4"],
-        bias_pot_params[0] : keep_dihedral_angle_v2_spring_const
-        bias_pot_params[1] : keep_dihedral_angle_v2_angle                  
+        Computes the smoothstep switching factor.
         """
-        fragm_1_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_dihedral_angle_v2_fragm1"]) - 1], dim=0)
-        fragm_2_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_dihedral_angle_v2_fragm2"]) - 1], dim=0)
-        fragm_3_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_dihedral_angle_v2_fragm3"]) - 1], dim=0)
-        fragm_4_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_dihedral_angle_v2_fragm4"]) - 1], dim=0)
-              
-        a1 = fragm_2_center - fragm_1_center
-        a2 = fragm_3_center - fragm_2_center
-        a3 = fragm_4_center - fragm_3_center
+        t = (val - self.COLLINEAR_CUT_MIN) / (self.COLLINEAR_CUT_MAX - self.COLLINEAR_CUT_MIN)
+        t = torch.clamp(t, 0.0, 1.0)
+        return t * t * (3.0 - 2.0 * t)
 
-        angle = torch.abs(torch_calc_dihedral_angle_from_vec(a1, a2, a3))
+    def calc_energy(self, geom_num_list, bias_pot_params=[]):
+        """
+        Calculates the potential energy for fragment centroids with collinearity smoothing.
+        """
         if len(bias_pot_params) == 0:
-            energy = 0.5 * self.config["keep_dihedral_angle_v2_spring_const"] * (angle - torch.deg2rad(torch.tensor(self.config["keep_dihedral_angle_v2_angle"]))) ** 2
+            k = self.config["keep_dihedral_angle_v2_spring_const"]
+            phi_0_deg = torch.tensor(self.config["keep_dihedral_angle_v2_angle"])
+            phi_0 = torch.deg2rad(phi_0_deg)
         else:
-            energy = 0.5 * bias_pot_params[0] * (angle - torch.deg2rad(bias_pot_params[1])) ** 2
-        return energy #hartree
+            k = bias_pot_params[0]
+            phi_0_deg = bias_pot_params[1]
+            if isinstance(phi_0_deg, torch.Tensor):
+                phi_0 = torch.deg2rad(phi_0_deg)
+            else:
+                phi_0 = torch.deg2rad(torch.tensor(phi_0_deg))
+
+        device = geom_num_list.device if isinstance(geom_num_list, torch.Tensor) else torch.device("cpu")
+        dtype = geom_num_list.dtype
+        
+        PI = torch.tensor(math.pi, device=device, dtype=dtype)
+        phi_0 = phi_0.to(device=device, dtype=dtype)
+
+        # Vector Calculations for Fragment Centers
+        def get_indices(key):
+            return torch.tensor(self.config[key], device=device, dtype=torch.long) - 1
+
+        fragm_1_center = torch.mean(geom_num_list[get_indices("keep_dihedral_angle_v2_fragm1")], dim=0)
+        fragm_2_center = torch.mean(geom_num_list[get_indices("keep_dihedral_angle_v2_fragm2")], dim=0)
+        fragm_3_center = torch.mean(geom_num_list[get_indices("keep_dihedral_angle_v2_fragm3")], dim=0)
+        fragm_4_center = torch.mean(geom_num_list[get_indices("keep_dihedral_angle_v2_fragm4")], dim=0)
+           
+        b1 = fragm_2_center - fragm_1_center
+        b2 = fragm_3_center - fragm_2_center
+        b3 = fragm_4_center - fragm_3_center
+
+        n1 = torch.linalg.cross(b1, b2)
+        n2 = torch.linalg.cross(b2, b3)
+
+        # Collinearity Guard (Smoothing)
+        n1_sq_norm = torch.sum(n1**2, dim=-1)
+        n2_sq_norm = torch.sum(n2**2, dim=-1)
+        
+        switch_1 = self._compute_switching(n1_sq_norm)
+        switch_2 = self._compute_switching(n2_sq_norm)
+        
+        n1_norm = torch.clamp(torch.sqrt(n1_sq_norm), min=1e-12)
+        n2_norm = torch.clamp(torch.sqrt(n2_sq_norm), min=1e-12)
+        
+        n1_hat = n1 / (n1_norm.unsqueeze(-1))
+        n2_hat = n2 / (n2_norm.unsqueeze(-1))
+        
+        b2_norm = torch.clamp(torch.linalg.norm(b2), min=1e-12)
+        b2_hat = b2 / (b2_norm.unsqueeze(-1))
+
+        # Angle Calculation
+        x = torch.sum(n1_hat * n2_hat, dim=-1)
+        m1 = torch.linalg.cross(n1_hat, n2_hat)
+        y = torch.sum(m1 * b2_hat, dim=-1)
+        
+        phi = torch.atan2(y, x)
+
+        # Energy Calculation
+        diff = phi - phi_0
+        diff = diff - 2.0 * PI * torch.round(diff / (2.0 * PI))
+        
+        raw_energy = 0.5 * k * diff**2
+        
+        # Apply Smoothing
+        energy = raw_energy * switch_1 * switch_2
+        
+        return energy
     
 class StructKeepDihedralAnglePotentialCos:
+    """
+    Computes a cosine-based dihedral potential energy for fragment centroids.
+    
+    Includes singularity smoothing to prevent gradient explosion when fragment centers
+    become collinear.
+    """
     def __init__(self, **kwarg):
         self.config = kwarg
         UVL = UnitValueLib()
         self.hartree2kcalmol = UVL.hartree2kcalmol 
         self.bohr2angstroms = UVL.bohr2angstroms 
         self.hartree2kjmol = UVL.hartree2kjmol 
+        
+        self.COLLINEAR_CUT_MIN = 1e-10
+        self.COLLINEAR_CUT_MAX = 1e-8
         return
+
+    def _compute_switching(self, val):
+        """
+        Computes the smoothstep switching factor.
+        """
+        t = (val - self.COLLINEAR_CUT_MIN) / (self.COLLINEAR_CUT_MAX - self.COLLINEAR_CUT_MIN)
+        t = torch.clamp(t, 0.0, 1.0)
+        return t * t * (3.0 - 2.0 * t)
+
     def calc_energy(self, geom_num_list, bias_pot_params=[]):
         """
-        # required variables: self.config["keep_dihedral_angle_cos_potential_const"],
-                             self.config["keep_dihedral_angle_cos_angle_const"], 
-                             self.config["keep_dihedral_angle_cos_angle"], 
-                             self.config["keep_dihedral_angle_cos_fragm1"],
-                             self.config["keep_dihedral_angle_cos_fragm2"],
-                             self.config["keep_dihedral_angle_cos_fragm3"],
-                             self.config["keep_dihedral_angle_cos_fragm4"],
-                             
+        Calculates the cosine-based potential energy.
         """
         potential_const = float(self.config["keep_dihedral_angle_cos_potential_const"])
         angle_const = float(self.config["keep_dihedral_angle_cos_angle_const"])
         
-        
         fragm_1_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_dihedral_angle_cos_fragm1"]) - 1], dim=0)
         fragm_2_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_dihedral_angle_cos_fragm2"]) - 1], dim=0)
         fragm_3_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_dihedral_angle_cos_fragm3"]) - 1], dim=0)
@@ -99,7 +298,19 @@ class StructKeepDihedralAnglePotentialCos:
         a2 = fragm_3_center - fragm_2_center
         a3 = fragm_4_center - fragm_3_center
 
+        # Explicitly compute cross products to determine switching factors
+        n1 = torch.linalg.cross(a1, a2)
+        n2 = torch.linalg.cross(a2, a3)
+        n1_sq_norm = torch.sum(n1**2, dim=-1)
+        n2_sq_norm = torch.sum(n2**2, dim=-1)
+
+        switch_1 = self._compute_switching(n1_sq_norm)
+        switch_2 = self._compute_switching(n2_sq_norm)
+
         angle = torch_calc_dihedral_angle_from_vec(a1, a2, a3)
-        energy = 0.5 * potential_const * (1.0 -1* torch.cos(angle_const * angle - (torch.deg2rad(torch.tensor(self.config["keep_dihedral_angle_cos_angle"])))))
+        raw_energy = 0.5 * potential_const * (1.0 - 1 * torch.cos(angle_const * angle - (torch.deg2rad(torch.tensor(self.config["keep_dihedral_angle_cos_angle"])))))
+        
+        # Apply smoothing
+        energy = raw_energy * switch_1 * switch_2
               
-        return energy #hartree
+        return energy

multioptpy/Potential/keep_outofplain_angle_potential.py

@@ -1,70 +1,277 @@
-
 from multioptpy.Parameters.parameter import UnitValueLib
 from multioptpy.Utils.calc_tools import torch_calc_outofplain_angle_from_vec
 import torch
-
+import math
 
 class StructKeepOutofPlainAnglePotential:
+    """
+    Class for calculating Out-of-Plane (Wilson) angle potential with robust singularity handling.
+    
+    Singularity Handling:
+    The Out-of-Plane angle measures the elevation of vector a1 from the plane defined by a2 and a3.
+    A singularity occurs when vectors a2 and a3 are collinear (angle 0 or 180).
+    In this case, the reference plane is undefined, and the normal vector vanishes.
+    
+    This implementation applies a 'Collinearity Guard' to force gradients to zero 
+    when the reference plane is undefined.
+    """
+
     def __init__(self, **kwarg):
         self.config = kwarg
         UVL = UnitValueLib()
         self.hartree2kcalmol = UVL.hartree2kcalmol 
         self.bohr2angstroms = UVL.bohr2angstroms 
         self.hartree2kjmol = UVL.hartree2kjmol 
+        
+        # Threshold for plane definition stability (Squared Norm of cross product)
+        # If the cross product of base vectors is smaller than this, 
+        # the plane is considered undefined.
+        self.COLLINEAR_CUT_SQ = 1e-8
+        
         return
+
     def calc_energy(self, geom_num_list, bias_pot_params=[]):
         """
-        # required variables: self.config["keep_out_of_plain_angle_spring_const"],
-                              self.config["keep_out_of_plain_angle_atom_pairs"]
-                              self.config["keep_out_of_plain_angle_angle"]
-        bias_pot_params[0] : keep_out_of_plain_angle_spring_const
-        bias_pot_params[1] : keep_out_of_plain_angle_angle 
-                        
+        Calculates Out-of-Plane angle energy.
+        
+        Args:
+            geom_num_list: Tensor of atomic coordinates (N_atoms, 3)
+            bias_pot_params: Optional bias parameters [k, theta_0]
+            
+        Definition:
+            Center atom: i (index 0 in pair list)
+            Neighbors: j, k, l (indices 1, 2, 3)
+            Vectors: a1 = r_j - r_i
+                     a2 = r_k - r_i
+                     a3 = r_l - r_i
+            
+            Angle is the deviation of a1 from the plane spanned by a2 and a3.
         """
-        a1 = geom_num_list[self.config["keep_out_of_plain_angle_atom_pairs"][1]-1] - geom_num_list[self.config["keep_out_of_plain_angle_atom_pairs"][0]-1]
-        a2 = geom_num_list[self.config["keep_out_of_plain_angle_atom_pairs"][2]-1] - geom_num_list[self.config["keep_out_of_plain_angle_atom_pairs"][0]-1]
-        a3 = geom_num_list[self.config["keep_out_of_plain_angle_atom_pairs"][3]-1] - geom_num_list[self.config["keep_out_of_plain_angle_atom_pairs"][0]-1]
-
-        angle = torch_calc_outofplain_angle_from_vec(a1, a2, a3)
+        
+        # ========================================
+        # 1. Parameter Retrieval
+        # ========================================
         if len(bias_pot_params) == 0:
-            energy = 0.5 * self.config["keep_out_of_plain_angle_spring_const"] * (angle - torch.deg2rad(torch.tensor(self.config["keep_out_of_plain_angle_angle"]))) ** 2
+            k = self.config["keep_out_of_plain_angle_spring_const"]
+            theta_0_deg = torch.tensor(self.config["keep_out_of_plain_angle_angle"])
+            theta_0 = torch.deg2rad(theta_0_deg)
         else:
-            energy = 0.5 * bias_pot_params[0] * (angle - torch.deg2rad(bias_pot_params[1])) ** 2
-        return energy #hartree 
+            k = bias_pot_params[0]
+            theta_0_deg = bias_pot_params[1]
+            if isinstance(theta_0_deg, torch.Tensor):
+                theta_0 = torch.deg2rad(theta_0_deg)
+            else:
+                theta_0 = torch.deg2rad(torch.tensor(theta_0_deg))
+
+        device = geom_num_list.device if isinstance(geom_num_list, torch.Tensor) else torch.device("cpu")
+        dtype = geom_num_list.dtype
+        theta_0 = theta_0.to(device=device, dtype=dtype)
+
+        # ========================================
+        # 2. Vector Calculations
+        # ========================================
+        # Indices are 1-based in config
+        # Atom 0 is the Central Atom based on the vector definition in your snippet
+        c_idx = self.config["keep_out_of_plain_angle_atom_pairs"][0] - 1
+        i1_idx = self.config["keep_out_of_plain_angle_atom_pairs"][1] - 1
+        i2_idx = self.config["keep_out_of_plain_angle_atom_pairs"][2] - 1
+        i3_idx = self.config["keep_out_of_plain_angle_atom_pairs"][3] - 1
+        
+        center_pos = geom_num_list[c_idx]
+        
+        # Vectors from Center to Neighbors
+        a1 = geom_num_list[i1_idx] - center_pos # The vector to measure (Probe)
+        a2 = geom_num_list[i2_idx] - center_pos # Plane definition vector 1
+        a3 = geom_num_list[i3_idx] - center_pos # Plane definition vector 2
+        
+        # ========================================
+        # 3. Plane Definition & Singularity Guard
+        # ========================================
+        # The plane is defined by a2 and a3.
+        # Normal vector n = a2 x a3
+        n = torch.linalg.cross(a2, a3)
+        
+        # Check squared norm of normal vector
+        # If n is near zero, a2 and a3 are collinear -> Plane Undefined
+        n_sq_norm = torch.sum(n**2, dim=-1)
+        
+        # Guard Mask
+        is_undefined_plane = (n_sq_norm < self.COLLINEAR_CUT_SQ)
+        
+        # Safe normalization
+        n_norm = torch.sqrt(n_sq_norm)
+        n_hat_demon = torch.clamp(n_norm.unsqueeze(-1), min=1e-12)
+        n_hat = n / n_hat_demon
+        
+        # ========================================
+        # 4. Angle Calculation (Robust atan2)
+        # ========================================
+        # We want the angle 'phi' between a1 and the plane (a2, a3).
+        # This is equivalent to 90 - angle(a1, n), or simply:
+        # sin(phi) = (a1 . n_hat) / |a1|
+        
+        # Height of a1 relative to the plane (Projection onto normal)
+        # h = |a1| * sin(phi)
+        h = torch.sum(a1 * n_hat, dim=-1)
+        
+        # Length of a1
+        a1_norm = torch.linalg.norm(a1) # add epsilon if atom overlap is a concern
+        
+        # Projected length of a1 onto the plane
+        # r_proj = |a1| * cos(phi) = sqrt(|a1|^2 - h^2)
+        # We clamp inside sqrt to avoid negative values due to numerical noise
+        r_proj_sq = a1_norm**2 - h**2
+        r_proj = torch.sqrt(torch.clamp(r_proj_sq, min=0.0))
+        
+        # Calculate angle using atan2(y, x) = atan2(height, projected_distance)
+        # This is valid for -90 to +90 degrees and stable at 90.
+        # Note: If a1 is perpendicular to plane, r_proj is 0, atan2(h, 0) gives +/- 90 correctly.
+        angle = torch.atan2(h, r_proj)
+        
+        # ========================================
+        # 5. Energy & Clamping
+        # ========================================
+        
+        # Difference from equilibrium
+        # Out-of-plane angles are usually non-periodic (limited to -90 to 90), 
+        # but if using generalized definition, simple subtraction is usually sufficient.
+        diff = angle - theta_0
+        
+        energy_harmonic = 0.5 * k * diff**2
+        
+        # Apply Guard:
+        # If plane is undefined (collinear base vectors), force energy/force to 0.
+        energy = torch.where(is_undefined_plane, torch.tensor(0.0, device=device, dtype=dtype), energy_harmonic)
+        
+        return energy
        
 class StructKeepOutofPlainAnglePotentialv2:
+    """
+    Class for calculating Out-of-Plane (Wilson) angle potential for fragment centroids
+    with robust singularity handling.
+    
+    This class calculates the angle of vector a1 (Frag1->Frag2) out of the plane 
+    defined by vectors a2 (Frag1->Frag3) and a3 (Frag1->Frag4).
+    
+    Singularity Handling:
+    - Plane Undefined: If Frag1, Frag3, and Frag4 are collinear, the reference plane 
+      cannot be defined (normal vector vanishes). A 'Collinearity Guard' forces 
+      gradients to zero in this region.
+    - Vertical Instability: Uses atan2(h, r_proj) instead of asin(h/r) to maintain 
+      numerical stability when the angle approaches +/- 90 degrees.
+    """
+
     def __init__(self, **kwarg):
         self.config = kwarg
         UVL = UnitValueLib()
         self.hartree2kcalmol = UVL.hartree2kcalmol 
         self.bohr2angstroms = UVL.bohr2angstroms 
         self.hartree2kjmol = UVL.hartree2kjmol 
+        
+        # Threshold for plane definition stability (Squared Norm of cross product)
+        # If the cross product of plane-defining vectors is smaller than this, 
+        # the plane is considered undefined.
+        self.COLLINEAR_CUT_SQ = 1e-8
+        
         return
     
     def calc_energy(self, geom_num_list, bias_pot_params=[]):
         """
-        # required variables: self.config["keep_out_of_plain_angle_v2_spring_const"], 
-                             self.config["keep_out_of_plain_angle_v2_angle"], 
-                             self.config["keep_out_of_plain_angle_v2_fragm1"],
-                             self.config["keep_out_of_plain_angle_v2_fragm2"],
-                             self.config["keep_out_of_plain_angle_v2_fragm3"],
-                             self.config["keep_out_of_plain_angle_v2_fragm4"],
-                             
+        Calculates Out-of-Plane angle energy for fragments.
+        
+        Args:
+            geom_num_list: Tensor of atomic coordinates (N_atoms, 3)
+            bias_pot_params: Optional bias parameters [k, theta_0]
         """
-        fragm_1_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_out_of_plain_angle_v2_fragm1"]) - 1], dim=0)
-        fragm_2_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_out_of_plain_angle_v2_fragm2"]) - 1], dim=0)
-        fragm_3_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_out_of_plain_angle_v2_fragm3"]) - 1], dim=0)
-        fragm_4_center = torch.mean(geom_num_list[torch.tensor(self.config["keep_out_of_plain_angle_v2_fragm4"]) - 1], dim=0)
+        
+        # ========================================
+        # 1. Parameter Retrieval
+        # ========================================
+        if len(bias_pot_params) == 0:
+            k = self.config["keep_out_of_plain_angle_v2_spring_const"]
+            theta_0_deg = torch.tensor(self.config["keep_out_of_plain_angle_v2_angle"])
+            theta_0 = torch.deg2rad(theta_0_deg)
+        else:
+            k = bias_pot_params[0]
+            theta_0_deg = bias_pot_params[1]
+            if isinstance(theta_0_deg, torch.Tensor):
+                theta_0 = torch.deg2rad(theta_0_deg)
+            else:
+                theta_0 = torch.deg2rad(torch.tensor(theta_0_deg))
+
+        # Device/Dtype handling
+        device = geom_num_list.device if isinstance(geom_num_list, torch.Tensor) else torch.device("cpu")
+        dtype = geom_num_list.dtype
+        theta_0 = theta_0.to(device=device, dtype=dtype)
+
+        # ========================================
+        # 2. Vector Calculations (Fragment Centroids)
+        # ========================================
+        
+        # Helper to get indices
+        def get_indices(key):
+            return torch.tensor(self.config[key], device=device, dtype=torch.long) - 1
 
-              
+        # Calculate centroids (Frag 1 is the Vertex/Center)
+        fragm_1_center = torch.mean(geom_num_list[get_indices("keep_out_of_plain_angle_v2_fragm1")], dim=0)
+        fragm_2_center = torch.mean(geom_num_list[get_indices("keep_out_of_plain_angle_v2_fragm2")], dim=0)
+        fragm_3_center = torch.mean(geom_num_list[get_indices("keep_out_of_plain_angle_v2_fragm3")], dim=0)
+        fragm_4_center = torch.mean(geom_num_list[get_indices("keep_out_of_plain_angle_v2_fragm4")], dim=0)
+
+        # Define vectors originating from Fragment 1
+        # a1: The "Probe" vector (whose angle we are measuring)
+        # a2, a3: The "Base" vectors (defining the reference plane)
         a1 = fragm_2_center - fragm_1_center
         a2 = fragm_3_center - fragm_1_center
         a3 = fragm_4_center - fragm_1_center
         
-        angle = torch_calc_outofplain_angle_from_vec(a1, a2, a3)
-        if len(bias_pot_params) == 0:
-            energy = 0.5 * self.config["keep_out_of_plain_angle_v2_spring_const"] * (angle - torch.deg2rad(torch.tensor(self.config["keep_out_of_plain_angle_v2_angle"]))) ** 2
-        else:
-            energy = 0.5 * bias_pot_params[0] * (angle - torch.deg2rad(bias_pot_params[1])) ** 2
-        return energy #hartree
-
+        # ========================================
+        # 3. Plane Definition & Singularity Guard
+        # ========================================
+        
+        # Normal vector to the plane spanned by a2 and a3
+        n = torch.linalg.cross(a2, a3)
+        
+        # Check if plane is defined (a2 and a3 are not collinear)
+        n_sq_norm = torch.sum(n**2, dim=-1)
+        is_undefined_plane = (n_sq_norm < self.COLLINEAR_CUT_SQ)
+        
+        # Safe normalization
+        n_norm = torch.sqrt(n_sq_norm)
+        n_hat_demon = torch.clamp(n_norm.unsqueeze(-1), min=1e-12)
+        n_hat = n / n_hat_demon
+        
+        # ========================================
+        # 4. Angle Calculation (Robust atan2)
+        # ========================================
+        
+        # Height of a1 relative to the plane (Projection onto normal)
+        # h = |a1| * sin(phi)
+        h = torch.sum(a1 * n_hat, dim=-1)
+        
+        # Length of a1
+        a1_norm = torch.linalg.norm(a1)
+        
+        # Projected length of a1 onto the plane
+        # r_proj = sqrt(|a1|^2 - h^2)
+        # Clamp to avoid negative sqrt due to numerical noise
+        r_proj_sq = a1_norm**2 - h**2
+        r_proj = torch.sqrt(torch.clamp(r_proj_sq, min=0.0))
+        
+        # Calculate angle using atan2(height, projected_distance)
+        # Valid range: -90 to +90 degrees (or -pi/2 to pi/2)
+        angle = torch.atan2(h, r_proj)
+        
+        # ========================================
+        # 5. Energy & Clamping
+        # ========================================
+        
+        diff = angle - theta_0
+        
+        energy_harmonic = 0.5 * k * diff**2
+        
+        # Apply Guard:
+        # If the reference plane is undefined, force energy to 0.0 to prevent gradient explosion.
+        energy = torch.where(is_undefined_plane, torch.tensor(0.0, device=device, dtype=dtype), energy_harmonic)
+        
+        return energy