loss_function

Module containing the loss functions for the ML-DFT model.

class CoefficientLoss(loss_function: Module = None, sample_weigher: SampleWeigher = None, reduction: str = 'mean')

Calculates the loss between the predicted difference of coefficients and the target difference of coefficients.

get_loss(batch: Batch, pred_diff: Tensor, **_) → Tensor

Returns the loss between the predicted difference of coefficients and the target difference of coefficients from the batch object.

Parameters:

• batch (Batch) – Batch object containing the target difference of coefficients

• pred_diff (Tensor) – the predicted difference of the ground state coefficients and the current coefficients

Returns:

loss between the predicted difference of coefficients and the target difference of coefficients,

per basis function

Return type:

Tensor

class EnergyGradientLoss(loss_function: Module = None, sample_weigher: SampleWeigher = None, reduction: str = 'mean')

Calculates the projected loss between the predicted gradients and the target gradients.

get_loss(batch: Batch, projected_gradient_difference: Tensor, **_) → Tensor

Computes the loss to the projected difference of the predicted and ground truth gradient. Since the input is already the difference to the ground truth, the label is set to zero in the loss function call.

Parameters:

• batch (Batch) – Batch object containing the target energy

• projected_gradient_difference (Tensor) – the difference of the predicted and real gradient, after projection

Returns:

loss function of the input

Return type:

Tensor

class EnergyLoss(loss_function: Module = None, sample_weigher: SampleWeigher = None, reduction: str = 'mean')

Calculates the loss between the predicted energy and the target energy.

get_loss(batch: OFData, pred_energy: Tensor, **_) → Tensor

Computes the loss between the predicted energy and target energy.

Parameters:

• batch (Batch) – Batch object containing the target energy

• pred_energy (Tensor) – Tensor containing the predicted energy

Returns:

the loss for each sample in the batch

Return type:

Tensor

class FullLoss

Previous version of the WeightedLoss module, which is no longer supported.

Included to make loading old checkpoints possible.

__init__()

Just raises an error.

class PerSampleWeightedPerCoeffLossFunction(loss_function: Module = None, sample_weigher: SampleWeigher = None, reduction: str = 'mean')

Base class for loss functions that compute a loss value per basis function.

weigh_loss(batch: OFData, loss: Tensor) → Tensor

Applies weights to the loss for each basis function.

class PerSampleWeightedPerSampleLossFunction(loss_function: Module = None, sample_weigher: SampleWeigher = None, reduction: str = 'mean')

Base class for loss functions that compute a loss value per sample.

weigh_loss(batch: OFData, loss: Tensor) → Tensor

Applies weights to the loss for each sample.

class SingleLossFunction(loss_function: Module = None, sample_weigher: SampleWeigher = None, reduction: str = 'mean')

Base class for loss functions that compute a single loss value.

__init__(loss_function: Module = None, sample_weigher: SampleWeigher = None, reduction: str = 'mean')

Initialize the LossFunction by setting the loss function, weighing and reduction.

Parameters:

• loss_function (nn.Module, optional) – Loss function to be used. Defaults to nn.L1Loss(). Important The loss function should not apply any reduction, as this is handled by the LossFunction after weighting the loss.

• sample_weigher (SampleWeigher, optional) – Sample weigher to be used. Defaults to None.

• reduction (str, optional) – Reduction type to be used. Defaults to “mean”.

forward(batch: OFData, **kwargs) → Tensor

Get the per-sample losses and apply the sample weights, as defined by the sample_weigher.

Parameters:

• batch (OFData) – The batch object, used in the loss calculation and for the sample weights

• **kwargs – Additional arguments to be passed to the loss function

Returns:

The scalar loss

Return type:

Tensor

get_loss(batch: OFData, **kwargs) → Tensor

Function that computes the loss.

Has to be implemented by subclasses.

weigh_loss(batch: OFData, loss: Tensor) → Tensor

Function that applies weights to the loss.

Has to be implemented by subclasses.

class WeightedLoss(**kwargs: Mapping[str, float | Module])

Module used to combine multiple losses with different weights.

The forward pass does not return a single scalar loss, but rather two dictionaries containing the weights and the losses for each individual component.

__init__(**kwargs: Mapping[str, float | Module])

Initialize the WeightedLoss object by passing a mapping of str to loss function and weight.

Parameters:

**kwargs – Mapping of loss names to the corresponding loss functions and weights. The names will be used for logging purposes. The values should be dictionaries with the keys “loss” and “weight”, the former containing the (nn.Module) loss function and the latter the (scalar) weight.

forward(batch: OFData, **kwargs) → tuple[dict[str, Tensor], dict[str, Tensor]]

The weighted sum of the energy loss, the gradient loss and the coefficient loss.

Loss = energy_weight * energy_loss + gradient_weight * gradient_loss + coefficient_weight * coefficient_loss.

Parameters:

• batch (Batch) – Batch object containing the target energy and target gradients

• **kwargs – Additional arguments to be passed to the loss functions

Returns:

Dictionary containing the weights for each loss component dict[str, Tensor]: Dictionary containing the losses for each loss component

Return type:

dict[str, Tensor]

project_gradient(gradient: Tensor, batch_or_sample: OFData) → Tensor

Calculates the projected gradient.

Note that when using non-orthogonal transformations we need the dual vector of \(w\), written as \(w^*\), since it will transform different to \(w^T\). For orthogonal transformations we have \(w^* = w^T\). We get rid of the matrix multiplications by rewriting the equation as follows:

\[\left( \boldsymbol{I} - \frac{\boldsymbol{w}\boldsymbol{w}^T}{\boldsymbol{w}^T \boldsymbol{w}} \right) \nabla_\boldsymbol{p} T = \nabla_\boldsymbol{p} T - \boldsymbol{w} \frac{\boldsymbol{w}^T \nabla_\boldsymbol{p} T} {{\boldsymbol{w}^T \boldsymbol{w}}}\]

Parameters:

• gradient (torch.Tensor) – The predicted gradients

• batch_or_sample (OFData) – The OFData object (can be a batch or a single sample) containing the basis integrals.

Returns:

The projected gradient

Return type:

torch.Tensor

project_gradient_difference(pred_gradients: Tensor, batch: OFData) → Tensor

Calculates the projected gradient difference and the absolute projected gradient error.

We get rid of the matrix multiplications by rewriting the equation as follows:

\[\left( \boldsymbol{I} - \frac{\boldsymbol{w}\boldsymbol{w}^T}{\boldsymbol{w}^T \boldsymbol{w}} \right) \left(\nabla_\boldsymbol{p} T_{\mathrm{pred}} - \nabla_\boldsymbol{p} T \right) = \left(\nabla_\boldsymbol{p} T_{\mathrm{pred}}-\nabla_\boldsymbol{p} T\right) - \boldsymbol{w} \frac{\boldsymbol{w}^T \left(\nabla_\boldsymbol{p} T_{\text{pred}}-\nabla_\boldsymbol{p} T\right)} {{\boldsymbol{w}^T \boldsymbol{w}}}\]

Parameters:

• pred_gradients (torch.Tensor) – The predicted gradients

• batch (Batch) – The batch object containing the target gradients and the basis integrals

Returns:

Unreduced tensor of projected differences