tfep.potentials.mimic.MiMiCPotentialEnergyFunc

class tfep.potentials.mimic.MiMiCPotentialEnergyFunc(*args, **kwargs)[source]

Bases: Function

PyTorch-differentiable QM/MM potential energy function wrapped around MiMiC.

The function calls MiMiC through the command line interface using user-prepared CPMD and GROMACS input files and reads the resulting energies and forces. The function can use the input files as templates and evaluate energies and forces for batches of configurations. To do this, the input files must be setup to generate the ENERGIES and FTRAJECTORY trajectory files, which can be generated when both &CPMD.MOLECULAR DYNAMICS and &CPMD.TRAJECTORY FORCES options are used in the CPMD input script. The returned energy and force will be read from the first entry in the respective trajectory file.

Note

With this mechanism, the returned energy and force will be associated to the configuration AFTER the first integration step, not to the configuration passed in batch_positions. However, this is necessary as currently there is no way in MiMiC to save the forces of all the atoms (including the MM atoms) with a single wavefunction optimization calculation.

To maximize the efficiency and reduce the error, we suggest setting &CPMD.TRAJECTORY FORCES and &CPMD.MAXSTEP to 1 and set a very small &CPMD.TIMESTEP (e.g., 0.000001 a.u.).

Note

The order of the atoms typically differs in GROMACS and CPMD. The input and output positions and forces of this function always use the GROMACS atom order.

If backpropagation is not necessary, &CPMD.TRAJECTORY FORCES can be left out but the option precompute_gradients must be set to False or the function will expect an output FTRAJECTORY file.

Because GROMACS does not support resuming from a coordinate file (it requires a full checkpoint file), to perform the calculation starting from an arbitrary batch data point with different coordinates than those in the tpr file, we need to regenerate a temporary tpr file with the correct coordinates with GROMPP.

Only orthorhombic boxes are currently supported by MiMiC. thus the box vectors are simply specified by the lengths of the box sides.

The function supports running each MiMiC execution in a separate working directory to safely support batch parallelization schemes through :class:tfep.utils.parallel.ParallelizationStrategy objects.

Sometimes the communication between GROMACS and CPMD can fail causing a crash. In this case, the MiMiC execution is attempted n_attempts times before raising a RuntimeError.

It is possible to handle in different ways the cases when the calculation does not converge the wavefunction within the number of SCF steps specified or when MiMiC terminates with an error, depending on whether the user wants to halt the program or continue with a NaN potential energy value.

Parameters:

ctx (torch.autograd.function._ContextMethodMixin) – A context to save information for the gradient.
batch_positions (torch.Tensor) – A tensor of positions in flattened format (i.e., with shape (batch_size, 3*n_atoms)).

Note that the order of the atoms is assumed to be that of the GROMACS input files, not the one used internally by CPMD.
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).
cpmd_cmd (tfep.potentials.mimic.Cpmd) – The CPMD command to be run for MiMiC’s execution that encapsulates the path to the CPMD input script and options.

The &MIMIC.PATHS option and atomic coordinates can be placeholders as they are automatically set by this function according to the working_dir_path and batch_positions arguments. All other options must be set correctly for the function to run successfully.
mdrun_cmd (tfep.potentials.mimic.GmxMdrun) – The GMX mdrun command to be run for MiMiC’s execution that encapsulates the path to the GROMACS input script and running options.

