credit.loss#

Attributes#

Classes#

LogCoshLoss

Log-Cosh Loss Function.

XTanhLoss

X-Tanh Loss Function.

XSigmoidLoss

X-Sigmoid Loss Function.

MSLELoss

Mean Squared Logarithmic Error (MSLE) Loss Function.

KCRPSLoss

Adapted from Nvidia Modulus

SpectralLoss2D

Spectral Loss in 2D.

PSDLoss

Power Spectral Density (PSD) Loss Function.

VariableTotalLoss2D

Custom loss function class for 2D geospatial data

Functions#

load_loss(loss_type[, reduction])

Load a specified loss function by its type.

latitude_weights(conf)

Calculate latitude-based weights for loss function.

variable_weights(conf, channels, frames)

Create variable-specific weights for different atmospheric

Module Contents#

credit.loss.logger#
credit.loss.load_loss(loss_type, reduction='mean')#

Load a specified loss function by its type. Helper function of VariableTotalLoss2D

This function returns a loss function based on the specified loss_type. It supports several common loss functions, including MSE, MAE, MSLE, Huber, Log-Cosh, X-Tanh, and X-Sigmoid. The loss function can also be customized to use different reduction methods (e.g., ‘mean’, ‘sum’). Use reduction=none if using latitude or variable weights

Parameters:
  • loss_type (str) – The type of loss function to load. Supported values are “mse”, “mae”, “msle”, “huber”, “logcosh”, “xtanh”, and “xsigmoid”.

  • reduction (str, optional) – Specifies the reduction to apply to the output: ‘mean’ (default) or ‘sum’.

Returns:

The corresponding loss function.

Return type:

torch.nn.Module

Raises:

ValueError – If the specified loss_type is not supported.

Example

>>> loss_fn = load_loss("mse")
>>> loss = loss_fn(pred, target)
class credit.loss.LogCoshLoss(reduction='mean')#

Bases: torch.nn.Module

Log-Cosh Loss Function.

This loss function computes the logarithm of the hyperbolic cosine of the prediction error. It is less sensitive to outliers compared to the Mean Squared Error (MSE) loss.

Parameters:

reduction (str) – Specifies the reduction to apply to the output. ‘mean’ | ‘none’. ‘mean’: the output is averaged; ‘none’: no reduction is applied.

reduction = 'mean'#
forward(y_t, y_prime_t)#

Forward pass for Log-Cosh loss.

Parameters:
  • y_t (torch.Tensor) – Target tensor.

  • y_prime_t (torch.Tensor) – Predicted tensor.

Returns:

Log-Cosh loss value.

Return type:

torch.Tensor

class credit.loss.XTanhLoss(reduction='mean')#

Bases: torch.nn.Module

X-Tanh Loss Function.

This loss function computes the element-wise product of the prediction error and the hyperbolic tangent of the error. This loss function aims to be more robust to outliers than traditional MSE.

Parameters:

reduction (str) – Specifies the reduction to apply to the output: ‘mean’ | ‘none’. ‘mean’: the output is averaged; ‘none’: no reduction is applied.

reduction = 'mean'#
forward(y_t, y_prime_t)#

Forward pass for X-Tanh loss.

Parameters:
  • y_t (torch.Tensor) – Target tensor.

  • y_prime_t (torch.Tensor) – Predicted tensor.

Returns:

X-Tanh loss value.

Return type:

torch.Tensor

class credit.loss.XSigmoidLoss(reduction='mean')#

Bases: torch.nn.Module

X-Sigmoid Loss Function.

This loss function computes a modified loss by using a sigmoid function transformation. It is designed to handle large errors in a non-linear fashion.

Parameters:

reduction (str) – Specifies the reduction to apply to the output. ‘mean’ | ‘none’. ‘mean’: the output is averaged; ‘none’: no reduction is applied.

reduction = 'mean'#
forward(y_t, y_prime_t)#

Forward pass for X-Sigmoid loss.

Parameters:
  • y_t (torch.Tensor) – Target tensor.

  • y_prime_t (torch.Tensor) – Predicted tensor.

Returns:

X-Sigmoid loss value.

Return type:

torch.Tensor

class credit.loss.MSLELoss(reduction='mean')#

Bases: torch.nn.Module

Mean Squared Logarithmic Error (MSLE) Loss Function.

This loss function computes the mean squared logarithmic error between the predicted and target values. It is useful for handling targets that span several orders of magnitude.

Parameters:

reduction (str) – Specifies the reduction to apply to the output. ‘mean’ | ‘none’. ‘mean’: the output is averaged; ‘none’: no reduction is applied.

reduction = 'mean'#
forward(prediction, target)#

Forward pass for MSLE loss.

Parameters:
  • prediction (torch.Tensor) – Predicted tensor.

  • target (torch.Tensor) – Target tensor.

Returns:

MSLE loss value.

Return type:

torch.Tensor

class credit.loss.KCRPSLoss(reduction, biased: bool = False)#

Bases: torch.nn.Module

Adapted from Nvidia Modulus pred : Tensor

Tensor containing the ensemble predictions. The ensemble dimension is assumed to be the leading dimension

obsUnion[Tensor, np.ndarray]

Tensor or array containing an observation over which the CRPS is computed with respect to.

