tfep.potentials.openmm.OpenMMPotential

class tfep.potentials.openmm.OpenMMPotential(system, platform, positions_unit: Unit | None = None, energy_unit: Unit | None = None, system_name: str | None = None, precompute_gradient: bool = False, parallelization_strategy: ParallelizationStrategy | None = None)[source]

Bases: PotentialBase

Potential energy and forces with OpenMM.

Wraps OpenMM to provides a differentiable potential energy function for training.

Warning

Currently double-backpropagation is not supported, which means force matching cannot be performed during training.

__init__(system, platform, positions_unit: Unit | None = None, energy_unit: Unit | None = None, system_name: str | None = None, precompute_gradient: bool = False, parallelization_strategy: ParallelizationStrategy | None = None)[source]

Constructor.

Parameters:

system (openmm.System) – The OpenMM System used to compute energies and forces.
platform (str or openmm.Platform, optional) – The OpenMM Platform object or its name. If None, the default platform is used.
positions_unit (pint.Unit, optional) – The unit of the positions passed to the class methods. Since input Tensor``s do not have units attached, this is used to appropriately convert ``batch_positions to OpenMM units. If None, no conversion is performed, which assumes that the input positions are in the same units used internally by OpenMM (nanometer).
energy_unit (pint.Unit, optional) – The unit used for the returned energies (and as a consequence forces). Since Tensor``s do not have units attached, this is used to appropriately convert OpenMM energies into the desired units. If ``None no conversion is performed, which means that energies and forces will be returned in OpenMM units (kJ/mol).
system_name (str, optional) – If given, the OpenMM Context will be cached in the global variable tfep.potentials.openmm.global_context_cache with system_name as a key so that the Context object will not be re-initialized at every energy evaluation.

Note that it is the user’s responsibility to clear the cache once the Context is not needed anymore.
precompute_gradient (bool, optional) – If True, the gradient is computed in the forward pass and saved to be consumed during the backward pass. This speeds up the training, but should be deactivated if gradients are not needed. Setting this to False (default) will cause an exception if a backward pass is attempted.
parallelization_strategy (tfep.utils.parallel.ParallelizationStrategy, optional) – The parallelization strategy used to distribute batches of energy and gradient calculations. By default, these are executed serially.

See also

OpenMMPotentialEnergyFunc: More details on input parameters and implementation details.

Methods

