neural_dynamics.models

class FeedForwardNetwork(torch.nn.modules.container.Sequential, neural_dynamics.models.base.NetworkProtocol):

Fully connected neural network with configurable activation functions and Xavier initialisation.

This class provides a standard feedforward architecture that integrates with the parameter system's NetworkArchitecture specification.

The network applies Xavier initialisation in good approximation to all layers and includes configurable activation functions between layers (except the final layer).

Arguments:
  • architecture: Network architecture specification defining layer dimensions
  • activation: Activation function. If None, defaults to nn.Tanh. Should return a new activation instance when called.
  • final_activation: Activation function for the final layer. If None, no activation is applied.
  • device: Computation device. Defaults to globally configured device.
Example:
>>> from neural_dynamics.core.hyperparameters import NetworkArchitecture
>>> architecture = NetworkArchitecture(input_size=3, hidden_sizes=[64, 32], output_size=2)
>>> network = FeedForwardNetwork(architecture, activation=lambda: nn.ReLU())
>>> output = network(torch.randn(10, 3))  # shape: [10, 2]
FeedForwardNetwork( architecture: neural_dynamics.core.hyperparameters.NetworkArchitecture, activation: Callable[[], torch.nn.modules.module.Module] = <class 'torch.nn.modules.activation.Tanh'>, final_activation: Optional[Callable[[], torch.nn.modules.module.Module]] = None, device: torch.device = device(type='cpu'))

Initialise the feedforward network with specified architecture.

Arguments:
  • architecture: Network layer configuration
  • activation: Activation function
  • final_activation: Optional activation for the final layer
  • device: Computation device
architecture: neural_dynamics.core.hyperparameters.NetworkArchitecture
input_size: int
output_size: int
layer_sizes: list[int]
def forward(self, input: torch.Tensor) -> torch.Tensor:

Forward pass through the network.

Arguments:
  • input: Input tensor of shape [batch_size, input_dim]
Returns:

Output tensor of shape [batch_size, output_dim]

Raises:
  • ValueError: If input dimension doesn't match architecture (maybe this will be raised by PyTorch in any case)
class DriftNet(neural_dynamics.models.FeedForwardNetwork):

Neural network approximating the drift term mu(x, t, u) in the neural SDE.

DriftNet( architecture: neural_dynamics.core.hyperparameters.NetworkArchitecture, activation: Callable[[], torch.nn.modules.module.Module] = <class 'torch.nn.modules.activation.Tanh'>, final_activation: Optional[Callable[[], torch.nn.modules.module.Module]] = None, device: torch.device = device(type='cpu'))

Initialise the drift network with specified architecture.

def compute_drift( self, state: torch.Tensor, time: torch.Tensor, external_inputs: Optional[torch.Tensor] = None) -> torch.Tensor:

Compute the drift term for given state, time, and inputs.

class NeuralODE(torch.nn.modules.module.Module):

Neural ODE model backed by a trained drift network.

NeuralODE( *, drift_net: DriftNet, hyperparameters: neural_dynamics.core.hyperparameters.Hyperparameters, time_grid: Optional[torch.Tensor] = None, training_losses: Optional[list[float]] = None, validation_losses: Optional[list[float]] = None)

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

drift_net: DriftNet
hyperparameters: neural_dynamics.core.hyperparameters.Hyperparameters
time_grid: Optional[torch.Tensor]
training_losses: list[float]
validation_losses: list[float]
device: torch.device
def forward(self, initial_state: torch.Tensor, t_span: torch.Tensor) -> torch.Tensor:

Define the computation performed at every call.

Should be overridden by all subclasses.

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.

def drift_function(self, t: float, state: torch.Tensor) -> torch.Tensor:
@classmethod
def train( cls, *, hyperparameters: neural_dynamics.core.hyperparameters.Hyperparameters, trajectories: torch.Tensor, time_grid: torch.Tensor, device: torch.device, validation_split: float = 0.2, early_stopping_patience: Optional[int] = None, wandb_run: Optional[Any] = None) -> NeuralODE:

Train the drift network using batched trajectory supervision.

Arguments:
  • hyperparameters: Training hyperparameters
  • trajectories: Training trajectories [batch, time, state]
  • time_grid: Time grid for trajectories
  • device: Device for training
  • validation_split: Fraction of data to use for validation
  • early_stopping_patience: Early stopping patience (None uses hyperparameter default)
  • wandb_run: Optional W&B run object for logging training progress
Returns:

Trained NeuralODE instance

class DiffusionNet(neural_dynamics.models.FeedForwardNetwork):