biased :

When False, uses the unbiased estimators described in (Zamo and Naveau, 2018):

E|X-y|/m - 1/(2m(m-1)) sum_(i,j=1)|x_i - x_j|

Unlike crps this is fair for finite ensembles. Non-fair crps favors less dispersive ensembles since it is biased high by E|X- X’|/ m where m is the ensemble size.

Estimate the CRPS from a finite ensemble

Computes the local Continuous Ranked Probability Score (CRPS) by using the kernel version of CRPS. The cost is O(m log m).

Creates a map of CRPS and does not accumulate over lat/lon regions. Approximates:

\[CRPS(X, y) = E[X - y] - 0.5 E[X-X']\]

with

\[sum_i=1^m |X_i - y| / m - 1/(2m^2) sum_i,j=1^m |x_i - x_j|\]
biased = False#
batched_forward#
forward(target, pred)#
single_sample_forward(target, pred)#

Forward pass for KCRPS loss for a single sample.

Parameters:
  • target (torch.Tensor) – Target tensor.

  • pred (torch.Tensor) – Predicted tensor.

Returns:

CRPS loss values at each lat/lon

Return type:

torch.Tensor

_kernel_crps_implementation(pred: torch.Tensor, obs: torch.Tensor, biased: bool) torch.Tensor#

An O(m log m) implementation of the kernel CRPS formulas

class credit.loss.SpectralLoss2D(wavenum_init=20, reduction='none')#

Bases: torch.nn.Module

Spectral Loss in 2D.

This loss function compares the spectral (frequency domain) content of the predicted and target outputs using FFT. It is useful for ensuring that the predicted output has similar frequency characteristics as the target.

Parameters:
  • wavenum_init (int) – The initial wavenumber to start considering in the loss calculation.

  • reduction (str) – Specifies the reduction to apply to the output: ‘mean’ | ‘none’. ‘mean’: the output is averaged; ‘none’: no reduction is applied.

wavenum_init = 20#
reduction = 'none'#
forward(output, target, weights=None, fft_dim=-1)#

Forward pass for Spectral Loss 2D.

Parameters:
  • output (torch.Tensor) – Predicted tensor.

  • target (torch.Tensor) – Target tensor.

  • weights (torch.Tensor, optional) – Latitude weights for the loss.

  • fft_dim (int) – The dimension to apply FFT.

Returns:

Spectral loss value.

Return type:

torch.Tensor

class credit.loss.PSDLoss(wavenum_init=20)#

Bases: torch.nn.Module

Power Spectral Density (PSD) Loss Function.

This loss function calculates the Power Spectral Density (PSD) of the predicted and target outputs and compares them to ensure similar frequency content in the predictions.

Parameters:

wavenum_init (int) – The initial wavenumber to start considering in the loss calculation.

wavenum_init = 20#
forward(target, pred, weights=None)#

Forward pass for PSD loss.

Parameters:
  • target (torch.Tensor) – Target tensor.

  • pred (torch.Tensor) – Predicted tensor.

  • weights (torch.Tensor, optional) – Latitude weights for the loss.

Returns:

PSD loss value.

Return type:

torch.Tensor

get_psd(f_x, device, dtype)#
credit.loss.latitude_weights(conf)#

Calculate latitude-based weights for loss function. This function calculates weights based on latitude values to be used in loss functions for geospatial data. The weights are derived from the cosine of the latitude and normalized by their mean.

Parameters:

conf (dict) – Configuration dictionary containing the path to the latitude weights file.

Returns:

A 2D tensor of weights with dimensions

corresponding to latitude and longitude.

Return type:

torch.Tensor

credit.loss.variable_weights(conf, channels, frames)#

Create variable-specific weights for different atmospheric and surface channels.

This function loads weights for different atmospheric variables (e.g., U, V, T, Q) and surface variables (e.g., SP, t2m) from the configuration file. It then combines them into a single weight tensor for use in loss calculations.

Parameters:
  • conf (dict) – Configuration dictionary containing the variable weights.

  • channels (int) – Number of channels for atmospheric variables.

  • frames (int) – Number of time frames.

Returns:

A tensor containing the combined weights for

all variables.

Return type:

torch.Tensor

class credit.loss.VariableTotalLoss2D(conf, validation=False)#

Bases: torch.nn.Module

Custom loss function class for 2D geospatial data with optional spectral and power loss components.

This class defines a loss function that combines a base loss (e.g., L1, MSE) with optional spectral and power loss components for 2D geospatial data. The loss function can incorporate latitude and variable-specific weights.

Parameters:
  • conf (dict) – Configuration dictionary containing loss function settings and weights.

  • validation (bool, optional) – If True, the loss function is used in validation mode. Defaults to False.

conf#
training_loss#
vars#
lat_weights = None#
var_weights = None#
use_spectral_loss#
use_power_loss#
validation = False#
forward(target, pred)#

Calculate the total loss for the given target and prediction.

This method computes the base loss between the target and prediction, applies latitude and variable weights, and optionally adds spectral and power loss components.

Parameters:
  • target (torch.Tensor) – Ground truth tensor.

  • pred (torch.Tensor) – Predicted tensor.

Returns:

The computed loss value.

Return type:

torch.Tensor