`__init__`(system, platform[, positions_unit, ...])	Constructor.
`add_module`(name, module)	Add a child module to the current module.
`apply`(fn)	Apply `fn` recursively to every submodule (as returned by `.children()`) as well as self.
`bfloat16`()	Casts all floating point parameters and buffers to `bfloat16` datatype.
`buffers`([recurse])	Return an iterator over module buffers.
`children`()	Return an iterator over immediate children modules.
`compile`(args, *kwargs)	Compile this Module's forward using `torch.compile()`.
`cpu`()	Move all model parameters and buffers to the CPU.
`cuda`([device])	Move all model parameters and buffers to the GPU.
`default_energy_unit`(unit_registry)	Return the default energy units.
`default_positions_unit`(unit_registry)	Return the default positions units.
`double`()	Casts all floating point parameters and buffers to `double` datatype.
`eval`()	Set the module in evaluation mode.
`extra_repr`()	Set the extra representation of the module.
`float`()	Casts all floating point parameters and buffers to `float` datatype.
`forward`(batch_positions[, batch_cell])	Compute a differential potential energy for a batch of configurations.
`get_buffer`(target)	Return the buffer given by `target` if it exists, otherwise throw an error.
`get_extra_state`()	Return any extra state to include in the module's state_dict.
`get_parameter`(target)	Return the parameter given by `target` if it exists, otherwise throw an error.
`get_submodule`(target)	Return the submodule given by `target` if it exists, otherwise throw an error.
`half`()	Casts all floating point parameters and buffers to `half` datatype.
`ipu`([device])	Move all model parameters and buffers to the IPU.
`load_state_dict`(state_dict[, strict, assign])	Copy parameters and buffers from `state_dict` into this module and its descendants.
`modules`()	Return an iterator over all modules in the network.
`mtia`([device])	Move all model parameters and buffers to the MTIA.
`named_buffers`([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
`named_children`()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
`named_modules`([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
`named_parameters`([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
`parameters`([recurse])	Return an iterator over module parameters.
`register_backward_hook`(hook)	Register a backward hook on the module.
`register_buffer`(name, tensor[, persistent])	Add a buffer to the module.
`register_forward_hook`(hook, *[, prepend, ...])	Register a forward hook on the module.
`register_forward_pre_hook`(hook, *[, ...])	Register a forward pre-hook on the module.
`register_full_backward_hook`(hook[, prepend])	Register a backward hook on the module.
`register_full_backward_pre_hook`(hook[, prepend])	Register a backward pre-hook on the module.
`register_load_state_dict_post_hook`(hook)	Register a post-hook to be run after module's `load_state_dict()` is called.
`register_load_state_dict_pre_hook`(hook)	Register a pre-hook to be run before module's `load_state_dict()` is called.
`register_module`(name, module)	Alias for `add_module()`.
`register_parameter`(name, param)	Add a parameter to the module.
`register_state_dict_post_hook`(hook)	Register a post-hook for the `state_dict()` method.
`register_state_dict_pre_hook`(hook)	Register a pre-hook for the `state_dict()` method.
`requires_grad_`([requires_grad])	Change if autograd should record operations on parameters in this module.
`set_extra_state`(state)	Set extra state contained in the loaded state_dict.
`set_submodule`(target, module)	Set the submodule given by `target` if it exists, otherwise throw an error.
`share_memory`()	See `torch.Tensor.share_memory_()`.
`state_dict`(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
`to`(args, *kwargs)	Move and/or cast the parameters and buffers.
`to_empty`(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
`train`([mode])	Set the module in training mode.
`type`(dst_type)	Casts all parameters and buffers to `dst_type`.
`xpu`([device])	Move all model parameters and buffers to the XPU.
`zero_grad`([set_to_none])	Reset gradients of all model parameters.

Attributes

`DEFAULT_ENERGY_UNIT`	The default energy unit.
`DEFAULT_POSITIONS_UNIT`	The default positions unit.
`T_destination`
`call_super_init`
`dump_patches`
`energy_unit`	The energy units of the returned potential.
`positions_unit`	The positions unit requested for the input.
`system`	The OpenMM System used to compute energies and forces.
`platform`	The OpenMM Platform.
`system_name`	The identifier of the system used as key for the context cache.
`precompute_gradient`	Whether to compute the gradients in the forward pass to speed up the backward pass.
`parallelization_strategy`	The strategy used to parallelize the single-point calculations.
`training`

DEFAULT_ENERGY_UNIT: str = 'kJ/mol': The default energy unit.

DEFAULT_POSITIONS_UNIT: str = 'nanometer': The default positions unit.

classmethod default_energy_unit(unit_registry) → Quantity: Return the default energy units.

classmethod default_positions_unit(unit_registry) → Quantity: Return the default positions units.

property energy_unit: Quantity: The energy units of the returned potential.

forward(batch_positions: Tensor, batch_cell: Tensor | None = None) → Tensor[source]

Compute a differential potential energy for a batch of configurations.

Parameters:

batch_positions (torch.Tensor) – Shape (batch_size, 3*n_atoms). The atoms positions in units of self.positions_unit.
batch_cell (torch.Tensor, optional) – Shape (batch_size, 6). Unitcell dimensions. For each data point, the first 3 elements represent the vector lengths in units of self.positions_unit and the last 3 their respective angles (in degrees) in the following order: [len(a), len(b), len(c), angle(b,c), angle(a,c), angle(a,b)]. The first vector will lie in x-direction, second in xy-plane, and the third one in the z-positive subspace.

Returns:

potential_energy – potential_energy[i] is the potential energy of configuration batch_positions[i] in units of self.energy_unit.

Return type:

torch.Tensor

parallelization_strategy: The strategy used to parallelize the single-point calculations.

platform: The OpenMM Platform.

property positions_unit: Quantity: The positions unit requested for the input.

precompute_gradient: Whether to compute the gradients in the forward pass to speed up the backward pass.

system: The OpenMM System used to compute energies and forces.

system_name: The identifier of the system used as key for the context cache.