pyscivr.coordinates.core module

Coordinate transformations for internal coordinates and quaternion coordinates.

See sections II B 1, II C, III B, III C, and appendix L of the paper for details.

Notation:

N: number of atoms.
K: number of all coordinates (rows of the coordinate system).
m: number of active coordinates. (m <= K)

pyscivr.coordinates.core.get_bond(q)[source]

Compute the bond length between two atoms.

Parameters:: q (numpy.typing.ArrayLike) – (2, 3) positions of the two atoms.
Returns:: Distance between q[0] and q[1].
Return type:: Array

pyscivr.coordinates.core.get_bond_angle(q)[source]

Compute the bond angle at q[1].

Parameters:: q (numpy.typing.ArrayLike) – (3, 3) positions of the three atoms.
Returns:: Angle between q[0]-q[1] and q[2]-q[1] in [0, pi].
Return type:: Array

pyscivr.coordinates.core.get_dihedral_angle(q)[source]

Compute the dihedral q[0]-q[1]-q[2]-q[3].

Parameters:: q (numpy.typing.ArrayLike) – (4, 3) positions of the four atoms.
Returns:: Dihedral angle in [-pi, pi).
Return type:: Array

pyscivr.coordinates.core.get_polar_angle(q)[source]

Compute the polar angle of the vector q[1]-q[0] from the z-axis.

Parameters:: q (numpy.typing.ArrayLike) – (2, 3) positions of the two atoms.
Returns:: Polar angle in [0, pi].
Return type:: Array

pyscivr.coordinates.core.get_azimuthal_angle(q)[source]

Compute the azimuthal angle of the vector q[1]-q[0] from the x-axis in the xy-plane.

Parameters:: q (numpy.typing.ArrayLike) – (2, 3) positions of the two atoms.
Returns:: Azimuthal angle in [-pi, pi).
Return type:: Array

pyscivr.coordinates.core.quat_to_rotmat(q)[source]

Convert a quaternion to a rotation matrix (Eq. (L6) in the paper).

Parameters:: q (numpy.typing.ArrayLike) – (4,) quaternion (q0, q1, q2, q3).
Returns:: (3, 3) rotation matrix C.
Return type:: Array

pyscivr.coordinates.core.get_rotation_matrix(posi_c, posi_ref_c, mass)[source]

Compute the rotation matrix minimizing the mass-weighted RMSD by Horn’s method.

The rotation matrix between posi_c (\(q_{c}\)) and posi_ref_c (\(q_{\text{ref}, c}\)) is computed following Horn, J. Opt. Soc. Am. A 4, 629-642 (1987).

Solves the optimization problem:

\[\min_{C \in SO(3)} \mathrm{tr}\left[(q_{\text{ref}, c} C^T - q_{c})^T M (q_{\text{ref}, c} C^T - q_{c} )\right]\]

with \(M = \text{diag}(\text{mass})\). \(C\) rotates \(q_{\text{ref}, c}\) to \(q_{c}\) (\(q_{\text{ref}, c} C^T \approx q_{c}\)).

Parameters:

posi_c (numpy.typing.ArrayLike) – (N, 3) mass-centered coordinates.
posi_ref_c (numpy.typing.ArrayLike) – (N, 3) mass-centered reference coordinates.
mass (numpy.typing.ArrayLike) – (N,) atomic masses.

Returns:

(3, 3) rotation matrix \(C\).

Return type:

Array

Note

\(C\) is obtained from the eigenvector \(q\) of the largest eigenvalue \(\lambda\) of the symmetric (4, 4) matrix \(K\). The derivative of \(q\) is implemented by a custom JVP rule (_top_eigvec),

\[(\lambda I - K + q q^T) \dot{q} = \dot{K} q - (q^T \dot{K} q) q\]

where the term \(q q^T\) makes the matrix nonsingular as long as the largest eigenvalue is nondegenerate. The resulting \(\dot{q}\) is orthogonal to \(q\), so the normalization of \(q\) is preserved.

The default derivative of jnp.linalg.eigh involves \(1 / (\lambda_i - \lambda_j)\) for all pairs of eigenvalues, and yields NaN when any two eigenvalues are degenerate (e.g., near the reference geometry of a spherical top molecule such as CH4), even though only the top eigenvector is needed here. The custom rule uses only the top eigenpair and remains finite in such cases. Since the rule itself is differentiable, reverse-mode and higher-order derivatives (used in the Hessian contractions) are also available.

pyscivr.coordinates.core.get_euler_angles(C)[source]

Compute the Euler angles (yaw, pitch, roll) in ZYX convention from the rotation matrix (Eq. (L7) in the paper).

Parameters:

C (numpy.typing.ArrayLike) – (3, 3) rotation matrix (from reference to current)

Returns:

(3,) Euler angles (yaw, pitch, roll), where

