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(*args, **kwargs)[source]
static backward(ctx, grad_output)[source]

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 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]

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

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

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 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