pyscivr.variations module

Propagation of the variations (monodromy matrices) along classical trajectories.

See section II C and appendices E and J 3 of the paper for details.

Notation:

• N: number of atoms.

• T: number of time steps of the trajectory.

• K: number of constraints.

• m: number of DOFs of the active coordinates.

• Redundant space: (3N,)-dimensional space of Cartesian coordinates with K constraints.

• Physical space: (m,)-dimensional space of the active coordinates.

pyscivr.variations.make_vari_propagator(coord_const, dt)

Make a function to propagate the variation along a SHAKE/RATTLE trajectory.

The linearized SHAKE/RATTLE equations (Eqs. (31)-(37) in the paper) are integrated along the trajectory.

Parameters:

• coord_const (InternalCoordinates) – InternalCoordinates defining the K constraints, the same as used in pyscivr.integrators.ShakeRattle.

• dt (float) – time step of the trajectory.

Returns:

Jitted function vari_propagator(posi, mome, lam, mu, hess, vari_init) returning the (T, 6N, 2m) propagated variation of the trajectory. Use it through propagate_variation().

Return type:

Callable[…, jax.Array]

pyscivr.variations.propagate_variation(vari_propagator, posi, mome, lam, mu, hess, vari_init)

Propagate the variation along a SHAKE/RATTLE trajectory by the linearized equations of motion.

Parameters:

• vari_propagator (Callable[..., jax.Array]) – Function made by make_vari_propagator().

• posi (ArrayLike) – (T, N, 3) positions of the trajectory.

• mome (ArrayLike) – (T, N, 3) momenta of the trajectory.

• lam (ArrayLike) – (T, K) Lagrange multipliers for positional constraints.

• mu (ArrayLike) – (T, K) Lagrange multipliers for momentum constraints.

• hess (ArrayLike) – (T, 3N, 3N) (projected) Hessian of the potential energy at each time step.

• vari_init (ArrayLike) – (6N, 2m) initial variation, e.g. get_jacobi_rev(posi[0], mome[0]) from make_jacobi_func().

Returns:

(T, 6N, 2m) propagated variation of the trajectory.

Return type:

NDArray[np.float64]

Note

Instead of the (6N, 6N) full monodromy matrix \(M_t^{\text{red}}\) (Eq. (32)), the variation in the redundant space is retained as the (6N, 2m) matrix product \(M_t^{\text{red}} J_{\text{rev}}(0)\) (Eq. (42)). The variation in the physical space can be obtained by \(J(t) (M_t^{\text{red}} J_{\text{rev}}(0))\) (Eq. (44)), and the initial variation in the redundant space is \(J_{\text{rev}}(0)\), since \(J_{\text{rev}}(0) J(0) J_{\text{rev}}(0) = J_{\text{rev}}(0)\) (Eq. (47)).

pyscivr.variations.make_vari_propagator_implicit(metric_func, dt, *, potential_func=None, normalize=False)

Make a function to propagate the variation along an implicit Verlet trajectory.

For \(H(Q, P) = \frac{1}{2} P^T G(Q) P + V(Q)\), with

\[X(Q, P) = \frac{\partial (G(Q) P)}{\partial Q}, \quad Y(Q, P) = \frac{\partial^2 H}{\partial Q^2}\]

the linearized implicit Verlet equations (Eqs. (J13)-(J14) in the paper) are

\[ \begin{align}\begin{aligned}&(I + \frac{\Delta t}{2} X(Q_n, P_{n+\frac{1}{2}})^T) \delta P_{n+\frac{1}{2}} = \delta P_{n} - \frac{\Delta t}{2} Y(Q_n, P_{n+\frac{1}{2}}) \delta Q_n\\&(I - \frac{\Delta t}{2} X(Q_{n+1}, P_{n+\frac{1}{2}})) \delta Q_{n+1} = (I + \frac{\Delta t}{2} X(Q_n, P_{n+\frac{1}{2}})) \delta Q_n + \frac{\Delta t}{2} (G_n + G_{n+1}) \delta P_{n+\frac{1}{2}}\\&\delta P_{n+1} = (I - \frac{\Delta t}{2} X(Q_{n+1}, P_{n+\frac{1}{2}})^T) \delta P_{n+\frac{1}{2}} - \frac{\Delta t}{2} Y(Q_{n+1}, P_{n+\frac{1}{2}}) \delta Q_{n+1}\end{aligned}\end{align} \]

