xnn.common.finetune.reference.estimate_atomic_energies#

xnn.common.finetune.reference.estimate_atomic_energies(model, structures, species=None, batch_size=32, device=None)[source]#

Model-aware reestimation of the per-element reference energies.

Solves E_i^ref - E_i^model = sum_Z n_iZ dE0_Z for the corrections dE0 in the least-squares (minimum-norm) sense (Tompa et al. eq. 6) and returns the model’s current references plus the corrections. A model that already predicts the interaction energies correctly keeps them and only its baseline moves, which is what fine-tuning across levels of theory needs.

Parameters:

model (torch.nn.Module) – The pretrained potential (bare or wrapped). For a MultiHead, call this inside using() on model.model.
structures (iterable of dict or AtomicDataset) – The fine-tuning structures with energies.
species (sequence of int, optional) – The elements to correct; by default those present.
batch_size (int, optional) – Structures per forward pass, by default 32.
device (torch.device, optional) – Where to evaluate.

Returns:

{Z: E0}, the new reference energies.

Return type:

dict of int to float