The mdrun_cmd.tpr_input_file_path can be left unset since a new .tpr file with the correct positions is automatically generated with gromp_cmd.
gromacs_to_cpmd_atom_indices (Dict[int, int]) – A dictionary associating atom indices in GROMACS to atom indices in CPMD.
grompp_cmd (tfep.potentials.mimic.GmxGrompp, optional) – This command is used to generate the the .tpr file with the correct coordinates. To do so, the batch positions are first stored in a .trr file which is then passed to grompp. Thus, the GmxGrompp.tpr_output_file_path and GmxGrompp.trajectory_input_file_path options can be None.
launcher (tfep.utils.cli.Launcher, optional) – The Launcher to use to run the cpmd_cmd and mdrun_cmd. If not passed, a new tfep.utils.cli.Launcher is created.
grompp_launcher (tfep.utils.cli.Launcher, optional) – The Launcher to use to run the grompp_cmd command. If not passed, a new tfep.utils.cli.Launcher is created.
positions_unit (pint.Unit, optional) – The unit of the positions passed. This is used to appropriately convert batch_positions to the units used by MiMiC. If None, no conversion is performed, which assumes that the input positions are in Bohr.
energy_unit (pint.Unit, optional) – The unit used for the returned energies (and as a consequence forces). This is used to appropriately convert MiMiC energies into the desired units. If None, no conversion is performed, which means that energies and forces will be in hartrees and hartrees/bohr respectively.
precompute_gradient (bool, optional) – If False, the FTRAJECTORY file is not read after executing MiMiC. This might save a small amount of time if backpropagation is not needed.
working_dir_path (str or List[str], optional) – The working directory to be used to run MiMiC and grompp. This must exist. If a list, batch_positions[i] is evaluated in the directory working_dir_path[i].
cleanup_working_dir (bool, optional) – If True and working_dir_path is passed, all the files inside the working directory are removed after executing MiMiC. The directory(s) itself is not deleted.
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.
launcher_kwargs (Dict, optional) – Other kwargs for launcher (with the exception of cwd which is automatically determined based on working_dir_path).
grompp_launcher_kwargs (Dict, optional) – Other kwargs for grompp_launcher.
n_attempts (int, optional) – Number of times MiMiC is restarted before raising a RuntimeError when MiMiC crashes without creating an error report in the LocalError-X-X-X.log file.
on_unconverged (str, optional) – Specifies how to handle the case in which the self-consistent calculation did not converge. It can have the following values: - 'raise': Raises a RuntimeError and halts the execution. - 'success': Treat the calculation as converged and return the latest

energy and force values.
- 'nan': Return float('nan') energy and zero forces.
If this is set to anything other than 'success', the stdout keyword argument must be included in launcher_kwargs and set to subprocess.PIPE so that Python can intercept and parse the output to detect the convergence warning message.
on_local_error (str, optional) – Specifies how to handle the case in which the calculation ends with an error and CPMD creates an error report in the LocalError-X-X-X.log file. It can have the following values: - 'raise': Raises a RuntimeError and halts the execution. - 'nan': Return float('nan') energy and zero forces.

Returns:

potentials – potentials[i] is the potential energy of configuration batch_positions[i].

Return type:

torch.Tensor

See also

MiMiCPotential: Module API for computing potential energies with MiMiC.

__init__(*args, **kwargs)

Methods

`__init__`(args, *kwargs)
`apply`(args, *kwargs)
`backward`(ctx, grad_output)	Compute the gradient of the potential energy.
`forward`(ctx, batch_positions, batch_cell, ...)	Compute the potential energy of the molecule with MiMiC.
`jvp`(ctx, *grad_inputs)	Define a formula for differentiating the operation with forward mode automatic differentiation.
`mark_dirty`(*args)	Mark given tensors as modified in an in-place operation.
`mark_non_differentiable`(*args)	Mark outputs as non-differentiable.
`mark_shared_storage`(*pairs)
`maybe_clear_saved_tensors`
`name`
`register_hook`
`register_prehook`
`save_for_backward`(*tensors)	Save given tensors for a future call to `backward()`.
`save_for_forward`(*tensors)	Save given tensors for a future call to `jvp()`.
`set_materialize_grads`(value)	Set whether to materialize grad tensors.
`setup_context`(ctx, inputs, output)	There are two ways to define the forward pass of an autograd.Function.
`vjp`(ctx, *grad_outputs)	Define a formula for differentiating the operation with backward mode automatic differentiation.
`vmap`(info, in_dims, *args)	Define the behavior for this autograd.Function underneath `torch.vmap()`.

Attributes

`dirty_tensors`
`generate_vmap_rule`
`materialize_grads`
`metadata`
`needs_input_grad`
`next_functions`
`non_differentiable`
`requires_grad`
`saved_for_forward`
`saved_tensors`
`saved_variables`
`to_save`

static backward(ctx, grad_output)[source]: Compute the gradient of the potential energy.

static forward(ctx: FunctionCtx, batch_positions: Tensor, batch_cell: Tensor, cpmd_cmd: Cpmd, mdrun_cmd: GmxMdrun, grompp_cmd: GmxGrompp, gromacs_to_cpmd_atom_indices: Dict[int, int], launcher: Launcher | None = None, grompp_launcher: Launcher | None = None, positions_unit: Unit | None = None, energy_unit: Unit | None = None, precompute_gradient: bool = True, working_dir_path: str | List[str] | None = None, cleanup_working_dir: bool = False, parallelization_strategy: ParallelizationStrategy | None = None, launcher_kwargs: Dict[str, Any] | None = None, grompp_launcher_kwargs: Dict[str, Any] | None = None, n_attempts: int = 1, on_unconverged: str = 'raise', on_local_error: str = 'raise')[source]: Compute the potential energy of the molecule with MiMiC.