Neural network approximating the diffusion term sigma(x, t, u).

DiffusionNet( architecture: neural_dynamics.core.hyperparameters.NetworkArchitecture, *, state_dimension: int, noise_dimension: int, activation: Callable[[], torch.nn.modules.module.Module] = <class 'torch.nn.modules.activation.Tanh'>, device: torch.device = device(type='cpu'))

Initialise the diffusion network with dimensionality checks.

state_dimension
noise_dimension
def compute_diffusion( self, state: torch.Tensor, time: torch.Tensor, external_inputs: Optional[torch.Tensor] = None) -> torch.Tensor:

Compute diffusion matrices for given state, time, and inputs.

class CriticNet(neural_dynamics.models.FeedForwardNetwork):

Neural network that scores trajectories for the WGAN-GP critic.

CriticNet( architecture: neural_dynamics.core.hyperparameters.NetworkArchitecture, trajectory_length: int, *, activation: Callable[[], torch.nn.modules.module.Module] = <class 'torch.nn.modules.activation.Tanh'>, device: torch.device = device(type='cpu'))

Initialise the critic with trajectory-specific validation.

Arguments:
  • architecture: Network architecture describing the critic network.
  • trajectory_length: Number of timesteps in each trajectory segment.
  • activation: Factory returning the activation module used between layers.
  • device: Device where the network parameters are stored.
Raises:
  • ValueError: If the provided configuration is inconsistent.
trajectory_length: int
def forward(self, trajectory_segment: torch.Tensor) -> torch.Tensor:

Forward pass through the network.

Arguments:
  • input: Input tensor of shape [batch_size, input_dim]
Returns:

Output tensor of shape [batch_size, output_dim]

Raises:
  • ValueError: If input dimension doesn't match architecture (maybe this will be raised by PyTorch in any case)
class NeuralSDE(torch.nn.modules.module.Module):

Neural SDE model composed of drift, diffusion, and critic networks.

NeuralSDE( drift_or_hyperparameters: DriftNet | neural_dynamics.core.hyperparameters.Hyperparameters, diffusion_net: DiffusionNet | None = None, hyperparameters: neural_dynamics.core.hyperparameters.Hyperparameters | None = None, *, critic_net: CriticNet | None = None)

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

hyperparameters: neural_dynamics.core.hyperparameters.Hyperparameters
state_dimension: int
input_dimension: int
timestep: float
drift_net: DriftNet
diffusion_net: Optional[DiffusionNet]
noise_dimension: int
critic_net: Optional[CriticNet]
trajectory_length: Optional[int]
generator_losses: list[float]
critic_losses: list[float]
device: torch.device
drift_losses: list[float]
@classmethod
def train( cls, *, hyperparameters: neural_dynamics.core.hyperparameters.Hyperparameters, neural_ode: NeuralODE, trajectories: torch.Tensor, time_grid: torch.Tensor, device: torch.device, enable_adversarial: bool = False, random_seed: int = 0, wandb_run: Optional[Any] = None) -> NeuralSDE:

Train a neural SDE starting from a pretrained neural ODE.

Arguments:
  • hyperparameters: Training hyperparameters
  • neural_ode: Pretrained neural ODE to initialize drift network
  • trajectories: Training trajectories [batch, time, state]
  • time_grid: Time grid for trajectories
  • device: Device for training
  • enable_adversarial: Whether to enable adversarial training
  • random_seed: Random seed for reproducibility
  • wandb_run: Optional W&B run object for logging training progress
Returns:

Trained NeuralSDE instance

def forward(self, *args: Any, **kwargs: Any) -> torch.Tensor:

Define the computation performed at every call.

Should be overridden by all subclasses.

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.

def sample_trajectory( self, initial_state: torch.Tensor, time_grid: torch.Tensor, drift_function: Callable[[float, torch.Tensor], torch.Tensor], diffusion_function: Callable[[float, torch.Tensor], torch.Tensor], device: torch.device) -> torch.Tensor:

Sample a single stochastic trajectory using custom drift/diffusion functions.

This method allows evaluating the SDE with time-varying external inputs by providing custom drift and diffusion functions that capture the full dynamics.

Arguments:
  • initial_state: Initial state [batch_size, state_dim] or [state_dim]
  • time_grid: Time points for integration [num_steps]
  • drift_function: Custom drift function f(t, y) -> drift
  • diffusion_function: Custom diffusion function g(t, y) -> diffusion matrix
  • device: Device for computation
Returns:

Trajectory tensor [batch_size, num_steps, state_dim]

def count_total_parameters(self) -> dict[str, int]:

Return the number of trainable parameters per network component.