yaw: rotation around z-axis in [-pi, pi).
pitch: rotation around y-axis, shifted to [0, pi].
roll: rotation around x-axis in [-pi, pi).

Return type:

Array

class pyscivr.coordinates.core.Coordinate(kind, atom_indices, value_ref, name, topology)[source]

Bases: object

Definition of a single coordinate.

Parameters:

kind (str)
atom_indices (tuple[int, ...])
value_ref (float)
name (str)
topology (str)

kind: str: Kind of the coordinate.

atom_indices: tuple[int, ...]: Indices of the atoms defining the coordinate (empty for global and constant coordinates).

value_ref: float: Value at the reference geometry.

name: str: Name of the coordinate.

topology: str: Topology of the coordinate.

class pyscivr.coordinates.core.InternalCoordinates(posi_ref, mass, move_reference=False)[source]

Bases: object

Internal coordinates defined from Cartesian coordinates.

The coordinate system is set up as follows:

Add coordinates using add_bond(), add_angle(), etc.
Call build() to finalize the coordinate system.
Pass the object to other functions to run dynamics.

Parameters:

posi_ref (ArrayLike) – (N, 3) reference positions of the atoms.
mass (ArrayLike) – (N,) atomic masses.
move_reference (bool) – If True, the reference positions are updated to the current positions at every step (time-dependent Eckart conditions Eq. (85) in the paper). The rotational constraints then do not affect the dynamics, and the angular momentum is conserved. They are used only to remove the rotational components in the coordinate transformation, assuming zero angular momentum (L = 0). See section III C in the paper. If False, posi_ref is used as a fixed reference. (Eckart conditions Eq. (84) in the paper)

add_local(kind, atoms, name=None)[source]

Add a local coordinate defined by a small number of atoms.

Parameters:

kind (str) – Kind of the coordinate.
atoms (tuple[int, ...]) – Indices of the atoms defining the coordinate, in the same order as required by the coordinate function.
name (str | None) – Name of the coordinate. If None, a default name is generated based on kind and atoms.

Return type:

None

add_bond(i, j, name=None)[source]

Add a bond length between atoms i and j.

Parameters:

i (int) – Index of the first atom.
j (int) – Index of the second atom.
name (str | None) – Name of the coordinate. If None, bond(i, j) is used as the default name.

Return type:

None

add_angle(i, j, k, name=None)[source]

Add a bond angle i-j-k at atom j.

Parameters:

i (int) – Index of the first atom.
j (int) – Index of the second atom (central atom of the angle).
k (int) – Index of the third atom.
name (str | None) – Name of the coordinate. If None, angle(i, j, k) is used as the default name.

Return type:

None

add_dihedral(i, j, k, l, name=None)[source]

Add a dihedral angle i-j-k-l.

Parameters:

i (int) – Index of the first atom.
j (int) – Index of the second atom.
k (int) – Index of the third atom.
l (int) – Index of the fourth atom.
name (str | None) – Name of the coordinate. If None, dihedral(i, j, k, l) is used as the default name.

Return type:

None

add_polar(i, j, name=None)[source]

Add a polar angle of the vector i-j from the z-axis.

Parameters:

i (int) – Index of the first atom.
j (int) – Index of the second atom.
name (str | None) – Name of the coordinate. If None, polar(i, j) is used as the default name.

Return type:

None

add_azimuthal(i, j, name=None)[source]

Add an azimuthal angle of the vector i-j from the x-axis in the xy-plane.

Parameters:

i (int) – Index of the first atom.
j (int) – Index of the second atom.
name (str | None) – Name of the coordinate. If None, azimuthal(i, j) is used as the default name.

Return type:

None

add_translation()[source]

Add three (Eckart) translational coordinates (x, y, z) of the center of mass.

\[\sum_i m_i (\mathbf{r}_i - \mathbf{R}_i)\]

where \(\mathbf{R}_i\) is the reference position.

Return type:: None

add_rotation()[source]

Add three (Eckart) rotational coordinates (x, y, z).

\[\sum_i m_i (\mathbf{R}_i - \mathbf{R}_{\text{com}}) \times (\mathbf{r}_i - \mathbf{r}_{\text{com}})\]

where \(\mathbf{R}_i\) is the reference position, and \(\mathbf{R}_{\text{com}}\) and \(\mathbf{r}_{\text{com}}\) are the centers of mass of the reference and current positions.

Return type:: None

add_eckart()[source]

Add Eckart translational and rotational coordinates.

Return type:: None

add_euler()[source]

Add three Euler angles (yaw, pitch, roll) of the overall rotation.

Return type:: None

add_normal_mode(mode_vector, name=None)[source]

Add a normal mode coordinate.

Parameters:

mode_vector (numpy.typing.ArrayLike) – (3N,) normal mode vector.
name (str | None) – Name of the coordinate. If None, a default name is generated based on the number of existing normal modes.