static jvp(ctx: Any, *grad_inputs: Any) → Any

Define a formula for differentiating the operation with forward mode automatic differentiation.

This function is to be overridden by all subclasses. It must accept a context ctx as the first argument, followed by as many inputs as the forward() got (None will be passed in for non tensor inputs of the forward function), and it should return as many tensors as there were outputs to forward(). Each argument is the gradient w.r.t the given input, and each returned value should be the gradient w.r.t. the corresponding output. If an output is not a Tensor or the function is not differentiable with respect to that output, you can just pass None as a gradient for that input.

You can use the ctx object to pass any value from the forward to this functions.

mark_dirty(*args: Tensor)

Mark given tensors as modified in an in-place operation.

This should be called at most once, in either the setup_context() or forward() methods, and all arguments should be inputs.

Every tensor that’s been modified in-place in a call to forward() should be given to this function, to ensure correctness of our checks. It doesn’t matter whether the function is called before or after modification.

Examples::

>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
>>> class Inplace(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         x_npy = x.numpy() # x_npy shares storage with x
>>>         x_npy += 1
>>>         ctx.mark_dirty(x)
>>>         return x
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, grad_output):
>>>         return grad_output
>>>
>>> a = torch.tensor(1., requires_grad=True, dtype=torch.double).clone()
>>> b = a * a
>>> Inplace.apply(a)  # This would lead to wrong gradients!
>>>                   # but the engine would not know unless we mark_dirty
>>> # xdoctest: +SKIP
>>> b.backward() # RuntimeError: one of the variables needed for gradient
>>>              # computation has been modified by an inplace operation

mark_non_differentiable(*args: Tensor)

Mark outputs as non-differentiable.

This should be called at most once, in either the setup_context() or forward() methods, and all arguments should be tensor outputs.

This will mark outputs as not requiring gradients, increasing the efficiency of backward computation. You still need to accept a gradient for each output in backward(), but it’s always going to be a zero tensor with the same shape as the shape of a corresponding output.

This is used e.g. for indices returned from a sort. See example::

>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         sorted, idx = x.sort()
>>>         ctx.mark_non_differentiable(idx)
>>>         ctx.save_for_backward(x, idx)
>>>         return sorted, idx
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):  # still need to accept g2
>>>         x, idx = ctx.saved_tensors
>>>         grad_input = torch.zeros_like(x)
>>>         grad_input.index_add_(0, idx, g1)
>>>         return grad_input

save_for_backward(*tensors: Tensor)

Save given tensors for a future call to backward().

save_for_backward should be called at most once, in either the setup_context() or forward() methods, and only with tensors.

All tensors intended to be used in the backward pass should be saved with save_for_backward (as opposed to directly on ctx) to prevent incorrect gradients and memory leaks, and enable the application of saved tensor hooks. See torch.autograd.graph.saved_tensors_hooks.

Note that if intermediary tensors, tensors that are neither inputs nor outputs of forward(), are saved for backward, your custom Function may not support double backward. Custom Functions that do not support double backward should decorate their backward() method with @once_differentiable so that performing double backward raises an error. If you’d like to support double backward, you can either recompute intermediaries based on the inputs during backward or return the intermediaries as the outputs of the custom Function. See the double backward tutorial for more details.

In backward(), saved tensors can be accessed through the saved_tensors attribute. Before returning them to the user, a check is made to ensure they weren’t used in any in-place operation that modified their content.

Arguments can also be None. This is a no-op.

See extending-autograd for more details on how to use this method.

Example::

>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x: torch.Tensor, y: torch.Tensor, z: int):
>>>         w = x * z
>>>         out = x * y + y * z + w * y
>>>         ctx.save_for_backward(x, y, w, out)
>>>         ctx.z = z  # z is not a tensor
>>>         return out
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, grad_out):
>>>         x, y, w, out = ctx.saved_tensors
>>>         z = ctx.z
>>>         gx = grad_out * (y + y * z)
>>>         gy = grad_out * (x + z + w)
>>>         gz = None
>>>         return gx, gy, gz
>>>
>>> a = torch.tensor(1., requires_grad=True, dtype=torch.double)
>>> b = torch.tensor(2., requires_grad=True, dtype=torch.double)
>>> c = 4
>>> d = Func.apply(a, b, c)

