From a8b0eedb50b132d0bc0ace6eb44a2bc74c874f5f Mon Sep 17 00:00:00 2001 From: gustavor101 Date: Tue, 17 May 2022 14:41:59 -0500 Subject: [PATCH 1/3] Fix radius of gyration computation The reference system for calculating the radius of gyration is respect to the barycenter or weighted barycenter in case of weighted radius of gyration. Also its value is the square root of previous values. --- pysages/colvars/shape.py | 44 ++++++++++++++++++++++++++++++---------- 1 file changed, 33 insertions(+), 11 deletions(-) diff --git a/pysages/colvars/shape.py b/pysages/colvars/shape.py index 817c36ca..ae453ff6 100644 --- a/pysages/colvars/shape.py +++ b/pysages/colvars/shape.py @@ -6,25 +6,35 @@ Collective Variables that are calcuated from the shape of group of atoms. """ -import jax.numpy as np +from jax import jit, numpy as np from jax.numpy import linalg from pysages.colvars.core import CollectiveVariable, AxisCV +from pysages.colvars.coordinates import barycenter, weighted_barycenter class RadiusOfGyration(CollectiveVariable): """ - Collective Variable that calculates the unweighted radius of gyration as CV. + Collective Variable that calculates the unweighted radius of gyration. Parameters ---------- indices: list[int], list[tuple(int)] Must be a list or tuple of atoms (ints or ranges) or groups of atoms. A group is specified as a nested list or tuple of atoms. - group_length: int, optional - Specify if a fixed group length is expected. + squared: Optional[bool] + Indicates whether to return the squared value. + weights: Optional[JaxArray] + If providede the weighted radius_of_gyration will be computed. """ + def __init__(self, indices, squared=False, 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.squared = squared + self.weights = np.asarray(weights) + @property def function(self): """ @@ -33,12 +43,17 @@ def function(self): Callable See `pysages.colvars.shape.radius_of_gyration` for details. """ - return radius_of_gyration + if self.weights: + rog = jit(lambda rs: weighted_radius_of_gyration(rs, self.weights)) + else: + rog = radius_of_gyration + + return rog if self.squared else jit(lambda rs: np.sqrt(rog(rs))) def radius_of_gyration(positions): """ - Calculate the radius of gyration for a group of atoms. + Calculate the radius of gyration (squared) for a group of atoms. Parameters ---------- @@ -48,19 +63,22 @@ def radius_of_gyration(positions): Returns ------- DeviceArray - Radius of gyration vector + Radius of gyration (scalar) """ group_length = positions.shape[0] rog = np.zeros((3,)) + r_c = barycenter(positions) # TODO: Replace by `np.sum` and `vmap` # pylint:disable=fixme for r in positions: + r -= r_c rog += np.dot(r, r) return rog / group_length def weighted_radius_of_gyration(positions, weights): """ - Calculate the radius of gyration for a group of atoms weighted by arbitrary weights. + Calculate the radius of gyration (squared) for a group of atoms + weighted by user provided weights. Parameters ---------- @@ -72,13 +90,14 @@ def weighted_radius_of_gyration(positions, weights): Returns ------- DeviceArray - Weighted radius of gyration vector + Weighted radius of gyration (scalar) """ group_length = positions.shape[0] rog = np.zeros((3,)) + r_c = weighted_barycenter(positions, weights) # TODO: Replace by `np.sum` and `vmap` # pylint:disable=fixme for i in range(group_length): - w, r = weights[i], positions[i] + w, r = weights[i], positions[i] - r_c rog += w * np.dot(r, r) return rog @@ -128,7 +147,9 @@ def gyration_tensor(positions): """ group_length = positions.shape[0] gyr = np.zeros((3, 3)) + r_c = barycenter(positions) for r in positions: + r -= r_c gyr += np.outer(r, r) return gyr / group_length @@ -151,8 +172,9 @@ def weighted_gyration_tensor(positions, weights): """ group_length = positions.shape[0] gyr = np.zeros((3, 3)) + r_c = weighted_barycenter(positions, weights) for i in range(group_length): - w, r = weights[i], positions[i] + w, r = weights[i], positions[i] - r_c gyr += w * np.outer(r, r) return gyr From 3f3c2eed01c083381204ba3cdb7e639a32169302 Mon Sep 17 00:00:00 2001 From: gustavor101 Date: Tue, 17 May 2022 07:52:09 -0500 Subject: [PATCH 2/3] Implement orientations-based collective variables Collective variables that measure orientation respect a reference using quaternions --- pysages/colvars/__init__.py | 19 +- pysages/colvars/orientation.py | 335 +++++++++++++++++++++++++++++++++ 2 files changed, 349 insertions(+), 5 deletions(-) create mode 100644 pysages/colvars/orientation.py diff --git a/pysages/colvars/__init__.py b/pysages/colvars/__init__.py index 565ee669..dee8f98c 100644 --- a/pysages/colvars/__init__.py +++ b/pysages/colvars/__init__.py @@ -16,6 +16,20 @@ DihedralAngle, ) +from .coordinates import ( + Component, + Distance, +) + +from .orientation import ( + RMSD, + RotationAngle, + RotationEigenvalues, + RotationProjection, + SpinAngle, + Tilt, +) + from .shape import ( RadiusOfGyration, PrincipalMoment, @@ -24,11 +38,6 @@ ShapeAnisotropy, ) -from .coordinates import ( - Component, - Distance, -) - from .utils import ( get_periods, wrap, diff --git a/pysages/colvars/orientation.py b/pysages/colvars/orientation.py new file mode 100644 index 00000000..ddc363b7 --- /dev/null +++ b/pysages/colvars/orientation.py @@ -0,0 +1,335 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2020-2021: PySAGES contributors +# See LICENSE.md and CONTRIBUTORS.md at https://github.com/SSAGESLabs/PySAGES + +""" +Collective variables for orientiations describe the orientation measures +of particles in the simulation respect to a reference. + +It is common to describe such orientations using RMSD, tilt, sping angle +and more to spot inside a molecule or protein a conformation change. +""" + +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 + + +QUATERNION_BASES = { + 0: (0, 1, 0, 0), + 1: (0, 0, 1, 0), + 2: (0, 0, 0, 1), +} + + +def quaternion_matrix(positions, references): + """ + Function to construct the quaternion matrix based on the reference positions. + + Parameters + ---------- + positions: np.array + atomic positions via indices. + 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. + """ + pos_b = barycenter(positions) + ref_b = barycenter(references) + R = np.zeros((3, 3)) + for pos, ref in zip(positions, references): + R += np.outer(pos - pos_b, 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] + S_03 = R[0, 1] - R[1, 0] + S_11 = R[0, 0] - R[1, 1] - R[2, 2] + S_12 = R[0, 1] + R[1, 0] + S_13 = R[0, 2] + R[2, 0] + S_22 = -R[0, 0] + R[1, 1] - R[2, 2] + S_23 = R[1, 2] + R[2, 1] + S_33 = -R[0, 0] + R[1, 1] + R[2, 2] + S = np.array( + [ + [S_00, S_01, S_02, S_03], + [S_01, S_11, S_12, S_13], + [S_02, S_12, S_22, S_23], + [S_03, S_13, S_23, S_33], + ] + ) + return S + + +class Tilt(AxisCV): + """ + Use a reference to fit the tilt rotation respect to an axis. + + Parameters + ---------- + indices: list[int], list[tuple(int)] + Select atom groups via indices. + axis: int + Cartesian coordinate axis of rotation `0` (X), `1` (Y), `2` (Z) that is requested in CV. + 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): + super().__init__(indices, axis) + self.references = np.asarray(references) + self.e = np.asarray(QUATERNION_BASES[axis]) + + @property + def function(self): + return lambda rs: tilt(rs, self.references, self.axis) + + +def tilt(r1, r2, e): + """ + Function to calculate the tilt rotation respect to an axis. + + Parameters + ---------- + r1: JaxArray + Atomic positions. + r2: JaxArray + 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. + e: JaxArray + Quaternion rotation axis. + """ + S = quaternion_matrix(r1, r2) + _, v = linalg.eigh(S) + v_dot_e = np.dot(v[:, 3], e) + cs = np.cos(np.arctan2(v_dot_e, v[0, 3])) + z = np.where(cs == 0, 0, v[0, 3] / cs) + return 2 * z * z - 1 + + +class RotationEigenvalues(CollectiveVariable): + """ + Calculate the eigenvalues of the rotation matrix respect to a reference. + + Parameters + ---------- + indices: list[int], list[tuple(int)] + Select atom groups via indices. + axis: int + Index for the eigenvalue of the rotation matrix (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): + super().__init__(indices) + self.axis = axis + self.references = np.asarray(references) + + @property + def function(self): + return lambda r: rotation_eigvals(r, self.references)[self.axis] + + +def rotation_eigvals(r1, r2): + """ + Returns the eigenvalue of the quaternion matrix rspect to an axis. + + Parameters + ---------- + r1: np.array + Atomic positions. + axis: int + select the eigenvalue of the matrix. + r2: np.array + Cartesian coordinates of the reference position of the atoms in r1. + """ + S = quaternion_matrix(r1, r2) + return linalg.eigvalsh(S) + + +class SpinAngle(AxisCV): + """ + Use a reference to fit the spin angle rotation respect to an axis. + + Parameters + ---------- + indices: list[int], list[tuple(int)] + Select atom groups via indices. + axis: int + Cartesian coordinate axis of rotation `0` (X), `1` (Y), `2` (Z) that is requested in CV. + references: list[tuple(float)] + Cartesian coordinates of the reference position of the atoms in indices. + The coordinates must match the ones of the atoms used in indices. + """ + + def __init__(self, indices, axis, references): + super().__init__(indices, axis) + self.references = np.asarray(references) + self.e = np.array(QUATERNION_BASES[axis]) + + @property + def function(self): + return lambda r: spin(r, self.references, self.e) + + +def spin(r1, r2, e): + """ + Calculate the spin angle rotation respect to an axis. + + Parameters + ---------- + r1: np.array + Atomic positions. + r2: np.array + Cartesian coordinates of the reference position of the atoms. + e: np.array + Rotation axis. + """ + S = quaternion_matrix(r1, r2) + _, v = linalg.eigh(S) + v_dot_e = np.dot(v[:, 3], e) + return 2 * np.arctan2(v_dot_e, v[0, 3]) + + +class RotationAngle(CollectiveVariable): + """ + Use a reference to fit the rotation angle of the atoms. + The angle varies from -pi to pi. + + Parameters + ---------- + indices: list[int], list[tuple(int)] + Select atom groups via indices. + references: list[tuple(float)] + Cartesian coordinates of the reference position of the atoms in indices. + The coordinates must match the ones of the atoms used in indices. + """ + + def __init__(self, indices, references): + super().__init__(indices) + self.references = np.asarray(references) + + @property + def function(self): + return lambda r: rotation_angle(r, self.references) + + +def rotation_angle(r1, r2): + """ + Calculate the rotation angle respect to a reference. + + Parameters + ---------- + r1: np.array + Atomic positions. + r2: np.array + Cartesian coordinates of the reference position of the atoms. + """ + S = quaternion_matrix(r1, r2) + _, v = linalg.eigh(S) + return 2 * np.arccos(v[0, 3]) + + +class RotationProjection(CollectiveVariable): + """ + Use a reference to fit the cosine of the angle rotation of the atoms. + The projection varies from -1 to 1. + + Parameters + ---------- + indices: list[int], list[tuple(int)] + Select atom groups via indices. + references: list[tuple(float)] + Cartesian coordinates of the reference position of the atoms in indices. + The coordinates must match the ones of the atoms used in indices. + """ + + def __init__(self, indices, references): + super().__init__(indices) + self.references = np.asarray(references) + + @property + def function(self): + return lambda r: rotation_projection(r, self.references) + + +def rotation_projection(r1, r2): + """ + Calculate the rotation angle projection. + + Parameters + ---------- + r1: np.array + Atomic positions. + r2: np.array + Cartesian coordinates of the reference position of the atoms. + """ + S = quaternion_matrix(r1, r2) + _, v = linalg.eigh(S) + return 2 * v[0, 3] * v[0, 3] - 1 + + +class RMSD(CollectiveVariable): + """ + Use a reference to calculate the RMSD of a set of atoms. + The algorithm is based on https://doi.org/10.1002/jcc.20110. + + Parameters + ---------- + indices: list[int], list[tuple(int)] + Select atom groups via indices. + references: list[tuple(float)] + Cartesian coordinates of the reference position of the atoms in indices. + The coordinates must match the ones of the atoms used in indices. + """ + + def __init__(self, indices, references): + super().__init__(indices) + self.references = np.asarray(references) + + @property + def function(self): + return lambda r: rmsd(r, self.references) + + +def sq_norm_rotation(positions, references): + """ + Calculate the squared norm of the atomic positions and references respect to barycenters. + + Parameters + ---------- + positions: np.array + Atomic positions. + references: np.array + Cartesian coordinates of the reference position of the atoms. + """ + pos_b = barycenter(positions) + ref_b = barycenter(references) + 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) + return R + + +def rmsd(r1, r2): + """ + Calculate the rmsd respect to a reference using quaternions. + + Parameters + ---------- + r1: np.array + Atomic positions. + r2: np.array + Cartesian coordinates of the reference position of the atoms. + """ + N = r1.shape[0] + S = quaternion_matrix(r1, r2) + w = linalg.eigvalsh(S) + norm_sq = sq_norm_rotation(r1, r2) + return np.sqrt((norm_sq - 2 * np.max(w)) / N) From c0fe37ee2172502963b0a2c93047a86ad5d76247 Mon Sep 17 00:00:00 2001 From: gustavor101 Date: Sat, 28 May 2022 08:58:32 -0500 Subject: [PATCH 3/3] Add optional weights and add quaternion component of rotation --- pysages/colvars/orientation.py | 176 +++++++++++++++++++++++++-------- pysages/colvars/shape.py | 4 +- 2 files changed, 136 insertions(+), 44 deletions(-) diff --git a/pysages/colvars/orientation.py b/pysages/colvars/orientation.py index ddc363b7..ba93cab3 100644 --- a/pysages/colvars/orientation.py +++ b/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 diff --git a/pysages/colvars/shape.py b/pysages/colvars/shape.py index ae453ff6..4a9db888 100644 --- a/pysages/colvars/shape.py +++ b/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):