Parameters:

• metric_func (Callable[[Array], Array]) – Function to compute the metric \(G(Q)\), (m,) -> (m, m), the same as used in pyscivr.integrators.ImplicitVerlet.

• dt (float) – Time step.

• potential_func (Callable[[Array], Array] | None) – Function to compute the potential \(V(Q)\), (m,) -> scalar. If None, the potential is assumed to be zero.

• normalize (bool) – If True, the variations are projected at every step to satisfy the linearized constraints q^T dq = 0 and q^T dp + p^T dq = 0 of q^T q = 1 and p^T q = 0. This is intended for the quaternion dynamics.

Returns:

Jitted function vari_propagator(posi, mome, mome_half, vari_init) returning the (T, 2m, 2m) propagated variation of the trajectory. Use it through propagate_variation_implicit().

Return type:

Callable[[…], Array]

pyscivr.variations.propagate_variation_implicit(vari_propagator, posi, mome, mome_half, vari_init)

Propagate the variation along an implicit Verlet trajectory by the linearized equations of motion.

Parameters:

• vari_propagator (Callable[..., jax.Array]) – Function made by make_vari_propagator_implicit().

• posi (ArrayLike) – (T, m) positions of the trajectory.

• mome (ArrayLike) – (T, m) momenta of the trajectory.

• mome_half (ArrayLike) – (T, m) momenta at half time steps of the trajectory.

• vari_init (ArrayLike) – (2m, 2m) initial variation, e.g. the identity matrix I_{2m}.

Returns:

(T, 2m, 2m) propagated variation of the trajectory (monodromy matrix).

Return type:

NDArray[np.float64]

pyscivr.variations.make_jacobi_func(coord_all, active_idx)

Make Jacobian functions between the redundant and physical spaces (Eq. (42) in the paper).

Parameters:

• coord_all (InternalCoordinates) – InternalCoordinates for the full system. It must define 3N coordinates for coordinate transformation.

• active_idx (ArrayLike) – (m,) row indices of the active coordinates in coord_all, in the row order after build().

Returns:

(get_jacobi, get_jacobi_rev, get_internal, get_monodromy), where all functions are jitted:

• get_jacobi(posi, mome): (2m, 6N) Jacobian J from the redundant space to the physical space.

• get_jacobi_rev(posi, mome): (6N, 2m) Jacobian J_rev from the physical space to the redundant space.

• get_internal(posi, mome): (m,) positions, (m,) momenta in the physical space.

• get_monodromy(posi_flat, mome_flat, vari): batched over T. Takes (T, 3N) positions, (T, 3N) momenta, and (T, 6N, 2m) variations in the redundant space, and returns (T, m) positions, (T, m) momenta, (T, 2m, 2m) monodromy matrices, and (T,) symplectic errors in the physical space. Use it through to_internal().

Return type:

tuple[Callable[…, jax.Array], Callable[…, jax.Array], Callable[…, tuple[jax.Array, jax.Array]], Callable[…, tuple[jax.Array, …]]]

pyscivr.variations.to_internal(monodromy_func, posi, mome, vari)

Transform the position, momentum, and variations from the redundant space to the physical space.

The monodromy matrix in the physical space is obtained as \(M^{\text{phys}} = J(t) (M_t^{\text{red}} J_{\text{rev}}(0))\).

Parameters:

• monodromy_func (Callable[..., tuple[jax.Array, ...]]) – get_monodromy function made by make_jacobi_func().

• posi (ArrayLike) – (T, N, 3) positions in the redundant space.

• mome (ArrayLike) – (T, N, 3) momenta in the redundant space.

• vari (ArrayLike) – (T, 6N, 2m) variations in the redundant space.

Returns:

(posi_internal, mome_internal, mono_internal, symplectic_err), where

• posi_internal: (T, m) positions in the physical space.

• mome_internal: (T, m) momenta in the physical space.

• mono_internal: (T, 2m, 2m) monodromy matrix in the physical space.

• symplectic_err: (T,) symplectic errors \(\epsilon_t = \max |M^{\text{phys}}(t)^T \Omega M^{\text{phys}}(t) - \Omega|\), where \(\Omega \equiv \begin{pmatrix} 0 & I \\ -I & 0 \end{pmatrix}\) is the symplectic matrix (Eq. (82)).

Return type:

tuple[NDArray[np.float64], …]