save_for_forward(*tensors: Tensor)

Save given tensors for a future call to jvp().

save_for_forward should be called at most once, in either the setup_context() or forward() methods, and all arguments should be tensors.

In jvp(), saved objects can be accessed through the saved_tensors attribute.

Arguments can also be None. This is a no-op.

See extending-autograd for more details on how to use this method.

Example::

>>> # xdoctest: +SKIP
>>> class Func(torch.autograd.Function):
>>>     @staticmethod
>>>     def forward(ctx, x: torch.Tensor, y: torch.Tensor, z: int):
>>>         ctx.save_for_backward(x, y)
>>>         ctx.save_for_forward(x, y)
>>>         ctx.z = z
>>>         return x * y * z
>>>
>>>     @staticmethod
>>>     def jvp(ctx, x_t, y_t, _):
>>>         x, y = ctx.saved_tensors
>>>         z = ctx.z
>>>         return z * (y * x_t + x * y_t)
>>>
>>>     @staticmethod
>>>     def vjp(ctx, grad_out):
>>>         x, y = ctx.saved_tensors
>>>         z = ctx.z
>>>         return z * grad_out * y, z * grad_out * x, None
>>>
>>>     a = torch.tensor(1., requires_grad=True, dtype=torch.double)
>>>     t = torch.tensor(1., dtype=torch.double)
>>>     b = torch.tensor(2., requires_grad=True, dtype=torch.double)
>>>     c = 4
>>>
>>>     with fwAD.dual_level():
>>>         a_dual = fwAD.make_dual(a, t)
>>>         d = Func.apply(a_dual, b, c)

set_materialize_grads(value: bool)

Set whether to materialize grad tensors. Default is True.

This should be called only from either the setup_context() or forward() methods.

If True, undefined grad tensors will be expanded to tensors full of zeros prior to calling the backward() and jvp() methods.

Example::

>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
>>> class SimpleFunc(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         return x.clone(), x.clone()
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):
>>>         return g1 + g2  # No check for None necessary
>>>
>>> # We modify SimpleFunc to handle non-materialized grad outputs
>>> class Func(Function):
>>>     @staticmethod
>>>     def forward(ctx, x):
>>>         ctx.set_materialize_grads(False)
>>>         ctx.save_for_backward(x)
>>>         return x.clone(), x.clone()
>>>
>>>     @staticmethod
>>>     @once_differentiable
>>>     def backward(ctx, g1, g2):
>>>         x, = ctx.saved_tensors
>>>         grad_input = torch.zeros_like(x)
>>>         if g1 is not None:  # We must check for None now
>>>             grad_input += g1
>>>         if g2 is not None:
>>>             grad_input += g2
>>>         return grad_input
>>>
>>> a = torch.tensor(1., requires_grad=True)
>>> b, _ = Func.apply(a)  # induces g2 to be undefined

static setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) → Any

There are two ways to define the forward pass of an autograd.Function.

Either:

Override forward with the signature forward(ctx, *args, **kwargs). setup_context is not overridden. Setting up the ctx for backward happens inside the forward.
Override forward with the signature forward(*args, **kwargs) and override setup_context. Setting up the ctx for backward happens inside setup_context (as opposed to inside the forward)

See torch.autograd.Function.forward() and extending-autograd for more details.

static vjp(ctx: Any, *grad_outputs: Any) → Any

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static vmap(info, in_dims, *args)

Define the behavior for this autograd.Function underneath torch.vmap().

For a torch.autograd.Function() to support torch.vmap(), you must either override this static method, or set generate_vmap_rule to True (you may not do both).

If you choose to override this staticmethod: it must accept

an info object as the first argument. info.batch_size specifies the size of the dimension being vmapped over, while info.randomness is the randomness option passed to torch.vmap().
an in_dims tuple as the second argument. For each arg in args, in_dims has a corresponding Optional[int]. It is None if the arg is not a Tensor or if the arg is not being vmapped over, otherwise, it is an integer specifying what dimension of the Tensor is being vmapped over.
*args, which is the same as the args to forward().

The return of the vmap staticmethod is a tuple of (output, out_dims). Similar to in_dims, out_dims should be of the same structure as output and contain one out_dim per output that specifies if the output has the vmapped dimension and what index it is in.

Please see func-autograd-function for more details.