sponge.optimizer.Updater

class sponge.optimizer.Updater(system: Molecule, controller: Union[Controller, List[Controller]] = None, time_step: float = 0.001, velocity: Union[Tensor, ndarray, List[float]] = None, weight_decay: float = 0.0, loss_scale: float = 1.0, **kwargs)[source]

Base class of the MindSPONGE updater, which is a special subclass of the mindspore.nn.Optimizer in MindSpore.

The sponge.optimizer.Updater updates the atomic coordinates of the simulation system. The updating of atomic coordinates requires atomic forces and atomic velocities, where the force is passed from outside and the velocity is the parameter of the sponge.optimizer.Updater itself. And in the case of periodic boundary conditions (PBC), the sponge.optimizer.Updater could also update the size of the PBC box by the virial of the simulation system.

The sponge.optimizer.Updater controls the values of seven variables during the simulation through a series of sponge.control.Controller : coordinates, velocity, force, energy, kinetics, virial and pbc_box. If more than one sponge.control.Controller is passed in, they will work in sequence.

Parameters

system (sponge.system.Molecule) – Simulation system.
controller (Union[sponge.control.Controller, List[sponge.control.Controller]], optional) – Controller or list of controllers to control the seven variables (coordinate, velocity, force, energy, kinetics, virial and pbc_box) of the simulation system. Default: None.
time_step (float, optional) – Time step. Default: 1e-3.
velocity (Union[Tensor, ndarray, List[float]], optional) – Array of atomic velocity. The shape of array is \((A, D)\) or \((B, A, D)\). Here \(A\) is the number of atoms, \(D\) is the spatial dimension of the simulation system, which is usually 3. The data type is float. Default: None.
weight_decay (float, optional) – An value for the weight decay. Default: 0.0.
loss_scale (float, optional) – A value for the loss scale. Default: 1.0.
kwargs (dict) – Other arguments.

Inputs:

energy (Tensor) - Energy of the system. Tensor of shape \((B, A, D)\). Data type is float.
force (Tensor) - Force of the system. Tensor of shape \((B, A, D)\). Data type is float.
virial (Tensor) - Virial of the system. Tensor of shape (B, A, D). Data type is float. Default: None.

Outputs:

success (bool) - whether successfully finish the current optimization step and move to next step.

Supported Platforms:

Ascend GPU

Examples

>>> from sponge.system import Molecule
>>> from sponge.optimizer import Updater
>>> from sponge.potential.forcefield import ForceField
>>> from sponge.core.sponge import Sponge
>>> system = Molecule(template='water.tip3p.yaml')
>>> potential = ForceField(system, parameters='SPCE')
>>> updater = Updater(system, controller=None, time_step=1e-3)
>>> print(system.coordinate.value())
>>> # [[[ 0. 0. 0. ]
>>> # [ 0.07907964 0.06120793 0. ]
>>> # [-0.07907964 0.06120793 0. ]]]
>>> md = Sponge(system, potential, updater)
>>> # As controller is None, no change on MD variables happens
>>> md.run(1000)
>>> # [MindSPONGE] Started simulation at 2024-04-29 01:01:34
>>> # [MindSPONGE] Finished simulation at 2024-04-29 01:01:37
>>> # [MindSPONGE] Simulation time: 2.54 seconds.
>>> print(system.coordinate.value())
>>> # [[[ 0. 0. 0. ]
>>> # [ 0.07907964 0.06120793 0. ]
>>> # [-0.07907964 0.06120793 0. ]]]

property boltzmann: float

Boltzmann constant in current unit.

Returns: float, Boltzmann constant in current unit.

decay_and_scale_grad(force: Tensor, virial: Tensor = None)[source]

Do weight decay and gradient scale for force and virial.

Parameters

force (Tensor) – Tensor of force. Data type is float.
virial (Tensor, optional) – Tensor of virial. Data type is float. Default: None.

Returns

Tensor, Tensor of force after weight decay and gradient scale.
Tensor, Tensor of virial after weight decay and gradient scale. If pbc_box is None, the output virial is the same as input.

get_dt()[source]

Get the learning rate of current step.

Returns: float, the learning rate of current step.

get_kinetics(velocity: Tensor)[source]

Get kinectics.

Parameters: velocity (Tensor) – Tensor of atom velocities. Data type is float.
Returns: Tensor, the kinectics of the system.

get_pressure(kinetics: Tensor, virial: Tensor, pbc_box: Tensor)[source]

Get pressure.

Parameters

kinetics (Tensor) – Tensor of kinetics. Data type is float.
virial (Tensor) – Tensor of virial. Data type is float.
pbc_box (Tensor) – Tensor of pbc_box. Data type is float.

Returns

Tensor, the pressure of the system.

get_temperature(kinetics: Tensor = None)[source]

Get temperature.

Parameters: kinetics (Tensor) – Tensor of kinetics. Data type is float. Default: None.
Returns: Tensor, the temperature of the system.

get_velocity()[source]

Get velocity.

Returns: Tensor, atom velocities of the system.

next_step(success: bool = True)[source]

Finish the current optimization step and move to next step.

Parameters: success (bool) – Whether to finish the current optimization step and move to next step. Default: True.
Returns: bool, whether successfully finish the current optimization step and move to next step.

property press_unit_scale: float

Reference value of pressure.

Returns: float, reference value of pressure.

set_degrees_of_freedom(dofs: int)[source]

Set degrees of freedom (DOFs)

Parameters: dofs (int) – Degrees of freedom.

set_step(step: int = 0)[source]

Set current step of the system.

Parameters: step (int) – Current step of the system. Default: 0

update_coordinate(coordinate: Tensor, success: bool = True)[source]

Update the parameters of coordinate

Parameters

coordinate (Tensor) – Tensor of atomic coordinates. Data type is float.
success (bool) – Whether to update the coordinate. Default: True.

Returns

bool, whether successfully update the coordinate.

update_kinetics(kinetics: Tensor, success: bool = True)[source]

Update the parameters of kinetics.

Parameters

kinetics (Tensor) – Tensor of kinetics. Data type is float.
success (bool) – Whether to update the kinetics. Default: True.

Returns

bool, whether successfully update the parameters of kinetics.

update_pbc_box(pbc_box: Tensor, success: bool = True)[source]

Update the parameters of PBC box.

Parameters

pbc_box (Tensor) – Tensor of PBC box. Data type is float.
success (bool) – Whether to update the pbc_box. Default: True.

Returns

bool, whether successfully update the parameters of PBC box.

update_pressure(pressure: Tensor, success: bool = True)[source]

Update the parameters of pressure.

Parameters

pressure (Tensor) – Tensor of pressure. Data type is float.
success (bool) – Whether to update the pressure. Default: True.

Returns

bool, whether successfully update the parameters of pressure.

update_temperature(temperature: Tensor, success: bool = True)[source]

Update the parameters of temperature.

Parameters

temperature (Tensor) – Tensor of temperature. Data type is float.
success (bool) – Whether to update the temperature. Default: True.

Returns

bool, whether successfully update the parameters of temperature.

update_velocity(velocity: Tensor, success: bool = True)[source]

Update the parameters of velocity.

Parameters

velocity (Tensor) – Tensor of atomic velocities. Data type is float.
success (bool) – Whether to update the velocities. Default: True.

Returns

bool, whether successfully update the parameters of atomic velocities.

update_virial(virial: Tensor, success: bool = True)[source]

Update the parameters of virial.

Parameters

virial (Tensor) – Tensor of virial. Data type is float.
success (bool) – Whether to update the virial. Default: True.

Returns

bool, whether successfully update the parameters of virial.