Loss Functions Module

The loss_functions module provides various loss functions for training Langevin dynamics simulations.

Basic Functions

gaussian(x, var=1.0, center=0.0)[source]

Gaussian function with mean center and variance var.

FT_pdf(pdf, kf, center=0.0, scale=5.0, steps=10000, kFs=None, kFsteps=100, args={}, device='cpu', **kwargs)[source]

Fourier transform (FT) of a probability density function (pdf).

Parameters:

  • pdf (callable) – the pdf whose Fourier transform is to be computed.

  • kf (float) – inverse of the variance of the pdf.

  • center (float) – mean of the pdf.

  • scale (float) – how many kf’s should we use to compute the FT.

  • steps (float) – steps of the x discretization.

  • kFs (torch.tensor) – range for the k values.

  • kFsteps (optional, int) – size for kFS if not provided.

  • args (optional, dict) – other keyword arguments taken by pdf.

  • device (str) – device to run on cpu or cuda.

  • **kwargs – extra keywords arguments ignored.

Returns:

FT of pdf, kFs

Return type:

tuple(torch.tensor, torch.tensor)

char_fn(xf, kf, scale=5.0, kFsteps=1000, device='cpu', **kwargs)[source]

Computes the characteristic function from samples xf.

Parameters:

  • xf (torch.tensor) – samples of positions

  • kf (float) – inverse variance of x

  • kFsteps (int) – number of k’s to compute the characteristic function

  • scale (float) – how many kf’s should the values of k spread.

  • **kwargs – extra keywords arguments ignored.

Returns:

characteristic function, corresponding values of k

Return type:

tuple(torch.tensor, torch.tensor)

Loss Functions

loss_fn_k(xf, kf, cf=0.0, device='cpu', scale=5.0, kFsteps=1000, x_steps=10000, **kwargs)[source]

Loss function comparing the L2 mean square loss of the characteristic function of the pdf to the target normal distribution with variance 1/kf and center cf. Any additional keyword arguments are ignored.

Parameters:

  • xf (torch.Tensor) – The final position of the particles.

  • kf (float) – The final stiffness value.

  • cf (float, optional) – The final center value. Defaults to 0.0.

  • device (torch.device, optional) – The device on which to perform the computations. Defaults to the current device.

  • scale (float, optional) – how many kf’s should the values of k spread. Defaults to 5.0.

  • x_steps (float) – steps of the x discretization.

  • kFs (torch.tensor) – range for the k values.

  • kFsteps (optional, int) – size for kFS if not provided.

  • **kwargs – extra keywords arguments ignored.

Returns:

The MSE loss between the computed characteristic function and the target characteristic function.

Return type:

float

loss_fn_variance(xf, kf, **kwargs)[source]

Loss function comparing square difference of theoretical variance 1/kf vs computed variance of xf.

Parameters:

  • xf (torch.Tensor) – The final position of the particles.

  • kf (float) – The final stiffness value.

  • **kwargs – extra keywords arguments ignored.

Returns:

The squared difference between the computed variance and the theoretical variance.

Return type:

float

loss_fn_mean(xf, cf, **kwargs)[source]

Loss function comparing square difference of theoretical mean cf vs computed mean of xf.

Parameters:

  • xf (torch.Tensor) – The final position of the particles.

  • cf (float) – The theoretical mean value.

  • **kwargs – extra keywords arguments ignored.

Returns:

The squared difference between the computed mean and the theoretical mean.

Return type:

float

loss_fn_grad_k(ki, kf, sim, **kwargs)[source]

Penalizes large variations of kappa (the stiffness of the harmonic potential). If the force is continuous, the function includes the edge values ki and kf in the computation. The function computes the difference between consecutive elements of ks (which represents kappa values), squares these differences, and then returns the mean of these squared differences.

Parameters:

  • ki (float) – The initial stiffness value.

  • kf (float) – The final stiffness value.

  • sim (Simulator) – The simulator object.

  • **kwargs – extra keywords arguments ignored.

Returns:

The mean of the squared differences between consecutive kappa values.

Return type:

float

loss_fn_control_k_vars(xf, kf, ki, sim, device='cpu', scale=5.0, kFsteps=1000, x_steps=1000, blend=0.001, **kwargs)[source]

Loss function that minimizes distance to equilibrium and penalizes strong variations of consecutive values of kappa. This function computes a weighted sum of two loss functions: loss_fn_k and loss_fn_grad_k. The weight of each loss function is determined by the blend parameter.

Parameters:

  • xf (torch.Tensor) – The final position of the particles.

  • kf (float) – The final stiffness value.

  • ki (float) – The initial stiffness value.

  • sim (Simulator) – The simulator object.

  • device (torch.device, optional) – The device on which to perform the computations. Defaults to the current device.

  • scale (float, optional) – how many kf’s should the values of k spread. Defaults to 5.0.

  • x_steps (float) – steps of the x discretization.

  • kFs (torch.tensor) – range for the k values.

  • kFsteps (optional, int) – size for kFS if not provided.

  • blend (float, optional) – The blending factor for the two loss functions. Defaults to 1e-3.

  • **kwargs – extra keywords arguments ignored.

Returns:

The weighted sum of the two loss functions.

Return type:

float

loss_fn_work(sim, **kwargs)[source]

Loss function to minimize work with respect to lower bound Delta F.

This function returns the mean work done during the simulation, which is stored in sim.w. The goal is to minimize this work with respect to a lower bound Delta F.

Parameters:

  • sim (Simulator) – The simulator object, which contains the work done during the simulation.

  • **kwargs – extra keywords arguments ignored.

Returns:

The mean work done during the simulation.

Return type:

float

loss_fn_eq_work(xf, kf, sim, device='cpu', scale=5.0, kFsteps=1000, x_steps=1000, blend=0.001, **kwargs)[source]

Loss function to minimize work and distance to equilibrium. This function computes a weighted sum of two loss functions: loss_fn_k and loss_fn_work. The weight of each loss function is determined by the blend parameter.

Parameters:

  • xf (torch.Tensor) – The final position of the particles.

  • kf (float) – The final stiffness value.

  • sim (Simulator) – The simulator object.

  • device (torch.device, optional) – The device on which to perform the computations. Defaults to the current device.

  • scale (float, optional) – how many kf’s should the values of k spread. Defaults to 5.0.

  • x_steps (float) – steps of the x discretization.

  • kFs (torch.tensor) – range for the k values.

  • kFsteps (optional, int) – size for kFS if not provided.

  • blend (float, optional) – The blending factor for the two loss functions. Defaults to 1e-3.

  • **kwargs – extra keywords arguments ignored.

Returns:

The weighted sum of the two loss functions: work and caracteristic function comparison to the equilibrium one.

Return type:

float

Example Usage

Here’s an example of how to use the loss functions:

import torch
from langesim_optim.loss_functions import loss_fn_k, gaussian

# Create some final positions
xf = torch.randn(1000)

# Calculate loss for a specific stiffness
kf = 2.0
loss = loss_fn_k(
    xf=xf,
    kf=kf,
    cf=0.0,  # Center
    scale=5.0,
    kFsteps=1000
)

# Calculate Gaussian distribution
x = torch.linspace(-3, 3, 100)
pdf = gaussian(x, var=1.0, center=0.0)