Added add_atom_to_zmat_with_arbitrary_parameters to ZMat

arc/species/zmat.py

@@ -30,6 +30,7 @@
 import numpy as np
 import operator
 import re
+from scipy.optimize import minimize
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
 
 from rmgpy.molecule.molecule import Molecule
@@ -2092,3 +2093,185 @@ def map_index_to_int(index: Union[int, str]) -> int:
     if isinstance(index, str) and all(char.isdigit() for char in index[1:]):
         return int(index[1:])
     raise TypeError(f'Expected either an int or a string on the format "X15", got {index}')
+
+
+def add_atom_to_zmat_with_arbitrary_parameters(zmat: dict,
+                                               element: str,
+                                               r_index: int,
+                                               a_indices: tuple,
+                                               d_indices: tuple,
+                                               r_value: float,
+                                               a_value: float,
+                                               d_value: float,
+                                               opt_method: str = 'Nelder-Mead',
+                                               ) -> dict:
+    """
+    Add an atom to a zmat. The atom may have random r, a, and d parameters
+    (not necessarily defined for the same respective atoms).
+    This function will compute the correct parameter assignments for the new atom.
+    This function assumes that the zmat is fully defined up to the 4th atom.
+
+    Args:
+        zmat (dict): The zmat to process.
+        element (str): The element of the atom to add.
+        r_index (int): The index of an atom R to define the distance parameter R-X w.r.t the newly added atom, X.
+        a_indices (tuple): The indices of two atoms, A and B, to define the angle A-B-X parameter w.r.t the newly added atom, X.
+        d_indices (tuple): The indices of three atoms, L M and N, to define the dihedral angle L-M-N-X parameter w.r.t the newly added atom, X.
+        r_value (float): The value of the R-X distance parameter, r.
+        a_value (float): The value of the A-B-X angle parameter, a.
+        d_value (float): The value of the L-M-N-X dihedral angle parameter, d.
+        opt_method (str, optional): The optimization method to use for finding the new atom's coordinates.
+
+    Returns:
+        dict: The updated zmat.
+    """
+    if len(zmat['symbols']) < 3:
+        raise ZMatError('The zmat must have at least 3 atoms to add a new atom with arbitrary parameters.')
+    if len(a_indices) != 2 or len(d_indices) != 3:
+        raise ZMatError('The indices of the atoms defining the angle and dihedral must be of length 2 and 3, '
+                        f'respectively, got {a_indices} and {d_indices}.')
+
+    coords, symbols = zmat_to_coords(zmat, keep_dummy=False)
+
+    atom_r_coord = coords[r_index]
+    atom_a_coord = coords[a_indices[0]]
+    atom_b_coord = coords[a_indices[1]]
+    atom_l_coord = coords[d_indices[0]]
+    atom_m_coord = coords[d_indices[1]]
+    atom_n_coord = coords[d_indices[2]]
+
+    sphere_eq = sphere_constraint(reference_coord=atom_r_coord, distance=r_value)
+    angle_eq = angle_constraint(atom_a=atom_a_coord, atom_b=atom_b_coord, angle=a_value)
+    dihedral_eq = dihedral_constraint(atom_a=atom_l_coord, atom_b=atom_m_coord, atom_c=atom_n_coord, dihedral=d_value)
+
+    def objective_func(coord: Tuple[float, float, float]) -> float:
+        """
+        The objective function to minimize to satisfy the sphere, angle, and dihedral constraints.
+
+        Args:
+            coord (Tuple[float, float, float]): The Cartesian coordinates of the new atom.
+
+        Returns:
+            float: The sum of the squared differences between the constraints and their desired values.
+        """
+        x, y, z = coord
+        sphere_constraint_ = sphere_eq(x, y, z)
+        angle_constraint_ = angle_eq(x, y, z)
+        dihedral_constraint_ = dihedral_eq(x, y, z)
+        return sphere_constraint_**2 + angle_constraint_**2 + dihedral_constraint_**2
+
+    initial_guess = np.array([atom_r_coord[0] + r_value * 1.0,
+                              atom_r_coord[1] + r_value * 0.0,
+                              atom_r_coord[2] + r_value * 0.0])
+
+    result = minimize(objective_func, initial_guess, method=opt_method)
+    if not result.success:
+        raise ValueError(f"Failed to find a solution: {result.message}")
+
+    coords.append(tuple(result.x))
+    symbols.append(element)
+    zmat = xyz_to_zmat(xyz={'symbols': symbols, 'coords': coords})
+    return zmat
+
+
+def sphere_constraint(reference_coord: tuple, distance: float):
+    """
+    Generate the sphere equation for a new atom at a specific distance from a reference atom.
+
+    Args:
+        reference_coord (tuple): The Cartesian coordinates (x1, y1, z1) of the reference atom.
+        distance (float): The fixed distance (radius of the sphere) in Angstroms.
+
+    Returns:
+        function: A lambda function representing the sphere equation.
+    """
+    x1, y1, z1 = reference_coord
+    r_squared = distance ** 2
+    sphere_eq = lambda x, y, z: (x - x1) ** 2 + (y - y1) ** 2 + (z - z1) ** 2 - r_squared
+    return sphere_eq
+
+
+def angle_constraint(atom_a: tuple, atom_b: tuple, angle: float):
+    """
+    Generate the angle constraint for a new atom with two other atoms in Cartesian space.
+
+    Args:
+        atom_a (tuple): Cartesian coordinates of the first reference atom (A).
+        atom_b (tuple): Cartesian coordinates of the second reference atom (B).
+        angle (float): Desired angle (in degrees).
+
+    Returns:
+        function: A lambda function representing the angle constraint equation.
+    """
+    x1, y1, z1 = atom_a
+    x2, y2, z2 = atom_b
+    cos_angle = math.cos(math.radians(angle))
+
+    def angle_eq(x, y, z):
+        """
+        Calculate the angle between the vectors AB and BC, and compare it to the desired angle.
+
+        Args:
+            x (float): x-coordinate of the new atom.
+            y (float): y-coordinate of the new atom.
+            z (float): z-coordinate of the new atom.
+
+        Returns:
+            float: The difference between the calculated cosine of the angle and the desired cosine of the angle.
+        """
+        AB = (x2 - x1, y2 - y1, z2 - z1)
+        BC = (x - x2, y - y2, z - z2)
+        dot_product = AB[0] * BC[0] + AB[1] * BC[1] + AB[2] * BC[2]  # Dot product of AB and BC
+        AB_length = math.sqrt(AB[0] ** 2 + AB[1] ** 2 + AB[2] ** 2)  # Magnitudes of AB and BC
+        BC_length = math.sqrt(BC[0] ** 2 + BC[1] ** 2 + BC[2] ** 2)
+        cos_calc = dot_product / (AB_length * BC_length)
+        return cos_calc - cos_angle
+
+    return angle_eq
+
+def dihedral_constraint(atom_a: tuple, atom_b: tuple, atom_c: tuple, dihedral: float):
+    """
+    Generate the dihedral angle constraint for a new atom with three other atoms in Cartesian space.
+
+    Args:
+        atom_a (tuple): Cartesian coordinates of the first atom (A).
+        atom_b (tuple): Cartesian coordinates of the second atom (B).
+        atom_c (tuple): Cartesian coordinates of the third atom (C).
+        dihedral (float): Desired dihedral angle (in degrees).
+
+    Returns:
+        function: A lambda function representing the dihedral constraint equation.
+    """
+    x1, y1, z1 = atom_a
+    x2, y2, z2 = atom_b
+    x3, y3, z3 = atom_c
+    cos_d = math.cos(math.radians(dihedral))
+    sin_d = math.sin(math.radians(dihedral))
+
+    def dihedral_eq(x, y, z):
+        """
+        Calculate the dihedral angle between the planes defined by vectors AB and BC, and compare it to the desired angle.
+        Return the combined equation: both cosine and sine should match.
+
+        Args:
+            x (float): x-coordinate of the new atom.
+            y (float): y-coordinate of the new atom.
+            z (float): z-coordinate of the new atom.
+
+        Returns:
+            float: The difference between the calculated cosine and sine of the dihedral angle
+                   and the desired cosine and sine of the dihedral angle.
+        """
+        AB = np.array([x2 - x1, y2 - y1, z2 - z1])
+        BC = np.array([x3 - x2, y3 - y2, z3 - z2])
+        CD = np.array([x - x3, y - y3, z - z3])
+        N1 = np.cross(AB, BC)  # Cross products
+        N2 = np.cross(BC, CD)
+        N1_norm = np.linalg.norm(N1)  # Magnitudes
+        N2_norm = np.linalg.norm(N2)
+        BC_norm = np.linalg.norm(BC)
+        cos_calc = np.dot(N1, N2) / (N1_norm * N2_norm)
+        sin_calc = np.dot(BC, np.cross(N1, N2)) / (BC_norm * N1_norm * N2_norm)
+        return (cos_calc - cos_d) ** 2 + (sin_calc - sin_d) ** 2
+
+    return dihedral_eq