-
Notifications
You must be signed in to change notification settings - Fork 3
EnergyTerm
An EnergyTerm
is a class that evaluates one or more energy functions on a PoseStack.
Multiple EnergyTerms can be combined within a ScoreFunction, which allows for a complete evaluation of the energy of a PoseStack.
EnergyTerms are evaluated in a block-centric manner. Each EnergyTerm breaks down the work on a per-block or per-block-pair basis, depending on if the EnergyTerm is a 1-body or 2-body term. This per-block and per-block-pair work is all dispatched simultaneously, letting the GPU handle the scheduling of the individual block-based calculations.
EnergyTerms may be coded entirely in pure torch
(such as the RefEnergyTerm), or may be coded in C++/CUDA to enhance performance (HBondEnergyTerm).
Pure torch EnergyTerms have their gradients calculated automatically by torch
autograd. For C++/CUDA, derivatives must be calculated manually.
The arguments for the exported functions must all be passed as Tensors or Ints. On the C++ side, the arguments will be unpacked and any tensors will be assigned their associated dimensionality and datatypes by the TCAST
function. This function infers the tensor's dimensionality and datatypes based on the argument type of the function it is being sent to.
Warning
It is critical that the order, dimensionality and type of the tensor arguments match in both the calling Python code and the C++ function which consumes them. Any mistakes here may cause very hard to read errors.