# Copyright (C) 2025 The authors of Py-ChemShell
#
# This file is part of Py-ChemShell.
#
# Py-ChemShell is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# Py-ChemShell is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with Py-ChemShell. If not, see
# <http://www.gnu.org/licenses/>.
r"""Restraint potentials.
Restraints are geometry-dependent potentials that are employed to bias
simulated systems towards selected structural features. These are used
whenever it is desirable to *encourage* rather than *enforce* a
particular geometry; rigid constraints, by contrast, eliminate the
corresponding degrees of freedom entirely. By introducing supplementary
energy terms that penalise deviations from specified reference values,
such as inter-atomic distances or bond angles, experimental
observations or prior theoretical insights may be honoured, enhanced
convergence can be promoted, and rare events may be explored within
practicable timescales; all whilst permitting natural fluctuations.
Typically, restraints are applied only to a limited selection of
atoms, leaving the remainder of the model to evolve according to the
underlying force-field. In this way, a balance is struck between
preserving the natural dynamical character of the system and focusing
computational effort on the configurational regions of greatest
significance to the study at hand.
Restraints are defined and applied at the :module:`chemsh.tasks` layer.
Upon execution, each task—such as :class:`chemsh.tasks.sp.SP` or
:class:`chemsh.tasks.opt.Opt`—automatically iterates over the supplied
restraints, evaluating them in sequence. All concrete restraint classes
inherit from the abstract base :class:`.Restraint`. Accordingly, an
instance of the appropriate concrete class is required for every
restraint to be used, and a list of these instances should be supplied
to the task via the `restraints` keyword argument.
The `chemsh` package offers harmonic restraint potentials for the six
following geometric properties:
- :class:`.Bond`: bond-distance between two atoms.
- :class:`.Angle`: bond‑angle involving three atoms.
- :class:`.Torsion`: torsion‑angle defined by four atoms.
- :class:`.BondDifference2`: difference of two bond distances.
- :class:`.BondDifference3`: difference of three bond distances.
- :class:`.BondDifference4`: difference of four bond distances.
Examples
--------
A minimum working example is provided below demonstrating the
application of three harmonic restraints to a single-point task:
>>> import chemsh
>>> from chemsh import Fragment, NWChem, SP
>>> from chemsh.base.restraint import Bond, Angle, Torsion
>>> fragment = Fragment("some_molecule_file.xyz")
>>> theory = NWChem(frag=fragment, method="hf", basis="3-21g")
>>> restraints = [
... Bond((0, 8), 2.4, 2.0),
... Angle((8, 0, 2), 115.0, 2.0),
... Torsion((0, 2, 3, 1), -50.0, 2.0)]
>>> task = SP(frag=fragment, theory=theory,
... restraints=restraints, gradients=True)
>>> task.run()
"""
__author__ = ('Adam J. McSloy <a.mcsloy@bristol.ac.uk')
from abc import ABC, abstractmethod
from typing import Optional, cast, TypeAlias
import numpy as np
from numpy.typing import NDArray
import chemsh
from chemsh.objects.fragment import Fragment
Indices2: TypeAlias = tuple[int, int]
Indices3: TypeAlias = tuple[int, int, int]
Indices4: TypeAlias = tuple[int, int, int, int]
Indices6: TypeAlias = tuple[int, int, int, int, int, int]
Indices8: TypeAlias = tuple[int, int, int, int, int, int, int, int]
[docs]
class Restraint(ABC):
"""Abstract base class for geometric restraints.
This abstract base class defines the interface method
":meth:`~chemsh.base.restraint.Restraint.accumulate`" which all
restraint classes must implement. It is through this interface that
restraints will be evaluated. The interface method is tasked with
1) evaluating the restraint energy and, optionally, its forces, and
2) accumulate the results into the supplied array(s).
"""
[docs]
@abstractmethod
def accumulate(
self, frag: Fragment, energy: NDArray[float],
gradient: Optional[NDArray[float]] = None) -> None:
"""Accumulate the restraint energy and, optionally, its forces.
Parameters
----------
frag : Fragment
The molecular fragment whose Cartesian coordinates are to
be restrained.
energy : NDArray[float]
A length‑1 array into which the restraint energy is to be
added.
gradient : Optional[NDArray[float]]
An optional array, of shape ``frag.coords``, into which the
associated gradients are to be added. If no array is
provided, as is the default, then the gradient values will
not be computed.
"""
pass
[docs]
class Bond(Restraint):
r"""A harmonic bond-distance restraint between two atoms.
The potential
.. math::
E = \tfrac12 k \left(r - r_{0}\right)^{2}
is applied, where *r* denotes the present inter‑atomic distance,
:math:`r_0` the target distance, and :math:`k` the spring constant.
Attributes
----------
r0 : float
Equilibrium distance :math:`r_{0}` in Bohr.
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
[docs]
def __init__(self, indices: Indices2, r0: float, k: float):
r"""Initialise harmonic bond distance restraint instance.
Parameters
----------
indices : Indices2
A tuple specifying the atomic indices of the two atoms to be
restrained.
r0 : float
Equilibrium distance :math:`r_{0}` (positive, Bohr).
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
if len(indices) != 2:
raise ValueError(
f"Two indices must be given not {len(indices)}.")
self._indices: NDArray[int] = np.array(indices, dtype=int)
self.r0: float = r0
self.k: float = k
@property
def indices(self) -> Indices2:
"""Atomic indices of the two target atoms."""
return cast(Indices2, tuple(self._indices.tolist()))
@indices.setter
def indices(self, indices: Indices2):
if len(indices) != 2:
raise ValueError(
f"Two indices must be given not {len(indices)}.")
self._indices[:] = indices
[docs]
def accumulate(
self, frag: Fragment, energy: NDArray[float],
gradient: Optional[NDArray[float]] = None) -> None:
"""Evaluate and accumulate the bond-distance restraint.
Parameters
----------
frag : Fragment
The molecular fragment whose Cartesian coordinates are to
be restrained.
energy : NDArray[float]
A length‑1 array into which the restraint energy is to be
added.
gradient : Optional[NDArray[float]]
An optional array, of shape ``frag.coords``, into which the
associated gradients are to be added. If no array is
provided, as is the default, then the gradient values will
not be computed.
"""
v12 = frag.coords[self._indices[0]] - frag.coords[self._indices[1]]
r12 = np.linalg.norm(v12)
dr12 = r12 - self.r0
# Evaluate the bond restraint energy and add the result to the
# supplied energy array.
energy[0] += 0.5 * self.k * (dr12 * dr12)
# Repeat this process with the forces, but only if requested
# to do so. Note that forces are expanded into their vector
# components.
if gradient is not None:
gradient_ = -self.k * dr12 * (v12 / r12)
gradient[self._indices[0]] -= gradient_
gradient[self._indices[1]] += gradient_
if chemsh.CHEMSH_DEBUG >= 5:
print(
f"Restraint - {self}: ",
f"r={r12:5.2f}, "
f"energy={0.5 * self.k * (dr12 * dr12):11.4E}, "
f"gradient={dr12 * self.k:11.4E}",
sep="")
def __repr__(self) -> str:
return (f'{self.__class__.__name__}('
f'{self.indices}, r0={self.r0:5.2f}, k={self.k:5.2f})')
[docs]
class Angle(Restraint):
r"""A harmonic bond‑angle restraint involving three atoms.
The potential
.. math::
E = \tfrac12 k (\theta - \theta_{0})^{2}
penalises deviations of the bond angle :math:`\theta` from the
target angle :math:`\theta_{0}`.
Attributes
----------
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{rad}^{2}`.
"""
[docs]
def __init__(self, indices: Indices3, theta0: float, k: float):
r"""Initialise harmonic bond angle restraint instance.
Parameters
----------
indices : Indices3
A tuple specifying the atomic indices of the three atoms to
be restrained. The bond angle is defined as the angle
between the vectors i-j & j-k.
theta0 : float
Target angle :math:`\theta_{0}` in degrees, restricted to
the interval [0°, 180°].
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{rad}^{2}`.
"""
if len(indices) != 3:
raise ValueError(
f"Three indices must be given not {len(indices)}.")
self._indices: NDArray[int] = np.array(indices, dtype=int)
self._theta0: float = self._sanitise_angle(theta0)
self.k: float = k
@property
def indices(self) -> Indices3:
"""Atomic indices of the three target atoms."""
return cast(Indices3, tuple(self._indices.tolist()))
@indices.setter
def indices(self, indices: Indices3):
if len(indices) != 3:
raise ValueError(
f"Three indices must be given not {len(indices)}.")
self._indices[:] = indices
@property
def theta0(self) -> float:
r"""Equilibrium angle, :math:`\theta_{0}`."""
return np.rad2deg(self._theta0)
@theta0.setter
def theta0(self, value: float):
self._theta0 = self._sanitise_angle(value)
[docs]
def accumulate(
self, frag: Fragment, energy: NDArray[float],
gradient: Optional[NDArray[float]] = None) -> None:
"""Evaluate and accumulate the bond-angle restraint.
Parameters
----------
frag : Fragment
The molecular fragment whose Cartesian coordinates are to
be restrained.
energy : NDArray[float]
A length‑1 array into which the restraint energy is to be
added.
gradient : Optional[NDArray[float]]
An optional array, of shape ``frag.coords``, into which the
associated gradients are to be added. If no array is
provided, as is the default, then the gradient values will
not be computed.
"""
i, j, k = self._indices
v12 = frag.coords[i] - frag.coords[j]
v32 = frag.coords[k] - frag.coords[j]
r12 = np.linalg.norm(v12)
r32 = np.linalg.norm(v32)
u12, u32 = v12 / r12, v32 / r32
cos_theta = u12 @ u32
theta = np.arccos(cos_theta)
sin_theta = np.sin(theta)
theta_d = theta - self._theta0
energy[0] += 0.5 * self.k * (theta_d * theta_d)
if gradient is not None:
gradient_i = ((u12 * cos_theta - u32) * theta_d
* self.k / (r12 * sin_theta))
gradient_k = ((u32 * cos_theta - u12) * theta_d
* self.k / (r32 * sin_theta))
gradient[i] += gradient_i
gradient[j] -= gradient_i + gradient_k
gradient[k] += gradient_k
if chemsh.CHEMSH_DEBUG >= 5:
print(
f"Restraint - {self}: ",
f"theta={np.rad2deg(theta):7.2f}, "
f"energy={0.5 * self.k * (theta_d * theta_d):11.4E}, "
f"gradient={theta_d * self.k:11.4E}",
sep="")
@staticmethod
def _sanitise_angle(angle: float):
r"""Validate `angle` and convert it to radians.
Parameters
----------
angle : float
Angle in degrees in the interval [0°, 180°].
Returns
-------
angle_sanitised : float
The sanitised angle in radians.
Raises
------
ValueError
If `angle` lies outside the closed interval [0°, 180°].
"""
if not (0.0 <= angle <= 180.0):
raise ValueError(f"Invalid bond restraint angle "
f"encountered; {angle}° ∉ [0°, 180°].")
return np.deg2rad(angle)
def __repr__(self) -> str:
return (f'{self.__class__.__name__}('
f'{self.indices}, theta0={self.theta0:7.2f}, k={self.k:5.2f})')
[docs]
class Torsion(Restraint):
r"""A harmonic torsion‑angle restraint involving four atoms.
The potential
.. math::
E = \tfrac12 k (\phi - \phi_{0})^{2}
is used, where :math:`\phi` is the proper dihedral angle and
:math:`\phi_{0}` the target value.
Attributes
----------
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{rad}^{2}`.
"""
[docs]
def __init__(self, indices: Indices4, phi0: float,
k: float):
r"""Initialise harmonic torsion-angle restraint instance.
Parameters
----------
indices : Indices4
A tuple specifying the atomic indices of the four atoms to
be restrained. The torsion angle is defined as the angle
between the plains i-j-k & j-k-l.
phi0 : float
Target dihedral :math:`\phi_{0}` in degrees, within the
interval [-180°, +180°].
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{rad}^{2}`.
"""
if len(indices) != 4:
raise ValueError(
f"Four indices must be given not {len(indices)}.")
self._indices = np.array(indices, dtype=int)
self.k: float = k
self._phi0: float = self._sanitise_angle(phi0)
@property
def indices(self) -> Indices4:
"""Atomic indices of the four target atoms."""
return cast(Indices4, tuple(self._indices.tolist()))
@indices.setter
def indices(self, indices: Indices4):
if len(indices) != 4:
raise ValueError(
f"Four indices must be given not {len(indices)}.")
self._indices[:] = indices
@property
def phi0(self) -> float:
r"""Equilibrium torsion angle, :math:`\phi_{0}`."""
return np.rad2deg(self._phi0)
@phi0.setter
def phi0(self, value: float):
self._phi0 = self._sanitise_angle(value)
[docs]
def accumulate(
self, frag: Fragment, energy: NDArray[float],
gradient: Optional[NDArray[float]] = None) -> None:
"""Evaluate and accumulate the torsion-angle restraint.
Parameters
----------
frag : Fragment
The molecular fragment whose Cartesian coordinates are to
be restrained.
energy : NDArray[float]
A length‑1 array into which the restraint energy is to be
added.
gradient : Optional[NDArray[float]]
An optional array, of shape ``frag.coords``, into which the
associated gradients are to be added. If no array is
provided, as is the default, then the gradient values will
not be computed.
"""
if gradient is not None:
phi, dphi_dr = self.torsion_angle_and_derivative(
frag.coords[self._indices])
else:
phi = self.torsion_angle(frag.coords[self._indices])
# Signed difference between the current and target torsion
# angle, wrapped into the interval [–π, π) ensuring minimum arc
# distance.
delta_phi = (phi - self._phi0 + np.pi) % (2 * np.pi) - np.pi
# Compute and accumulate energy
energy[0] += 0.5 * self.k * (delta_phi * delta_phi)
if gradient is not None:
gradient[self._indices] += self.k * delta_phi * dphi_dr
if chemsh.CHEMSH_DEBUG >= 5:
print(
f"Restraint - {self}: ",
f"phi={np.rad2deg(phi):7.2f}, "
f"energy={0.5 * self.k * (delta_phi * delta_phi):11.4E}, "
f"gradient={delta_phi * self.k:11.4E}",
sep="")
[docs]
@staticmethod
def torsion_angle(r: NDArray[float]):
"""Proper dihedral angle of atoms A–B–C–D.
Parameters
----------
r : NDArray[float]
A 4×3 array specifying the cartesian coordinates of atoms
A, B, C, and D.
Returns
-------
phi : float
Dihedral angle in radians, in the range [-π, π].
"""
# bond vectors
r_kj = r[2] - r[1]
# normals to the two planes
r_mj = np.cross(r[0] - r[1], r_kj)
r_nk = np.cross(r_kj, r[2] - r[3])
# dihedral angle via atan2 (robust against round-off)
phi = np.arctan2(
(r_kj / np.linalg.norm(r_kj)) @ np.cross(r_mj, r_nk),
r_mj @ r_nk)
return phi
[docs]
@staticmethod
def torsion_angle_and_derivative(r: NDArray[float]):
r"""Proper dihedral angle and its cartesian gradient.
Parameters
----------
r : NDArray[float]
A 4×3 array specifying the cartesian coordinates of atoms
A, B, C, and D.
Returns
-------
phi : float
Dihedral angle in radians, in the range [-π, π].
dphi_dr : NDArray[float]
A 4×3 array storing the analytical gradients
:math:`\tfrac{\partial\phi}{r}`, row-wise for A, B, C,
and D.
"""
# bond vectors
r_ij = r[0] - r[1]
r_kj = r[2] - r[1]
r_kl = r[2] - r[3]
# normals to the two planes
r_mj = np.cross(r_ij, r_kj)
r_nk = np.cross(r_kj, r_kl)
len_kj = np.linalg.norm(r_kj)
len_mj = np.linalg.norm(r_mj)
len_nk = np.linalg.norm(r_nk)
# dihedral angle via atan2 (robust against round-off)
x = r_mj @ r_nk
y = (r_kj / len_kj) @ np.cross(r_mj, r_nk)
phi = np.arctan2(y, x)
# Analytical gradient (MODELLER manual 9v6 A.57-A.60)
r_ij_kj_kj2 = (r_ij @ r_kj) / (len_kj * len_kj)
r_kl_kj_kj2 = (r_kl @ r_kj) / (len_kj * len_kj)
dp_dri = (len_kj / (len_mj * len_mj)) * r_mj
dp_drl = -(len_kj / (len_nk * len_nk)) * r_nk
dp_drj = (r_ij_kj_kj2 - 1) * dp_dri - r_kl_kj_kj2 * dp_drl
dp_drk = (r_kl_kj_kj2 - 1) * dp_drl - r_ij_kj_kj2 * dp_dri
dphi_dr = np.stack((dp_dri, dp_drj, dp_drk, dp_drl))
return phi, dphi_dr
@staticmethod
def _sanitise_angle(angle: float):
r"""Validate & convert angle to the internal representation.
Parameters
----------
angle : float
A dihedral in degrees, within the interval [-180°, +180°].
Returns
-------
angle_sanitised : float
The sanitised torsion angle in radians.
Raises
------
ValueError
If ``angle`` lies outside the interval [-180°, +180°].
"""
if not (-180 <= angle <= +180):
raise ValueError(f"Invalid torsion restraint angle "
f"encountered; {angle}° ∉ [-180°, +180°].")
return np.deg2rad(angle)
def __repr__(self) -> str:
return (f'{self.__class__.__name__}('
f'{self.indices}, phi0={self.phi0:7.2f}, k={self.k:5.2f})')
def _bond_vector_helper(
array: NDArray[float], indices: NDArray[int]
) -> tuple[NDArray[float], NDArray[float], NDArray[float]]:
"""Compute bond vectors, distances, and their unit vectors.
This primarily operates as a helper function for the bond
distance differences restraint classes.
Parameters
----------
array : NDArray[float]
An Nx3 array of positions from which atomic coordinates can
be sourced.
indices : NDArray[int]
A flattened array of atomic indices providing a pair of indices
for each bond whose properties are to be computed.
Returns
-------
vs : NDArray[float]
A Bx3 array containing the bond vectors, one for each
bond "B".
us : NDArray[float]
The associated unit vectors.
rs : NDArray[float]
An array of length "B" reporting the bond distances for
each evaluated bond "B".
"""
vs = np.subtract.reduce(array[indices.reshape(-1, 2)], axis=1)
rs = np.linalg.norm(vs, axis=-1)
us = vs / rs[:, np.newaxis]
return vs, us, rs
[docs]
class BondDifference2(Restraint):
r"""A harmonic restraint on the difference of two bond distances.
The potential of the two bond difference restraint is evaluated as
follows:
.. math::
E = \tfrac12 k\left(\left(r_{12}-r_{34}\right)-r_0\right)^2
Where :math:`r_{ij}` is the distance between the two points
:math:`\bf{r}_i` and :math:`\bf{r}_i`, :math:`k` is the spring
constant, and :math:`r_0` is the equilibrium distance difference.
Attributes
----------
r0 : float
Equilibrium distance difference, in units of Bohr, computed as
:math:`r_{12}-r_{34}`.
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
[docs]
def __init__(self, indices: Indices4, r0: float, k: float):
"""Initialise `BondDifference2` restraint instance.
Parameters
----------
indices : Indices4
A tuple specifying the atomic indices of the four target
atoms ordered such that the first two indices designate the
atoms of the first bond, and the last two the second bond.
r0 : float
Equilibrium distance difference, in units of Bohr, computed
as :math:`r_{12}-r_{34}`.
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
if len(indices) != 4:
raise ValueError(
f"Four indices must be given not {len(indices)}.")
self._indices = np.array(indices, dtype=int)
self.r0: float = r0
self.k: float = k
@property
def indices(self) -> Indices4:
"""Atomic indices of the four target atoms."""
return cast(Indices4, tuple(self._indices.tolist()))
@indices.setter
def indices(self, indices: Indices4):
if len(indices) != 4:
raise ValueError(
f"Four indices must be given not {len(indices)}.")
self._indices[:] = indices
[docs]
def accumulate(
self, frag: Fragment, energy: NDArray[float],
gradient: Optional[NDArray[float]] = None) -> None:
"""Evaluate & accumulate two bond distance difference restraint.
Parameters
----------
frag : Fragment
The molecular fragment whose Cartesian coordinates are to
be restrained.
energy : NDArray[float]
A length‑1 array into which the restraint energy is to be
added.
gradient : Optional[NDArray[float]]
An optional array, of shape ``frag.coords``, into which the
associated gradients are to be added. If no array is
provided, as is the default, then the gradient values will
not be computed.
"""
# Compute bond vectors, unit vectors, and distances for each
# atom pair. Then evaluate the deviation between the current
# and equilibrium reference value.
vs, us, rs = _bond_vector_helper(frag.coords, self._indices)
dr = rs[0] - rs[1] - self.r0
# Evaluate the restraint energy and accumulate the result into
# the supplied energy array.
energy[0] += 0.5 * self.k * (dr * dr)
if gradient is not None:
gradient_ = -self.k * dr * us
np.add.at(gradient, self._indices[[1, 2]], gradient_)
np.subtract.at(gradient, self._indices[[0, 3]], gradient_)
if chemsh.CHEMSH_DEBUG >= 5:
print(
f"Restraint - {self}: ",
f"rd2={rs[0] - rs[1]:5.2f}, "
f"energy={0.5 * self.k * (dr * dr):11.4E}, "
f"gradient={dr * self.k:11.4E}",
sep="")
def __repr__(self) -> str:
return (f'{self.__class__.__name__}('
f'{self.indices}, r0={self.r0:5.2f}, k={self.k:5.2f})')
[docs]
class BondDifference3(Restraint):
r"""A harmonic restraint on the difference of three bond distances.
The potential of the three bond difference restraint is evaluated
as follows:
.. math::
E = \tfrac12 k\left(\left(r_{12}-r_{34}-r_{56}\right)-r_0\right)^2
Where :math:`r_{ij}` is the distance between the two points
:math:`\bf{r}_i` and :math:`\bf{r}_i`, :math:`k` is the spring
constant, and :math:`r_0` is the equilibrium distance difference.
Attributes
----------
r0 : float
Equilibrium distance difference, in units of Bohr, computed as
:math:`r_{12}-r_{34}-r_{56}`.
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
[docs]
def __init__(self, indices: Indices6, r0: float, k: float):
"""Initialise `BondDifference3` restraint instance.
Parameters
----------
indices : Indices6
A tuple specifying the atomic indices of the six target
atoms ordered such that the first two indices designate
the atoms of the first bond, the next two the second bond,
and the final pair those in the third bond.
r0 : float
Equilibrium distance difference, in units of Bohr, computed
as :math:`r_{12}-r_{34}-r_{56}`.
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
if len(indices) != 6:
raise ValueError(
f"Six indices must be given not {len(indices)}.")
self._indices = np.array(indices, dtype=int)
self.r0: float = r0
self.k: float = k
@property
def indices(self) -> Indices6:
"""Atomic indices of the six target atoms."""
return cast(Indices6, tuple(self._indices.tolist()))
@indices.setter
def indices(self, indices: Indices6):
if len(indices) != 6:
raise ValueError(
f"Six indices must be given not {len(indices)}.")
self._indices[:] = indices
[docs]
def accumulate(
self, frag: Fragment, energy: NDArray[float],
gradient: Optional[NDArray[float]] = None) -> None:
"""Evaluate & accumulate three bond distance difference restraint.
Parameters
----------
frag : Fragment
The molecular fragment whose Cartesian coordinates are to
be restrained.
energy : NDArray[float]
A length‑1 array into which the restraint energy is to be
added.
gradient : Optional[NDArray[float]]
An optional array, of shape ``frag.coords``, into which the
associated gradients are to be added. If no array is
provided, as is the default, then the gradient values will
not be computed.
"""
# Compute bond vectors, unit vectors, and distances for each
# atom pair. Then evaluate the deviation between the current
# and equilibrium reference value.
vs, us, rs = _bond_vector_helper(frag.coords, self._indices)
dr = rs[0] - rs[1] - rs[2] - self.r0
# Evaluate the restraint energy and accumulate the result into
# the supplied energy array.
energy[0] += 0.5 * self.k * (dr * dr)
if gradient is not None:
gradient_ = -self.k * dr * us
np.add.at(gradient, self._indices[[1, 2, 4]], gradient_)
np.subtract.at(gradient, self._indices[[0, 3, 5]], gradient_)
if chemsh.CHEMSH_DEBUG >= 5:
print(
f"Restraint - {self}: ",
f"rd3={rs[0] - rs[1] - rs[2]:5.2f}, "
f"energy={0.5 * self.k * (dr * dr):11.4E}, "
f"gradient={dr * self.k:11.4E}",
sep="")
def __repr__(self) -> str:
return (f'{self.__class__.__name__}('
f'{self.indices}, r0={self.r0:5.2f}, k={self.k:5.2f})')
[docs]
class BondDifference4(Restraint):
r"""A harmonic restraint on the difference of four bond distances.
The potential of the four bond difference restraint is evaluated
as follows:
.. math::
E = \tfrac12 k\left(\left(r_{12}+r_{34}-r_{56}-r_{78}\right)-r_0\right)^2
Where :math:`r_{ij}` is the distance between the two points
:math:`\bf{r}_i` and :math:`\bf{r}_i`, :math:`k` is the spring
constant, and :math:`r_0` is the equilibrium distance difference.
Attributes
----------
r0 : float
Equilibrium distance difference, in units of Bohr, computed as
:math:`r_{12}+r_{34}-r_{56}-r_{78}`.
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
[docs]
def __init__(self, indices: Indices8, r0: float, k: float):
"""Initialise `BondDifference4` restraint instance.
Parameters
----------
indices : Indices8
A tuple specifying the atomic indices of the six target
atoms ordered such that the first two indices designate the
atoms of the first bond, the next two the second bond, and
so on.
r0 : float
Equilibrium distance difference, in units of Bohr, computed
as :math:`r_{12}+r_{34}-r_{56}-r_{78}`.
k : float
Harmonic spring constant :math:`k` in units of
:math:`\text{Hartree}/\text{Bohr}^{2}`.
"""
if len(indices) != 8:
raise ValueError(
f"Eight indices must be given not {len(indices)}.")
self._indices = np.array(indices, dtype=int)
self.r0: float = r0
self.k: float = k
@property
def indices(self) -> Indices8:
"""Atomic indices of the eight target atoms."""
return cast(Indices8, tuple(self._indices.tolist()))
@indices.setter
def indices(self, indices: Indices8):
if len(indices) != 8:
raise ValueError(
f"Eight indices must be given not {len(indices)}.")
self._indices[:] = indices
[docs]
def accumulate(
self, frag: Fragment, energy: NDArray[float],
gradient: Optional[NDArray[float]] = None) -> None:
"""Evaluate & accumulate four bond distance difference restraint.
Parameters
----------
frag : Fragment
The molecular fragment whose Cartesian coordinates are to
be restrained.
energy : NDArray[float]
A length‑1 array into which the restraint energy is to be
added.
gradient : Optional[NDArray[float]]
An optional array, of shape ``frag.coords``, into which the
associated gradients are to be added. If no array is
provided, as is the default, then the gradient values will
not be computed.
"""
# Compute bond vectors, unit vectors, and distances for each
# atom pair. Then evaluate the deviation between the current
# and equilibrium reference value.
vs, us, rs = _bond_vector_helper(frag.coords, self._indices)
dr = rs[0] + rs[1] - rs[2] - rs[3] - self.r0
# Evaluate the restraint energy and accumulate the result into
# the supplied energy array.
energy[0] += 0.5 * self.k * (dr * dr)
if gradient is not None:
gradient_ = -self.k * dr * us
np.add.at(gradient, self._indices[[1, 3, 4, 6]], gradient_)
np.subtract.at(gradient, self._indices[[0, 2, 5, 7]], gradient_)
if chemsh.CHEMSH_DEBUG >= 5:
print(
f"Restraint - {self}: ",
f"rd4={rs[0] + rs[1] - rs[2] - rs[3]:5.2f}, "
f"energy={0.5 * self.k * (dr * dr):11.4E}, "
f"gradient={dr * self.k:11.4E}",
sep="")
def __repr__(self) -> str:
return (f'{self.__class__.__name__}('
f'{self.indices}, r0={self.r0:5.2f}, k={self.k:5.2f})')