mtenn.strategy.ComplexOnlyStrategy

class mtenn.strategy.ComplexOnlyStrategy(energy_func)

Bases: Strategy

Strategy to only predict based on the complex representation. This is useful if you want to make a prediction on just the ligand or just the protein, and essentially just reduces to a standard version of whatever your underlying model is.

Methods

__init__(energy_func)	Store module for predicting an energy from representation.
forward(comp, *parts)	Make energy prediction for the complex representation.

__init__(energy_func)

Store module for predicting an energy from representation.

Parameters:

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

forward(comp, *parts)

Make energy prediction for the complex representation.

Parameters:

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

• parts (list[torch.Tensor], optional) – Ignored, but present just to match the signatures

Returns:

Predicted value

Return type:

torch.Tensor