Skip to content

EnergyTerm

Jeff Flatten edited this page Jun 6, 2024 · 5 revisions

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.

Torch vs C++/CUDA

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.

Clone this wiki locally