Return type:

None

build()[source]

Finalize the coordinate system.

The coordinates are sorted in the order of KIND_ORDER, so the row order after build() can differ from the order in which they were added. Use find_coordinate() or inv_argsort to find the new row indices.

After the call, the following attributes are available (x is the flattened (3N,) positions):

names, index, topology, atom_indices: coordinate information.
argsort, inv_argsort: mapping between the original and sorted orders.
get_value(x): (K,) raw coordinate values.
get_const(x): (K,) residuals value - value_ref, wrapped to [-pi, pi) for S1 topology.
get_const_grad(x): (K, 3N) gradient B of the residuals (Eq. (15)).
get_const_hess_kj_contract(x, w, V): (3N, X) contraction of the second derivatives (Eq. (34)).
get_const_hess_ij_contract(x, w, V): (K, X) contraction of the second derivatives (Eq. (36)).
get_value_batch, get_const_batch, get_const_grad_batch: batched versions over the first axis.

Return type:: None

find_coordinate(kind, atoms)[source]

Find a local coordinate after build() from its kind and atom indices.

Parameters:

kind (str) – Kind of the coordinate (e.g. "bond", "dihedral").
atoms (Sequence[int]) – Indices of the atoms defining the coordinate, in the same order as given when it was added.

Returns:

(index, value_ref, name, topology), where index is the row index after build(), which can be used in active_idx.

Raises:

RuntimeError – If build() has not been called.
KeyError – If the coordinate is not found.
ValueError – If more than one coordinate matches.

Return type:

tuple[int, float, str, str]

register_kind(kind, func, n_atoms, topology)[source]

Register a new kind of local coordinate.

Parameters:

kind (str) – Kind of the coordinate (e.g. "bond", "dihedral").
func (Callable[[Array], Array]) – Function that computes the coordinate value from the positions of the atoms. It should take an array of shape (n_atoms, 3) and return a scalar.
n_atoms (int) – Number of atoms defining the coordinate.
topology (str) – Topology of the coordinate (e.g. S1, R, etc.).

Return type:

None

Examples

For example, to fix the difference of two bond lengths, one can define a new coordinate kind "bond_diff" as follows:

>>> def bond_diff(p: jax.Array) -> jax.Array:
...     return jnp.linalg.norm(p[1] - p[0]) - jnp.linalg.norm(p[2] - p[0])
>>> coord = InternalCoordinates(posi_ref, mass)
>>> coord.register_kind("bond_diff", bond_diff, 3, R)
>>> coord.add_local("bond_diff", (0, 1, 2))

class pyscivr.coordinates.core.QuaternionCoordinates[source]

Bases: object

Euler angles defined from quaternion coordinates.

This object can be used in place of InternalCoordinates to define the Euler angles.

The coordinates are (yaw, pitch, roll, norm), where norm is the normalization condition (q^T q - 1)/2 of the quaternion. The functions take the (4,) quaternion in place of the (3N,) positions.

build()[source]

Do nothing, since the coordinate system is already defined.

Return type:: None

pyscivr.coordinates.core.make_converters(coord_all, active_idx=None)[source]

Make coordinate converters for the active (physical) coordinates.

Parameters:

coord_all (InternalCoordinates | QuaternionCoordinates) – InternalCoordinates for the full system. It must define 3N coordinates for coordinate transformation.
active_idx (TypeAliasForwardRef('numpy.typing.ArrayLike') | None) – (m,) row indices of the active coordinates. If None, all coordinates are regarded as active.

Returns:

(to_physical_posi, to_physical_mome, to_redundant_mome, get_R_active), where the first three are jitted and batched over the first axis:

to_physical_posi(posi): (n_batch, 3N) -> (n_batch, m).
to_physical_mome(posi, mome): (n_batch, 3N), (n_batch, 3N) -> (n_batch, m), P_A = A_A^T p.
to_redundant_mome(posi, mome_int): (n_batch, 3N), (n_batch, m) -> (n_batch, 3N), p = R_A P_A.
get_R_active(x): (3N,) -> (3N, m) (Eq. (41)), not batched.

For QuaternionCoordinates, 3N is replaced by 4.

Return type:

tuple[Callable[[Array], Array], Callable[[Array, Array], Array], Callable[[Array, Array], Array], Callable[[Array], Array]]

pyscivr.coordinates.core.quaternion_metric(inertia)[source]

Make the metric function for the quaternion dynamics (Eq. (L5) in the paper).

Parameters:: inertia (numpy.typing.ArrayLike) – (3,) principal moments of inertia.
Returns:: Function to compute the metric, (4,) quaternion -> (4, 4) matrix. It can be passed to pyscivr.integrators.ImplicitVerlet as metric_func.
Return type:: Callable[[Array], Array]