Simulator and Forces Module

The simulator_forces module provides the core simulation functionality and force implementations for Langevin dynamics.

Base Classes

class BaseForce[source]

Base class for force on the particle. All forces should inherit from it. This base class returns a zero force at all times.

__init__()[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

force(x, t)[source]

This should be implemented by inherited classes. This base class returns a zero force.

Parameters:
  • x (torch.tensor) – position where the force is applied.

  • t (torch.tensor) – time at which the force is applied.

Returns:

The force.

Return type:

torch.tensor

potential(x, t)[source]

This should be implemented by inherited classes. Although force = - grad potential, it is more efficient (speed and memory) to compute them apart analytically. This base class returns a zero potential.

Parameters:
  • x (torch.tensor) – position where the potential energy is computed.

  • t (torch.tensor) – time at which the potential energy is computed.

Returns:

The potential energy.

Return type:

torch.tensor

forward(x, t)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class BaseHarmonicForce[source]

Bases: BaseForce

Base class for harmonic oscillator force.

__init__()[source]

Initializes the force.

kappa(t)[source]

Stiffness of the harmonic force. This should be defined in inherited classes. By default it is 0.0.

Parameters:

t (torch.tensor) – current time

Returns:

the stiffness at time t.

Return type:

torch.tensor

center(t)[source]

Center of the harmonic force. This should be defined in inherited classes. Defaults to 0.0.

Parameters:

t (torch.tensor) – current time

Returns:

the center at time t.

Return type:

torch.tensor

force(x, t)[source]

Computes and returns a harmonic oscillator force with time dependent stiffness kappa(t) and center(t)

Parameters:
  • x (torch.tensor) – position where the force is applied.

  • t (torch.tensor) – time at which the force is applied.

Returns:

The force.

Return type:

torch.tensor

potential(x, t)[source]

Computes and returns a harmonic oscillator potential with time dependent stiffness kappa(t) and center(t)

Parameters:
  • x (torch.tensor) – position where the force is applied.

  • t (torch.tensor) – time at which the force is applied.

Returns:

The potential energy.

Return type:

torch.tensor

Simulator

class Simulator(dt=0.001, tot_steps=10000, noise_scaler=1.0, force=BaseForce(), device='cpu', compute_work_heat=False)[source]

Simulator class for Langevin dynamics of a harmonic oscillator.

This class implements a numerical solver for the overdamped Langevin equation: dx/dt = F(x,t) + sqrt(2*noise_scaler)*dW(t) where F(x,t) is the applied force, noise_scaler is diffusion coefficient D=k_B*T/gamma, and dW(t) is Gaussian white noise.

__init__(dt=0.001, tot_steps=10000, noise_scaler=1.0, force=BaseForce(), device='cpu', compute_work_heat=False)[source]

Initializes the Simulator.

Parameters:
  • tot_sims (int, optional) – total number of simulations. Defaults to 1000.

  • dt (float, optional) – time step. Defaults to 0.001.

  • tot_steps (int, optional) – total steps of each simulation. Defaults to 10000.

  • noise_scaler (float, optional) – brownian noise scale k_B T. Defaults to 1.0.

  • force (BaseForce, optional) – force exerted on the particle. Defaults to 0.0.

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

  • compute_work_heat (bool) – compute work and heat during simulation. Defaults to False.

forward(x0)[source]

Forward pass: Performs in parallel N simulations with initial condition x0 where N=shape(x0)

Parameters:
  • x0 (torch.tensor) – tensor containing the initial condition.

  • be (Its shape should)

Returns:

final positions of the particles, if compute_work_heat=False Tuple[torch.tensor, torch.tensor, torch.tensor]: final positions of the particles, final work, final heat, if compute_work_heat=True

Return type:

torch.tensor

Force Implementations

class VariableStiffnessHarmonicForce(kappai, kappaf, tf, steps=None, k=None, continuous=True)[source]

Bases: BaseHarmonicForce

Harmonic oscillator force with a variable stiffness that is a learnable parameter. The stiffness is made of linear interpolation of segments. The center is fixed at zero. To do: change to a variable center also.

Parameters:
  • kappai (float)

  • kappaf (float)

  • tf (float)

  • steps (int | None)

  • k (List | None)

__init__(kappai, kappaf, tf, steps=None, k=None, continuous=True)[source]

Initializes the force with a stiffness that is defined at steps values of time and linearly interpolated between.

Parameters:
  • kappai (float) – initial stiffness

  • kappaf (float) – final stiffness

  • tf (float) – final time of the protocol

  • steps (int, optional) – number time steps where the stiffness value is given

  • k (list, optional) – the initial stiffness given by a list of steps values

  • countinuous (bool) – whether the stiffness is continuous at initial and final times,

  • kappaf. (thus equal to kappai and)

kappa(t)[source]

Stiffness given as an interpolation between the values given by the list k.

Parameters:

t – time to compute the stiffness

Returns:

the stiffness value at time t

Return type:

torch.tensor

class VariableCenterHarmonicForce(centeri, centerf, tf, steps=None, center_list=None, kappa0=1.0, continuous=True)[source]

Bases: BaseHarmonicForce

Harmonic oscillator force with a variable center that is a learnable parameter.

Parameters:
  • centeri (float)

  • centerf (float)

  • tf (float)

  • steps (int | None)

  • center_list (List | None)

  • kappa0 (float)

__init__(centeri, centerf, tf, steps=None, center_list=None, kappa0=1.0, continuous=True)[source]

Initializes the force with a center that is defined at steps values of time and linearly interpolated between.

Parameters:
  • centeri (float) – initial center

  • centerf (float) – final center

  • tf (float) – final time of the protocol

  • steps (int, optional) – number time steps where the center value is given

  • center (list, optional) – the initial center given by a list of

  • values (steps)

  • kappa0 (float) – fixed stiffness

  • countinuous (bool) – whether the center is continuous at initial and final times,

  • centerf. (thus equal to centeri and)

  • center_list (List | None)

center(t)[source]

Center given as an interpolation between the values given by the list center.

Parameters:

t – time to compute the center

Returns:

the center value at time t

Return type:

torch.tensor

kappa(t)[source]

Stiffness given as a constant value kappa0.

Parameters:

t – time to compute the stiffness

class VariableStiffnessCenterHarmonicForce(kappai, kappaf, centeri, centerf, tf, steps=None, k=None, center_list=None, continuous=True)[source]

Bases: VariableStiffnessHarmonicForce, VariableCenterHarmonicForce

Harmonic oscillator force with a variable stiffness and center that are learnable parameters.

Parameters:
  • kappai (float)

  • kappaf (float)

  • centeri (float)

  • centerf (float)

  • tf (float)

  • steps (int | None)

  • k (List | None)

  • center_list (List | None)

__init__(kappai, kappaf, centeri, centerf, tf, steps=None, k=None, center_list=None, continuous=True)[source]

Initializes the force with a stiffness and center that are defined at steps values of time and linearly interpolated between.

Parameters:
  • kappai (thus equal to) – initial stiffness

  • kappaf (float) – final stiffness

  • centeri (float) – initial center

  • centerf (float) – final center

  • tf (float) – final time of the protocol

  • steps (int, optional) – number time steps where the stiffness value is given

  • k (list, optional) – the initial stiffness given by a list of steps values

  • center_list (list, optional) – the initial center given by a list of steps values

  • continuous (bool) – whether the stiffness and center are continuous at initial and final times,

  • kappai

  • kappaf

  • centerf. (centeri and)

Helper Functions

validate_init_interpolation_list(yi, yf, ylist=None, steps=None, noise=0.25, ylist_name='ylist')[source]

Validate and initialize a list of initial parameters to use for interpolation

Parameters:
  • yi (float) – initial value

  • yf (float) – final value

  • ylist (list, optional) – list of values. Defaults to None.

  • steps (int, optional) – number of steps. Defaults to None.

Raises:

ValueError – If the list of values is not provided and the number of steps is not an integer greater than 1.

Returns:

the list of values to interpolate between. To be used in nn.Parameter.

Return type:

ylist_ini

Example Usage

Here’s an example of how to use the simulator with a harmonic force:

import torch
from langesim_optim.simulator_forces import (
    Simulator,
    VariableStiffnessHarmonicForce
)

# Create a force with variable stiffness
force = VariableStiffnessHarmonicForce(
    kappai=1.0,  # Initial stiffness
    kappaf=2.0,  # Final stiffness
    tf=1.0,      # Final time
    steps=100    # Number of interpolation points
)

# Create a simulator
sim = Simulator(
    dt=0.001,           # Time step
    tot_steps=1000,     # Total simulation steps
    noise_scaler=1.0,   # Noise strength
    force=force         # Force object
)

# Run simulation
x0 = torch.randn(1000)  # Initial positions
xf = sim(x0)           # Final positions