Utilities Module

The utilities module provides essential functions for training, visualization, and analysis of Langevin simulations.

Training Functions

train_loop(epochs, sim, tot_sims, optimizer, loss_fn, ki=1.0, kf=1.0, ci=0.0, cf=0.0, scheduler=None, device='cpu', init_epoch=0, **kwargs)[source]

Trains model sim, returns the loss list. Perform the training loop for a given number of epochs.

Parameters:

  • epochs (int) – The number of training epochs.

  • sim (Simulator) – The simulator object used for training.

  • tot_sims (int) – The total number of simulations.

  • ki – The initial value of stiffness.

  • kf – The final value of stiffness.

  • optimizer – The optimizer used for training.

  • scheduler – The learning rate scheduler.

  • loss_fn – The loss function used for training. Takes as arguments (xf, kf, ki, sim, device, **kwargs)

  • device (str, optional) – The device where the training will be performed. Defaults to the ‘device’ constant.

  • init_epoch (int, optional) – initial epoch number for this round. Defaults to 0.

  • **kwargs – Additional keyword arguments specific to the loss function.

Returns:

A list of the loss function values for each epoch.

Return type:

List

train_loop_snapshots(epochs, sim, tot_sims, ki, kf, tf, optimizer, loss_fn, scheduler=None, device='cpu', snapshot_step=1, xrange=1.0, bins=200, times=None, time_ticks=None, yrange=None, y_ticks=None, y_ticklabels=None, grid_histo=False, grid_protocol=True, **kwargs)[source]

Trains model sim, returns the loss list. Perform the training loop for a given number of epochs.

Parameters:

  • epochs (int) – The number of training epochs.

  • sim (Simulator) – The simulator object used for training.

  • tot_sims (int) – The total number of simulations.

  • ki – The initial value of stiffness.

  • kf – The final value of stiffness.

  • optimizer – The optimizer used for training.

  • scheduler – The learning rate scheduler.

  • loss_fn – The loss function used for training. Takes as arguments (xf, kf, ki, sim, device, **kwargs)

  • device (str, optional) – The device where the training will be performed. Defaults to the ‘device’ constant.

  • snapshot_step (int, optional) – Makes a snapshot of the protocol and graphs of the final distribution every snapshot_step epochs. Defaults to 1.

  • **kwargs – Additional keyword arguments specific to the loss function.

Returns:

(lossl, protocols, plots), with lossl = A list of the loss function values for each epoch. protocols = A list of the snapshot of the protocols plots = A list of the plots of the final distribution at each snapshot step.

Return type:

Tuple

Visualization Functions

plot_test_hist(tot_sims, sim, ki, kf, ci=0.0, cf=0.0, device='cpu', xrange=1.0, xpoints=200, init_label='Initial equilibrium PDF', expect_label='Expected final equilibrium PDF', final_label='Final PDF after simulation', title='Probability density function (PDF) of the position', grid=False, bins=200)[source]

Plots a histogram of final PDF after a forward pass on the Simulator sim of an initial equilibrium PDF with stiffness ki.

Parameters:

  • tot_sims (int) – total simulations to be performed.

  • ki (float) – initial stiffness

  • kf (float) – final stiffness

  • ci (float, optional) – initial center. Defaults to 0.0.

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

  • sim (Simulator) – simulator module

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

plot_protocols(sim, ki, kf, tf, k_comp=None, t_steps=200, sim_label='Trained', comp_label='Theoretical', title='Stiffness $\\\\kappa$', times=None, yrange=None, ylabel='$\\\\kappa$', grid=True, time_ticks=None, y_ticks=None, y_ticklabels=None, protocol_name='kappa')[source]

Plots a protocol of the simulator sim.force.protocol_name(t) and the theoretical protocol k_comp(t).

Parameters:

  • sim (Simulator) – The simulator object.

  • ki – The initial value of stiffness.

  • kf – The final value of stiffness.

  • tf – The final time of the protocol.

  • k_comp – A function that computes the theoretical value of kappa(t) given the initial and final stiffness and the final time.

  • t_steps – The number of steps to compute the protocol.

  • sim_label – The label for the simulator protocol.

  • comp_label – The label for the theoretical protocol.

  • title – The title of the plot.

  • times – The times to compute the protocol.

  • yrange – The range of the y-axis.

  • grid – Whether to show the grid.

  • time_ticks – The ticks for the time axis.

  • y_ticks – The ticks for the y-axis.

  • y_ticklabels – The labels for the y-axis ticks.

  • protocol_name – The name of the protocol to plot. Defaults to “kappa”.

Returns:

The matplotlib figure object.

Return type:

Figure

Helper Functions

k_from_sim(sim)[source]

Extracts ki, kf, tf and k(t) from a simulator.

Parameters:

sim (Simulator)

Example Usage

Here’s a basic example of how to use the training loop:

import langesim_optim as lso
from langesim_optim.utilities import train_loop

# Create a simulator
sim = lso.Simulator(...)

# Set up optimizer
optimizer = torch.optim.Adam(sim.parameters(), lr=0.001)

# Train the model
losses = train_loop(
    epochs=100,
    sim=sim,
    tot_sims=1000,
    optimizer=optimizer,
    loss_fn=your_loss_function,
    ki=1.0,
    kf=2.0
)