Loss Functions

class ModifiedForceField(*args, alternative_residue_names=None, **kwargs)[source]

Takes all *args and **kwargs of openmm.app.ForceField, plus an optional parameter described here.

Parameters:

alternative_residue_names (dict) – aliases for resnames, e.g., {‘HIS’:’HIE’}.

class OpenMMPluginScoreSoftForceField(mol=None, platform='CUDA', atoms=['CA', 'C', 'N', 'CB', 'O'])[source]
Parameters:

  • mol (biobox.Molecule) – if pldataloader is not given, then a biobox object will be taken from this parameter. If neither are given then an error will be thrown.

  • xml_file (str) – xml parameter file

  • platform (str) – ‘CUDA’ or ‘Reference’.

  • remove_NB (bool) – if True remove NonbondedForce, CustomGBForce and CMMotionRemover, else just remove CustomGBForce

  • alternative_residue_names (dict) – aliases for resnames, e.g., {‘HIS’:’HIE’}.

  • atoms –

  • soft –

class OpenmmPluginScore(mol=None, xml_file=['amber14-all.xml'], platform='CUDA', remove_NB=False, alternative_residue_names={'HIS': 'HIE', 'HSE': 'HIE'}, atoms=['CA', 'C', 'N', 'CB', 'O'], soft=False)[source]

This will use the new OpenMM Plugin to calculate forces and energy. The intention is that this will be fast enough to be able to calculate forces and energy during training. N.B.: The current torchintegratorplugin only supports float on GPU and double on CPU.

Parameters:

  • mol (biobox.Molecule) – if pldataloader is not given, then a biobox object will be taken from this parameter. If neither are given then an error will be thrown.

  • xml_file (str) – xml parameter file

  • platform (str) – ‘CUDA’ or ‘Reference’.

  • remove_NB (bool) – if True remove NonbondedForce, CustomGBForce and CMMotionRemover, else just remove CustomGBForce

  • alternative_residue_names (dict) – aliases for resnames, e.g., {‘HIS’:’HIE’}.

  • atoms –

  • soft –

execute(x)[source]
Parameters:

x (torch.Tensor) – shape [b, N, 3]. dtype=float. device = gpu

get_energy(pos_ptr, force_ptr, energy_ptr, n_particles, batch_size)[source]
Parameters:

  • pos_ptr – tensor.data_ptr()

  • force_ptr – tensor.data_ptr()

  • energy_ptr – tensor.data_ptr()

  • n_particles (int) – number of particles

  • batch_size (int) – batch size

class OpenmmTorchEnergyMinimizer(mol=None, xml_file=['amber14-all.xml'], platform='CUDA', remove_NB=False, alternative_residue_names={'HIS': 'HIE', 'HSE': 'HIE'}, atoms=['CA', 'C', 'N', 'CB', 'O'], soft=False)[source]
Parameters:

  • mol (biobox.Molecule) – if pldataloader is not given, then a biobox object will be taken from this parameter. If neither are given then an error will be thrown.

  • xml_file (str) – xml parameter file

  • platform (str) – ‘CUDA’ or ‘Reference’.

  • remove_NB (bool) – if True remove NonbondedForce, CustomGBForce and CMMotionRemover, else just remove CustomGBForce

  • alternative_residue_names (dict) – aliases for resnames, e.g., {‘HIS’:’HIE’}.

  • atoms –

  • soft –

class openmm_clamped_energy_function[source]
static backward(ctx, grad_output)[source]

Defines a formula for differentiating the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs did forward() return, 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.

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 computated w.r.t. the output.

static forward(ctx, plugin, x, clamp)[source]
Parameters:

  • plugin – OpenmmPluginScore instance

  • x (torch.Tensor) – dtype = float, shape = [B, N, 3], device = cuda

Returns:

energy tensor, dtype = double, shape = [B], device CPU

class openmm_energy(mol, std, clamp=None, **kwargs)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

class openmm_energy_function[source]
static backward(ctx, grad_output)[source]

Defines a formula for differentiating the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs did forward() return, 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.

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 computated w.r.t. the output.

static forward(ctx, plugin, x)[source]
Parameters:

  • plugin – OpenmmPluginScore instance

  • x (torch.Tensor) – dtype = float, shape = [B, N, 3], device = any

Returns:

energy tensor, dtype = float, shape = [B], device = any