# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
# Peking University &
# Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Collective variables by position
"""
from inspect import signature
import mindspore.numpy as msnp
from mindspore import Tensor
from mindspore import ops, nn
from mindspore.ops import functional as F
from ...function import any_none, any_not_none
from ...function import get_integer, check_broadcast
from ..colvar import Colvar
from ..atoms import AtomsBase, Vector, get_atoms
[docs]class Torsion(Colvar):
r"""Colvar for torsional (dihedral) angle.
Args:
atoms (AtomsBase): Atoms of shape `(..., 4, D)` to form a torsional angle of shape `(...)` or `(..., 1)`.
Cannot be used with `atoms_a` or `atoms_b`.
Default: ``None``. `D` means spatial dimension of the simulation system. Usually is 3.
atoms_a (AtomsBase): Atoms A with shape `(..., D)` to form a torsional angle of shape `(...)` or `(..., 1)`.
Must be used with `atoms_b`, `atoms_c` and `atoms_d`. Cannot be used with `atoms`.
Default: ``None``.
atoms_b (AtomsBase): Atoms B with shape `(..., D)` to form a torsional angle
of shape `(...)` or `(.., 1)`.
Must be used with `atoms_a`, `atoms_c` and `atoms_d`. Cannot be used with `atoms`.
Default: ``None``.
atoms_c (AtomsBase): Atoms C with shape `(..., D)` to form a torsional angle
of shape `(...)` or `(..., 1)`.
Must be used with `atoms_a`, `atoms_b` and `atoms_d`. Cannot be used with `atoms`.
Default: ``None``.
atoms_d (AtomsBase): Atoms D with shape `(..., D)` to form a torsional angle of shape `(...)` or `(..., 1)`.
Must be used with `atoms_a`, `atoms_b` and `atoms_c`. Cannot be used with `atoms`.
Default: ``None``.
vector1 (Vector): Vector 1 of shape `(..., D)` to form of a torsional angle with shape
`(...)` or `(..., 1)`. Must be used with `vector2`. Cannot be used with Atoms.
Default: ``None``.
vector2 (Vector): Vector 2 of shape `(..., D)` to form of a torsional angle with shape
`(...)` or `(..., 1)`. Must be used with `vector1`. Cannot be used with Atoms.
Default: ``None``.
axis_vector (Vector): Axis vector of shape `(..., D)` to form of a torsional angle with shape
`(...)` or `(..., 1)`. Must be used with `vector1`. Cannot be used with Atoms.
Default: ``None``.
use_pbc (bool): Whether to calculate distance under periodic boundary condition.
Default: ``None``.
batched (bool): Whether the first dimension of the input index in atoms is the batch size.
Default: ``False``.
keepdims (bool): Whether to keep the dimension of the last dimension of vector.
Default: ``False``.
axis (int): Axis to gather the points from coordinate of atoms. Default: -2.
name (str): Name of the Colvar. Default: 'torsion'.
Supported Platforms:
``Ascend`` ``GPU``
Examples:
>>> from sponge import Sponge
>>> from sponge.colvar import Torsion
>>> from sponge.callback import RunInfo
>>> cv_dihedral = Torsion([0, 1, 2, 3])
>>> # system is the Molecule object defined by user.
>>> # energy is the Energy object defined by user.
>>> # opt is the Optimizer object defined by user.
>>> md = Sponge(system, potential=energy, optimizer=opt, metrics={'dihedral': cv_dihedral})
>>> run_info = RunInfo(1000)
>>> md.run(2000, callbacks=[run_info])
[MindSPONGE] Started simulation at 2024-02-19 15:43:11
[MindSPONGE] Step: 1000, E_pot: -117.30916, dihedral: 2.1488183
[MindSPONGE] Step: 2000, E_pot: -131.60872, dihedral: 2.143513
[MindSPONGE] Finished simulation at 2024-02-19 15:44:03
[MindSPONGE] Simulation time: 51.27 seconds.
"""
def __init__(self,
atoms: AtomsBase = None,
atoms_a: AtomsBase = None,
atoms_b: AtomsBase = None,
atoms_c: AtomsBase = None,
atoms_d: AtomsBase = None,
vector1: Vector = None,
vector2: Vector = None,
axis_vector: Vector = None,
use_pbc: bool = None,
batched: bool = False,
keepdims: bool = None,
axis: int = -2,
name: str = 'torsion'
):
super().__init__(
periodic=True,
use_pbc=use_pbc,
name=name,
unit='rad',
)
if any_not_none([atoms, atoms_a, atoms_b, atoms_c, atoms_d]) and \
any_not_none([vector1, vector2]):
raise ValueError('The atoms and vector cannot be used at same time!')
axis = get_integer(axis)
self.atoms = None
self.vector1 = None
self.vector2 = None
self.axis_vector = None
self.split4 = None
if any_not_none([vector1, vector2]):
if any_none([vector1, vector2]):
raise ValueError('vector1 must be used with vector2!')
self.vector1 = vector1
self.vector2 = vector2
self.axis_vector = axis_vector
else:
if atoms is None:
if not all([atoms_a, atoms_b, atoms_c]):
raise ValueError
self.vector1 = Vector(atoms0=atoms_b,
atoms1=atoms_a,
batched=batched,
use_pbc=use_pbc,
axis=axis,
)
self.vector2 = Vector(atoms0=atoms_c,
atoms1=atoms_d,
batched=batched,
use_pbc=use_pbc,
axis=axis,
)
self.axis_vector = Vector(atoms0=atoms_b,
atoms1=atoms_c,
batched=batched,
use_pbc=use_pbc,
axis=axis,
)
else:
if any_not_none([atoms_a, atoms_b, atoms_c, atoms_d]):
raise ValueError('atoms cannot be used with atoms_a, atoms_b and atoms_c!')
self.atoms = get_atoms(atoms, batched)
shape = (1,) + self.atoms.shape
if shape[axis] != 4:
raise ValueError(f'The axis {axis} of atoms must be 3 but got: {shape[axis]}')
self.split4 = ops.Split(axis, 4)
if self.atoms is None:
if self.vector1.ndim > self.vector2.ndim:
new_shape = (1,) * (self.vector1.ndim - self.vector2.ndim)
self.vector2.reshape(new_shape)
if self.vector1.ndim < self.vector2.ndim:
new_shape = (1,) * (self.vector2.ndim - self.vector1.ndim)
self.vector1.reshape(new_shape)
if keepdims is None:
if len(shape) > 1:
keepdims = False
else:
keepdims = True
# (..., D)
shape = check_broadcast(self.vector1.shape, self.vector2.shape)
# (...)
shape = shape[:-1]
if keepdims:
# (..., 1)
shape += (1,)
self._set_shape(shape)
else:
if keepdims is None:
if self.atoms.ndim > 2:
keepdims = False
else:
keepdims = True
# (1, ..., 4, D)
shape = (1,) + self.atoms.shape
# (1, ..., D)
shape = shape[:axis] + shape[axis+1:]
# (...)
shape = shape[1:-1]
if keepdims:
# (..., 1)
shape += (1,)
self._set_shape(shape)
self.squeeze = ops.Squeeze(axis)
self.reduce_sum = ops.ReduceSum(keepdims)
self.keepdims = True
if self.atoms is None and self.axis_vector is None:
self.keepdims = keepdims
self.norm_last_dim = None
# MindSpore < 2.0.0-rc1
if 'ord' not in signature(ops.norm).parameters.keys():
self.norm_last_dim = nn.Norm(-1, self.keepdims)
self.atan2 = ops.Atan2()
def construct(self, coordinate: Tensor, pbc_box: bool = None):
r"""calculate torsional angle.
Args:
coordinate (Tensor): Tensor of shape (B, A, D). Data type is float.
`B` means batchsize, i.e. number of walkers in simulation.
`A` means number of atoms in system.
pbc_box (Tensor): Tensor of shape (B, D). Data type is float.
Default: ``None``.
Returns:
torsion (Tensor): Tensor of shape (B, ...) or (B, ..., 1). Data type is float.
"""
axis_vector = None
if self.atoms is None:
# (B, ..., D)
vector1 = self.vector1(coordinate, pbc_box)
vector2 = self.vector2(coordinate, pbc_box)
if self.axis_vector is not None:
axis_vector = self.axis_vector(coordinate, pbc_box)
else:
# (B, ..., 4, D)
atoms = self.atoms(coordinate, pbc_box)
# (B, ..., 1, D)
pos_a, pos_b, pos_c, pos_d = self.split4(atoms)
# (B, ..., D) <- (B, ..., 1, D)
pos_a = self.squeeze(pos_a)
pos_b = self.squeeze(pos_b)
pos_c = self.squeeze(pos_c)
pos_d = self.squeeze(pos_d)
# (B, ..., D)
vector1 = self.get_vector(pos_b, pos_a, pbc_box)
vector2 = self.get_vector(pos_c, pos_d, pbc_box)
axis_vector = self.get_vector(pos_b, pos_c, pbc_box)
if self.atoms is None and self.axis_vector is None:
# (B, ...) or (B, ..., 1) <- (B, ..., D)
if self.norm_last_dim is None:
dis1 = ops.norm(vector1, None, -1, self.keepdims)
dis2 = ops.norm(vector2, None, -1, self.keepdims)
else:
dis1 = self.norm_last_dim(vector1)
dis2 = self.norm_last_dim(vector2)
dot12 = self.reduce_sum(vector1*vector2, -1)
# (B, ...) or (B, ..., 1)
cos_theta = dot12 / dis1 / dis2
return F.acos(cos_theta)
# (B, ..., D)
vec_a = msnp.cross(vector1, axis_vector)
vec_b = msnp.cross(vector2, axis_vector)
cross_ab = msnp.cross(vec_a, vec_b)
# (B, ..., 1) <- (B, ..., D)
if self.norm_last_dim is None:
axis_dis = ops.norm(axis_vector, None, -1, self.keepdims)
else:
axis_dis = self.norm_last_dim(axis_vector)
# (B, ..., D) = (B, ..., D) / (B, ...,1)
axis_vector *= msnp.reciprocal(axis_dis)
# (B, ...) or (B, ..., 1)
sin_phi = self.reduce_sum(axis_vector*cross_ab, -1)
cos_phi = self.reduce_sum(vec_a*vec_b, -1)
return self.atan2(sin_phi, cos_phi)