mtenn.strategy.SplitDeltaStrategy

class mtenn.strategy.SplitDeltaStrategy(complex_energy_func, ligand_energy_func=None, protein_energy_func=None)

Bases: Strategy

Simple strategy for subtracting the sum of the individual component energies from the complex energy. This Strategy requires an individual energy_func (\(\phi: \mathbb{R}^n \rightarrow \mathbb{R}\) that maps from an n-dimensional vector representation (output from a Representation block) to a scalar-value energy prediction) for the complex, ligand, and protein representations. As with SplitModel, the complex_energy_func is required, and any missing values in the ligand_energy_func or protein_energy_func will be filled in with the complex_energy_func.

The point of this class is to be able to handle different dimensionalities in the different representations that may be generated by a SplitModel. In the simplest case, one could like ligand_energy_func and protein_energy_func blank to achieve the same behavior as the standard :py:class:DeltaStrategy.

\[ \begin{align}\begin{aligned}\mathrm{G_{complex}} &= \phi_{\mathrm{complex}} (\mathrm{\boldsymbol{x}_{complex}})\\\mathrm{G_{ligand}} &= \phi_{\mathrm{ligand}} (\mathrm{\boldsymbol{x}_{ligand}})\\\mathrm{G_{protein}} &= \phi_{\mathrm{protein}} (\mathrm{\boldsymbol{x}_{protein}})\\\Delta \mathrm{G_{pred}} &= \mathrm{G_{complex}} - (\mathrm{G_{ligand}} + \mathrm{G_{protein}})\end{aligned}\end{align} \]

Methods

__init__(complex_energy_func[, ...])	Store module for predicting an energy from representation.
forward(comp, prot, lig)	Make energy predictions for each representation, and then perform the delta calculation.

__init__(complex_energy_func, ligand_energy_func=None, protein_energy_func=None)

Store module for predicting an energy from representation.

Parameters:

• complex_energy_func (torch.nn.Module) – Some torch module that will predict an energy from an n-dimension vector representation of a structure

• ligand_energy_func (torch.nn.Module, optional) – Some torch module that will predict an energy from an n-dimension vector representation of a structure

• protein_energy_func (torch.nn.Module, optional) – Some torch module that will predict an energy from an n-dimension vector representation of a structure

forward(comp, prot, lig)

Make energy predictions for each representation, and then perform the delta calculation.

Parameters:

• comp (torch.Tensor) – Complex representation that will be passed to self.energy_func

• parts (list[torch.Tensor], optional) – Representations for all individual parts of the complex (eg ligand and protein separately) that will be passed to self.energy_func

Returns:

Predicted \(\Delta G\) value

Return type:

torch.Tensor