examples.systems.duffing_oscillator

Stochastic Duffing Oscillator

This module implements the stochastic Duffing oscillator as a system of first-order stochastic differential equations (SDEs) with both white and coloured noise.

The dimensionless SDE system is:

dq(t) = v(t) dt dv(t) = (-alpha q(t) - beta q(t)^3 - mu v(t) + gamma cos(Omega t) + xi(t)) dt + sqrt(2 mu Theta) \circ dW(t) dxi(t) = -1/t_correl_tilde * xi(t) dt + sqrt(2 D_tilde / t_correl_tilde) \circ dW(t)

where the following are dimensionless parameters: q(t): position v(t): velocity xi(t): auxiliary variable representing Ornstein-Uhlenbeck coloured noise alpha, beta: linear and nonlinear stiffness parameters mu: damping coefficient gamma: forcing amplitude Omega: forcing frequency Theta: temperature (white noise) D_tilde: coloured noise intensity t_correl_tilde: correlation time

class DuffingOscillator:

Stochastic Duffing oscillator implementation using PyTorch.

Implements the dimensionless SDE system with state vector [q, v, xi] where:

  • q(t) is dimensionless position
  • v(t) is dimensionless velocity
  • xi(t) is the Ornstein-Uhlenbeck process for coloured noise
DuffingOscillator( params: examples.systems.parameters.duffing_oscillator.DuffingOscillatorParameters)

Initialise the stochastic Duffing oscillator.

Arguments:
  • params: Dimensionless system parameters
params
num_steps: int
time_grid: torch.Tensor
has_white_noise
has_coloured_noise
has_noise
def get_state_dimension(self) -> int:

Get the dimension of the state vector.

Returns:

State dimension (always 3 for [q, v, xi])

def get_initial_state(self, batch_size: int = 1) -> torch.Tensor:

Get the initial state vector [q_0, v_0, xi_0], batched for efficiency.

Arguments:
  • batch_size: The number of initial states to generate.
Returns:

A batch of initial state tensors of shape [batch_size, 3]

def external_forcing(self, t: float) -> torch.Tensor:

External periodic forcing function gamma cos(Omega t).

Arguments:
  • t: Dimensionless time (scalar)
Returns:

Forcing value at time t

def drift_function(self, t: float, state: torch.Tensor) -> torch.Tensor:

Compute the drift term for the SDE system. This function is vectorized to handle batches of states efficiently.

Implements the drift vector field:

[dq/dt, dv/dt, dxi/dt] = [v, -alpha q - beta q^3 - mu v + gamma cos(Omega t) + xi, -xi/t_correl_tilde]

Arguments:
  • t: Current dimensionless time
  • state: Current state [q, v, xi] of shape [batch_size, 3] or [3]
Returns:

Drift term vector of the same shape as state.

def diffusion_function(self, t: float, state: torch.Tensor) -> torch.Tensor:

Compute the diffusion term for the SDE system. This function is vectorized to handle batches of states efficiently.

Implements the diffusion matrix:

[0, sqrt(2 mu Theta), sqrt(2 D_tilde / t_correl_tilde)]

Arguments:
  • t: Current dimensionless time
  • state: Current state [q, v, xi] of shape [batch_size, 3] or [3]
Returns:

Diffusion term vector of shape [batch_size, 3] or [3]

def integrate_sde( self, initial_state: Optional[torch.Tensor] = None, batch_size: int = 1) -> Tuple[torch.Tensor, torch.Tensor]:

Integrate the stochastic Duffing oscillator SDE system for a batch of trajectories.

This method was refactored to generate multiple trajectories in a batched manner for significantly improved performance, as the underlying drift and diffusion functions are vectorized. Generating trajectories one-by-one in a Python loop is inefficient.

Arguments:
  • initial_state: Initial conditions of shape [batch_size, 3]. If None, uses default parameters.
  • batch_size: The number of trajectories to generate. Ignored if initial_state is provided.
Returns:

Tuple of (time_grid, state_trajectory) state_trajectory has shape [batch_size, num_steps+1, 3]

def get_position_velocity(self, trajectory: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

Extract position and velocity from state trajectory.

Arguments:
  • trajectory: State trajectory from integrate_sde of shape [batch_size, num_steps+1, 3]
Returns:

Tuple of (position, velocity) tensors

def get_coloured_noise(self, trajectory: torch.Tensor) -> torch.Tensor:

Extract coloured noise component from state trajectory.

Arguments:
  • trajectory: State trajectory from integrate_sde
Returns:

Coloured noise time series xi(t)

def potential_energy(self, q: torch.Tensor) -> torch.Tensor:

Compute the Duffing potential energy.

V(q) = alpha q^2/2 + beta q^4/4

Arguments:
  • q: Dimensionless position
Returns:

Potential energy V(q)

def kinetic_energy(self, v: torch.Tensor) -> torch.Tensor:

Compute kinetic energy.

T = v^2/2

Arguments:
  • v: Dimensionless velocity
Returns:

Kinetic energy T

def total_energy(self, q: torch.Tensor, v: torch.Tensor) -> torch.Tensor:

Compute total mechanical energy.

E = T + V = v^2/2 + alpha q^2/2 + beta q^4/4

Arguments:
  • q: Dimensionless position
  • v: Dimensionless velocity
Returns:

Total energy E = T + V

def effective_potential(self, q: torch.Tensor, xi: torch.Tensor) -> torch.Tensor:

Compute effective potential including coloured noise contribution.

This is useful for visualising how coloured noise modifies the potential landscape.

Arguments:
  • q: Dimensionless position
  • xi: Coloured noise value (Ornstein-Uhlenbeck process)
Returns:

Effective potential V_eff(q, xi) = V(q) - xi·q

@dataclass(frozen=True)
class DuffingDataConfig:

Configuration for generating Duffing training data.

DuffingDataConfig( seed: int = 42, n_trajectories: int = 4096, initial_condition_scale: float = 0.5, noise_disabled: bool = False, data_dir: pathlib._local.Path = <factory>)
seed: int = 42
n_trajectories: int = 4096
initial_condition_scale: float = 0.5
noise_disabled: bool = False
data_dir: pathlib._local.Path
def cache_path( self, *, physical_params: examples.systems.parameters.duffing_oscillator.PhysicalDuffingParameters) -> pathlib._local.Path:

Return cache path with key parameters directly in filename.

def generate_duffing_training_data( *, system: DuffingOscillator, n_trajectories: int, initial_condition_scale: float, device: torch.device = device(type='cpu')) -> neural_dynamics.core.TrajectoryDataset:

Generate batched Duffing trajectories for training.

def load_duffing_training_data( path: pathlib._local.Path, *, device: torch.device = device(type='cpu')) -> neural_dynamics.core.TrajectoryDataset:

Load cached Duffing trajectories from disk.

def save_duffing_training_data( data: neural_dynamics.core.TrajectoryDataset, path: pathlib._local.Path) -> None:

Persist Duffing training data to a torch archive.

def prepare_duffing_training_data( *, config: DuffingDataConfig, physical_params: Optional[examples.systems.parameters.duffing_oscillator.PhysicalDuffingParameters] = None, regenerate: bool = False, device: torch.device = device(type='cpu')) -> tuple[pathlib._local.Path, neural_dynamics.core.TrajectoryDataset, DuffingOscillator, examples.systems.parameters.duffing_oscillator.PhysicalDuffingParameters]:

Generate or load Duffing training data and return associated artefacts.