Add optional weights and add quaternion component of rotation

Author: gustavor101

pysages/colvars/orientation.py

@@ -13,18 +13,18 @@
 from jax import numpy as np
 from jax.numpy import linalg
 
-from pysages.collective_variables.core import CollectiveVariable, AxisCV
-from pysages.collective_variables.coordinates import barycenter
+from pysages.colvars.core import CollectiveVariable, AxisCV
+from pysages.colvars.coordinates import barycenter, weighted_barycenter
 
 
-QUATERNION_BASES = {
+_QUATERNION_BASES = {
     0: (0, 1, 0, 0),
     1: (0, 0, 1, 0),
     2: (0, 0, 0, 1),
 }
 
 
-def quaternion_matrix(positions, references):
+def quaternion_matrix(positions, references, weights):
     """
     Function to construct the quaternion matrix based on the reference positions.
 
@@ -35,12 +35,24 @@ def quaternion_matrix(positions, references):
     references: np.array
        Cartesian coordinates of the reference positions of the atoms in indices.
        The number of coordinates must match the ones of the atoms used in indices.
+    weights: np.array
+       Weights for the barycenter calculation
     """
-    pos_b = barycenter(positions)
-    ref_b = barycenter(references)
+    if len(positions) != len(references):
+        raise RuntimeError("References must be of the same length as the positions")
+    pos_b = np.where(weights is None,
+                     barycenter(positions),
+                     weighted_barycenter(positions, weights))
+    ref_b = np.where(weights is None,
+                     barycenter(references),
+                     weighted_barycenter(references, weights))
     R = np.zeros((3, 3))
-    for pos, ref in zip(positions, references):
-        R += np.outer(pos - pos_b, ref - ref_b)
+    if weights is None:
+        for pos, ref in zip(positions, references):
+            R += np.outer(pos - pos_b, ref - ref_b)
+    else:
+        for pos, ref, w in zip(positions, references, weights):
+            R += np.outer(w*pos - pos_b, w*ref - ref_b)
     S_00 = R[0, 0] + R[1, 1] + R[2, 2]
     S_01 = R[1, 2] - R[2, 1]
     S_02 = R[2, 0] - R[0, 2]
@@ -77,17 +89,20 @@ class Tilt(AxisCV):
        of coordinates must match the ones of the atoms used in indices.
     """
 
-    def __init__(self, indices, axis, references):
+    def __init__(self, indices, axis, references, weights=None):
+        if weights is not None and len(indices) != len(weights):
+            raise RuntimeError("Indices and weights must be of the same length")
         super().__init__(indices, axis)
         self.references = np.asarray(references)
-        self.e = np.asarray(QUATERNION_BASES[axis])
+        self.weights = np.asarray(weights)
+        self.e = np.asarray(_QUATERNION_BASES[axis])
 
     @property
     def function(self):
-        return lambda rs: tilt(rs, self.references, self.axis)
+        return lambda rs: tilt(rs, self.references, self.axis, self.weights)
 
 
-def tilt(r1, r2, e):
+def tilt(r1, r2, e, w):
     """
     Function to calculate the tilt rotation respect to an axis.
 
@@ -100,8 +115,10 @@ def tilt(r1, r2, e):
        The number of coordinates must match the ones of the atoms used in indices.
     e: JaxArray
        Quaternion rotation axis.
+    w: JaxArray
+       Atomic weights.
     """
-    S = quaternion_matrix(r1, r2)
+    S = quaternion_matrix(r1, r2, w)
     _, v = linalg.eigh(S)
     v_dot_e = np.dot(v[:, 3], e)
     cs = np.cos(np.arctan2(v_dot_e, v[0, 3]))
@@ -124,17 +141,20 @@ class RotationEigenvalues(CollectiveVariable):
        The number of coordinates must match the ones of the atoms used in indices.
     """
 
-    def __init__(self, indices, axis, references):
+    def __init__(self, indices, axis, references, weights=None):
+        if weights is not None and len(indices) != len(weights):
+            raise RuntimeError("Indices and weights must be of the same length")
         super().__init__(indices)
         self.axis = axis
         self.references = np.asarray(references)
+        self.weights = np.asarray(weights)
 
     @property
     def function(self):
-        return lambda r: rotation_eigvals(r, self.references)[self.axis]
+        return lambda r: rotation_eigvals(r, self.references, self.weights)[self.axis]
 
 
-def rotation_eigvals(r1, r2):
+def rotation_eigvals(r1, r2, w):
     """
     Returns the eigenvalue of the quaternion matrix rspect to an axis.
 
@@ -147,7 +167,7 @@ def rotation_eigvals(r1, r2):
     r2: np.array
        Cartesian coordinates of the reference position of the atoms in r1.
     """
-    S = quaternion_matrix(r1, r2)
+    S = quaternion_matrix(r1, r2, w)
     return linalg.eigvalsh(S)
 
 
@@ -166,17 +186,20 @@ class SpinAngle(AxisCV):
        The coordinates must match the ones of the atoms used in indices.
     """
 
-    def __init__(self, indices, axis, references):
+    def __init__(self, indices, axis, references, weights=None):
+        if weights is not None and len(indices) != len(weights):
+            raise RuntimeError("Indices and weights must be of the same length")
         super().__init__(indices, axis)
         self.references = np.asarray(references)
-        self.e = np.array(QUATERNION_BASES[axis])
+        self.e = np.array(_QUATERNION_BASES[axis])
+        self.weights = np.asarray(weights)
 
     @property
     def function(self):
-        return lambda r: spin(r, self.references, self.e)
+        return lambda r: spin(r, self.references, self.e, self.weights)
 
 
-def spin(r1, r2, e):
+def spin(r1, r2, e, w):
     """
     Calculate the spin angle rotation respect to an axis.
 
@@ -189,7 +212,7 @@ def spin(r1, r2, e):
     e: np.array
         Rotation axis.
     """
-    S = quaternion_matrix(r1, r2)
+    S = quaternion_matrix(r1, r2, w)
     _, v = linalg.eigh(S)
     v_dot_e = np.dot(v[:, 3], e)
     return 2 * np.arctan2(v_dot_e, v[0, 3])
@@ -209,16 +232,19 @@ class RotationAngle(CollectiveVariable):
        The coordinates must match the ones of the atoms used in indices.
     """
 
-    def __init__(self, indices, references):
+    def __init__(self, indices, references, weights=None):
+        if weights is not None and len(indices) != len(weights):
+            raise RuntimeError("Indices and weights must be of the same length")
         super().__init__(indices)
         self.references = np.asarray(references)
+        self.weights = np.asarray(weights)
 
     @property
     def function(self):
-        return lambda r: rotation_angle(r, self.references)
+        return lambda r: rotation_angle(r, self.references, self.weights)
 
 
-def rotation_angle(r1, r2):
+def rotation_angle(r1, r2, w):
     """
     Calculate the rotation angle respect to a reference.
 
@@ -229,7 +255,7 @@ def rotation_angle(r1, r2):
     r2: np.array
        Cartesian coordinates of the reference position of the atoms.
     """
-    S = quaternion_matrix(r1, r2)
+    S = quaternion_matrix(r1, r2, w)
     _, v = linalg.eigh(S)
     return 2 * np.arccos(v[0, 3])
 
@@ -248,16 +274,19 @@ class RotationProjection(CollectiveVariable):
        The coordinates must match the ones of the atoms used in indices.
     """
 
-    def __init__(self, indices, references):
+    def __init__(self, indices, references, weights=None):
+        if weights is not None and len(indices) != len(weights):
+            raise RuntimeError("Indices and weights must be of the same length")
         super().__init__(indices)
         self.references = np.asarray(references)
+        self.weights = np.asarray(weights)
 
     @property
     def function(self):
-        return lambda r: rotation_projection(r, self.references)
+        return lambda r: rotation_projection(r, self.references, self.weights)
 
 
-def rotation_projection(r1, r2):
+def rotation_projection(r1, r2, w):
     """
     Calculate the rotation angle projection.
 
@@ -268,7 +297,7 @@ def rotation_projection(r1, r2):
     r2: np.array
        Cartesian coordinates of the reference position of the atoms.
     """
-    S = quaternion_matrix(r1, r2)
+    S = quaternion_matrix(r1, r2, w)
     _, v = linalg.eigh(S)
     return 2 * v[0, 3] * v[0, 3] - 1
 
@@ -287,16 +316,19 @@ class RMSD(CollectiveVariable):
        The coordinates must match the ones of the atoms used in indices.
     """
 
-    def __init__(self, indices, references):
+    def __init__(self, indices, references, weights=None):
+        if weights is not None and len(indices) != len(weights):
+            raise RuntimeError("Indices and weights must be of the same length")
         super().__init__(indices)
         self.references = np.asarray(references)
+        self.weights = np.asarray(weights)
 
     @property
     def function(self):
-        return lambda r: rmsd(r, self.references)
+        return lambda r: rmsd(r, self.references, self.weights)
 
 
-def sq_norm_rotation(positions, references):
+def sq_norm_rotation(positions, references, weights):
     """
     Calculate the squared norm of the atomic positions and references respect to barycenters.
 
@@ -307,17 +339,31 @@ def sq_norm_rotation(positions, references):
     references: np.array
        Cartesian coordinates of the reference position of the atoms.
     """
-    pos_b = barycenter(positions)
-    ref_b = barycenter(references)
+    if len(positions) != len(references):
+        raise RuntimeError("References must be of the same length as the positions")
+    pos_b = np.where(weights is None,
+                     barycenter(positions),
+                     weighted_barycenter(positions, weights))
+    ref_b = np.where(weights is None,
+                     barycenter(references),
+                     weighted_barycenter(references, weights))
     R = 0.0
-    for pos, ref in zip(positions, references):
-        pos -= pos_b
-        ref -= ref_b
-        R += np.dot(pos, pos) + np.dot(ref, ref)
+    if weights is None:
+        for pos, ref in zip(positions, references):
+            pos -= pos_b
+            ref -= ref_b
+            R += np.dot(pos, pos) + np.dot(ref, ref)
+    else:
+        for pos, ref, w in zip(positions, references, weights):
+            pos *= w
+            ref *= w
+            pos -= pos_b
+            ref -= ref_b
+            R += np.dot(pos, pos) + np.dot(ref, ref)
     return R
 
 
-def rmsd(r1, r2):
+def rmsd(r1, r2, w_0):
     """
     Calculate the rmsd respect to a reference using quaternions.
 
@@ -329,7 +375,53 @@ def rmsd(r1, r2):
        Cartesian coordinates of the reference position of the atoms.
     """
     N = r1.shape[0]
-    S = quaternion_matrix(r1, r2)
+    S = quaternion_matrix(r1, r2, w_0)
     w = linalg.eigvalsh(S)
-    norm_sq = sq_norm_rotation(r1, r2)
+    norm_sq = sq_norm_rotation(r1, r2, w_0)
     return np.sqrt((norm_sq - 2 * np.max(w)) / N)
+
+
+class QuaternionComponent(CollectiveVariable):
+    """
+    Calculate the quaternion component of the rotation respect to a reference.
+
+    Parameters
+    ----------
+    indices: list[int], list[tuple(int)]
+       Select atom groups via indices.
+    axis: int
+       Index for the component of the quaternion (a value form `0` to `3`).
+    references: list[tuple(float)]
+       Cartesian coordinates of the reference position of the atoms in indices.
+       The number of coordinates must match the ones of the atoms used in indices.
+    """
+
+    def __init__(self, indices, axis, references, weights=None):
+        if weights is not None and len(indices) != len(weights):
+            raise RuntimeError("Indices and weights must be of the same length")
+        super().__init__(indices)
+        self.axis = axis
+        self.references = np.asarray(references)
+        self.weights = np.asarray(weights)
+
+    @property
+    def function(self):
+        return lambda r: quaternion_component(r, self.references, self.weights)[self.axis]
+
+
+def quaternion_component(r1, r2, w):
+    """
+    Returns the eigenvalue of the quaternion matrix rspect to an axis.
+
+    Parameters
+    ----------
+    r1: np.array
+       Atomic positions.
+    axis: int
+       select the component of the quaternion for rotation.
+    r2: np.array
+       Cartesian coordinates of the reference position of the atoms in r1.
+    """
+    S = quaternion_matrix(r1, r2, w)
+    _, v = linalg.eigh(S)
+    return v

pysages/colvars/shape.py

@@ -25,7 +25,7 @@ class RadiusOfGyration(CollectiveVariable):
     squared: Optional[bool]
         Indicates whether to return the squared value.
     weights: Optional[JaxArray]
-        If providede the weighted radius_of_gyration will be computed.
+        If providede the weighted radius of gyration will be computed.
     """
 
     def __init__(self, indices, squared=False, weights=None):
@@ -48,7 +48,7 @@ def function(self):
         else:
             rog = radius_of_gyration
 
-        return rog if self.squared else jit(lambda rs: np.sqrt(rog(rs)))
+        return jit(rog) if self.squared else jit(lambda rs: np.sqrt(rog(rs)))
 
 
 def radius_of_gyration(positions):