# 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.
# ============================================================================
"""
Controller
"""
from typing import Union, Tuple, List
from numpy import ndarray
import mindspore as ms
from mindspore import Tensor
from mindspore.nn import Cell
from mindspore import ops
from mindspore.ops import functional as F
from mindspore.numpy.utils_const import _raise_value_error
from ..system import Molecule
from ..function import functions as func
from ..function.functions import get_integer, get_ms_array, get_arguments
[docs]class Controller(Cell):
r"""
Base class for the controller module in MindSPONGE.
The :class:`sponge.control.Controller` is used in
:class:`sponge.control.Updater` to control the values of
seven variables during the simulation process, including
coordinate, velocity, force, energy, kinetics, virial and pbc_box.
Args:
system(:class:`sponge.system.Molecule`): Simulation system.
control_step(int, optional): Step interval for controller execution. Default: ``1``.
kwargs(dict): Other parameters for extension.
Inputs:
- **coordinate** (Tensor) - Coordinate. Tensor of shape :math:`(B, A, D)`.
Data type is float.
Here :math:`B` is the number of walkers in simulation,
:math:`A` is the number of atoms and
:math:`D` is the spatial dimension of the simulation system, which is usually 3.
- **velocity** (Tensor) - Velocity. Tensor of shape :math:`(B, A, D)`. Data type is float.
- **force** (Tensor) - Force. Tensor of shape :math:`(B, A, D)`. Data type is float.
- **energy** (Tensor) - Energy. Tensor of shape :math:`(B, 1)`. Data type is float.
- **kinetics** (Tensor) - Kinetics. Tensor of shape :math:`(B, D)`. Data type is float.
- **virial** (Tensor) - Virial. Tensor of shape :math:`(B, D)`. Data type is float.
- **pbc_box** (Tensor) - Pressure boundary condition box. Tensor of shape :math:`(B, D)`.
Data type is float.
- **step** (int) - Simulation step. Default: ``0``.
Outputs:
- coordinate, Tensor of shape :math:`(B, A, D)`. Coordinate. Data type is float.
- velocity, Tensor of shape :math:`(B, A, D)`. Velocity. Data type is float.
- force, Tensor of shape :math:`(B, A, D)`. Force. Data type is float.
- energy, Tensor of shape :math:`(B, 1)`. Energy. Data type is float.
- kinetics, Tensor of shape :math:`(B, D)`. Kinetics. Data type is float.
- virial, Tensor of shape :math:`(B, D)`. Virial. Data type is float.
- pbc_box, Tensor of shape :math:`(B, D)`. Periodic boundary condition box.
Data type is float.
Supported Platforms:
``Ascend`` ``GPU``
Examples:
>>> from mindspore import Tensor
>>> from sponge import Sponge, Molecule, ForceField, UpdaterMD, WithEnergyCell
>>> from sponge.control import Controller
>>> from sponge.callback import RunInfo
>>> system = Molecule(template='water.tip3p.yaml')
>>> potential = ForceField(system, parameters='SPCE')
>>> withenergy = WithEnergyCell(system, potential)
>>> class MyController(Controller):
... def construct(self,
... coordinate: Tensor,
... velocity: Tensor,
... force: Tensor,
... energy: Tensor,
... kinetics: Tensor,
... virial: Tensor = None,
... pbc_box: Tensor = None,
... step: int = 0,
... **kwargs):
... return super().construct(coordinate, velocity/100,
... force, energy, kinetics,
... virial, pbc_box, step,
... **kwargs)
>>> velocity = Tensor([[0.1008,0.,0.],[-0.8,0.,0.],[-0.8,0.,0.]])
>>> updater = UpdaterMD(
... system=system,
... time_step=1e-3,
... velocity=velocity,
... integrator='velocity_verlet',
... temperature=300,
... controller=MyController(system)
... )
>>> mini = Sponge(withenergy, optimizer=updater)
>>> run_info = RunInfo(1)
>>> mini.run(5, callbacks=[run_info])
"""
def __init__(self,
system: Molecule,
control_step: int = 1,
**kwargs,
):
super().__init__(auto_prefix=False)
self._kwargs = get_arguments(locals(), kwargs)
self.system = system
self.num_walker = self.system.num_walker
self.num_atoms = system.num_atoms
self.dimension = system.dimension
self.sys_dofs = system.degrees_of_freedom
self.degrees_of_freedom = system.degrees_of_freedom
self.time_step = Tensor(1e-3, ms.float32)
self._coordinate = self.system.coordinate
self._pbc_box = self.system.pbc_box
self.units = self.system.units
self.kinetic_unit_scale = Tensor(self.units.kinetic_ref, ms.float32)
self.press_unit_scale = Tensor(self.units.pressure_ref, ms.float32)
# (B, A)
self.atom_mass = self.system.atom_mass
self.inv_mass = self.system.inv_mass
# (B, A, 1)
self._atom_mass = F.expand_dims(self.atom_mass, -1)
self._inv_mass = F.expand_dims(self.inv_mass, -1)
# (B, 1)
self.system_mass = self.system.system_mass
self.system_natom = self.system.system_natom
self.control_step = get_integer(control_step)
if self.control_step <= 0:
raise ValueError('The "control_step" must be larger than 0!')
self.num_constraints = 0
self.identity = ops.Identity()
self.keepdims_sum = ops.ReduceSum(True)
@property
def boltzmann(self) -> float:
"""
Boltzmann constant in current unit.
Returns:
float, Boltzmann constant in current unit.
"""
return self.units.boltzmann
[docs] def set_time_step(self, dt: float):
r"""
Set simulation time step.
Args:
dt (float): Time step.
"""
self.time_step = get_ms_array(dt, ms.float32)
return self
[docs] def set_degrees_of_freedom(self, dofs: int):
"""
Set degrees of freedom (DOFs).
Args:
dofs (int): Degrees of freedom.
"""
self.degrees_of_freedom = get_integer(dofs)
return self
[docs] def update_coordinate(self, coordinate: Tensor) -> Tensor:
r"""
Update the coordinate of the simulation system.
Args:
coordinate (Tensor): Tensor of atomic coordinates. Tensor shape is :math:`(B, A, D)`.
Data type is float.
Returns:
Tensor, has the same data type and shape as original `coordinate`.
"""
return F.assign(self._coordinate, coordinate)
[docs] def update_pbc_box(self, pbc_box: Tensor) -> Tensor:
r"""
Update the parameter of PBC box.
Args:
pbc_box (Tensor): Tensor of PBC box. Tensor shape is :math:`(B, D)`.
Data type is float.
Returns:
Tensor, has the same data type and shape as original `pbc_box`.
"""
if self._pbc_box is None:
return pbc_box
return F.assign(self._pbc_box, pbc_box)
[docs] def get_kinetics(self, velocity: Tensor) -> Tensor:
r"""
Calculate kinetics according to velocity.
Args:
velocity (Tensor): Tensor of atomic velocities. Tensor shape is :math:`(B, A, D)`.
Data type is float.
Returns:
Tensor, kinetics. Tensor shape is :math:`(B, D)`. Data type is float.
"""
if velocity is None:
return None
# (B, A, D) * (B, A, 1)
k = 0.5 * self._atom_mass * velocity**2
# (B, D) <- (B, A, D)
kinetics = F.reduce_sum(k, -2)
return kinetics * self.kinetic_unit_scale
[docs] def get_temperature(self, kinetics: Tensor = None) -> Tensor:
r"""
Calculate temperature according to velocity.
Args:
kinetics (Tensor): Kinetics. Tensor shape is :math:`(B, D)`.
Data type is float. Default: ``None``.
Returns:
Tensor, temperature. Tensor shape is :math:`(B)`. Data type is float.
"""
if kinetics is None:
return None
# (B) <- (B, D)
kinetics = F.reduce_sum(kinetics, -1)
return 2 * kinetics / self.degrees_of_freedom / self.boltzmann
[docs] def get_volume(self, pbc_box: Tensor) -> Tensor:
r"""
Calculate volume according to PBC box.
Args:
pbc_box (Tensor): Tensor of PBC box. Tensor shape is :math:`(B, D)`. Data type is float.
Returns:
Tensor, Tensor of volume. Shape is :math:`(B)`. The data type is float.
"""
if self._pbc_box is None:
return None
# (B, 1) <- (B, D)
return func.keepdims_prod(pbc_box, -1)
[docs] def get_pressure(self, kinetics: Tensor, virial: Tensor, pbc_box: Tensor) -> Tensor:
r"""
Calculate pressure according to kinetics, viral and PBC box.
Args:
kinetics (Tensor): Kinetics. Tensor shape is :math:`(B, D)`. Data type is float.
virial (Tensor): Virial. Tensor shape is :math:`(B, D)`. Data type is float.
pbc_box (Tensor): PBC box. Tensor shape is :math:`(B, D)`. Data type is float.
Returns:
Tensor, pressure. Tensor shape is :math:`(B, D)`. Data type is float.
"""
if self._pbc_box is None:
return None
volume = func.keepdims_prod(pbc_box, -1)
# (B, D) = ((B, D) - (B, D)) / (B, 1)
pressure = 2 * (kinetics - virial) / volume
return pressure * self.press_unit_scale
[docs] def get_com(self, coordinate: Tensor, keepdims: bool = True) -> Tensor:
r"""
Get coordinate of center of mass.
Args:
coordinate (Tensor): Atomic coordinates. Tensor shape is :math:`(B, A, D)`.
Data type is float.
keepdims (bool, optional): If this is set to ``True``, the second axis will be left
in the result as dimensions with size one.
Default: ``True``.
Returns:
Tensor, the coordinate of the center of mass.
The shape is :math:`(B, A, D)` or :math:`(B, D)`.
Data type is float.
"""
# (B, A, D) = (B, A, D) * (B, A, 1)
weight_coord = coordinate * self._atom_mass
if keepdims:
# (B, 1, D) <- (B, A, D)
tot_coord = self.keepdims_sum(weight_coord, -2)
# (B, 1, 1) <- (B, 1)
tot_mass = F.expand_dims(self.system_mass, -1)
else:
# (B, D) <- (B, A, D)
tot_coord = F.reduce_sum(weight_coord, -2)
# (B, 1)
tot_mass = self.system_mass
# (B, 1, D) = (B, 1, D) / (B, 1, 1)
# OR
# (B, D) = (B, D) / (B, 1)
com = tot_coord / tot_mass
return com
[docs] def get_com_velocity(self, velocity: Tensor, keepdims: bool = True) -> Tensor:
r"""
Calculate velocity of center of mass.
Args:
velocity (Tensor): Velocity. The shape is :math:`(B, A, D)`. Data type is float.
keepdims(bool): If this is set to ``True``, the second axis will be left
in the result as dimensions with size one. Default: ``True``.
Returns:
Tensor, Tensor of the velocity of the center of mass.
The shape is :math:`(B, A, D)` or :math:`(B, D)`. Data type is float.
"""
# (B, A, D) = (B, A, D) * (B, A, 1)
weight_vel = velocity * self._atom_mass
if keepdims:
# (B, 1, D) <- (B, A, D)
tot_vel = self.keepdims_sum(weight_vel, -2)
# (B, 1, 1) <- (B, 1)
tot_mass = F.expand_dims(self.system_mass, -1)
else:
# (B, D) <- (B, A, D)
tot_vel = F.reduce_sum(weight_vel, -2)
# (B, 1)
tot_mass = self.system_mass
# (B, 1, D) = (B, 1, D) / (B, 1, 1)
# OR
# (B, D) = (B, D) / (B, 1)
com_vel = tot_vel / tot_mass
return com_vel
def construct(self,
coordinate: Tensor,
velocity: Tensor,
force: Tensor,
energy: Tensor,
kinetics: Tensor,
virial: Tensor = None,
pbc_box: Tensor = None,
step: int = 0,
) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
r"""
Control the parameters during the simulation
Args:
coordinate (Tensor): Tensor of shape :math:`(B, A, D)`. Data type is float.
velocity (Tensor): Tensor of shape :math:`(B, A, D)`. Data type is float.
force (Tensor): Tensor of shape :math:`(B, A, D)`. Data type is float.
energy (Tensor): Tensor of shape :math:`(B, 1)`. Data type is float.
kinetics (Tensor): Tensor of shape :math:`(B, D)`. Data type is float.
virial (Tensor): Tensor of shape :math:`(B, D)`. Data type is float.
pbc_box (Tensor): Tensor of shape :math:`(B, D)`. Data type is float.
step (int): Simulation step. Default: ``0``.
Returns:
- **coordinate** (Tensor) - Tensor of shape :math:`(B, A, D)`. Data type is float.
- **velocity** (Tensor) - Tensor of shape :math:`(B, A, D)`. Data type is float.
- **force** (Tensor) - Tensor of shape :math:`(B, A, D)`. Data type is float.
- **energy** (Tensor) - Tensor of shape :math:`(B, 1)`. Data type is float.
- **kinetics** (Tensor) - Tensor of shape :math:`(B, D)`. Data type is float.
- **virial** (Tensor) - Tensor of shape :math:`(B, D)`. Data type is float.
- **pbc_box** (Tensor) - Tensor of shape :math:`(B, D)`. Data type is float.
Note:
:math:`B` is the number of walkers in simulation.
:math:`A` is the number of atoms.
:math:`D` is the spatial dimension of the simulation system. Usually is 3.
"""
#pylint: disable=unused-argument
return coordinate, velocity, force, energy, kinetics, virial, pbc_box
def _get_mw_tensor(self,
value: Union[float, ndarray, Tensor, List[float]],
name: str) -> Tensor:
"""get tensor for multiple walkers"""
value = func.get_tensor(value, ms.float32)
if value.size == 1:
# ()
return F.reshape(value, ())
if value.size != self.num_walker:
error_info = f'The size of {name} must be equal to 1 or ' \
f'the number of multiple walker ({self.num_walker}) but got '
_raise_value_error(error_info, value.size)
# (B, 1)
return F.reshape(value, (self.num_walker, 1))