Data loaders

Module for handling datasets and computing prediction errors in population dynamics.

This module provides several dataset classes designed for loading and accessing different formats of the population trajectory data, including trajectory data and coupling data, and testing fit and prediction errors.

Classes

• PopulationDataset

  Handles loading and batching of particle trajectory data. The single unit if a particle trajectory.

• CouplingsDataset

  Loads coupling data for trajectory models, including weights, features, and densities. The single unit is a coupling.

• LinearParametrizationDataset

  Loads data for the linear parametrization. The single unit is the entire dataset.

• PopulationEvalDataset

  Facilitates evaluation of model predictions using particle trajectories and computes prediction errors such as the Wasserstein distance.

class dataset.CouplingsDataset(dataset_name: str)

Bases: Dataset

Dataset class for loading and accessing couplings data.

The dataset is expected to be located in a directory named ‘data/{dataset_name}’ and consist of multiple .npy files. It provides access to input features, target features, time labels, weights, density values, and density gradients.

weight

Array of weights extracted from the couplings data.

Type:

np.ndarray

x

Array of input features extracted from the couplings data.

Type:

np.ndarray

y

Array of target features extracted from the couplings data.

Type:

np.ndarray

time

Array of time labels extracted from the couplings data.

Type:

np.ndarray

densities

Array of density values extracted from the densities files.

Type:

np.ndarray

densities_grads

Array of gradients of densities extracted from the densities files.

Type:

np.ndarray

__getitem__(idx: int) → Tuple[Array, Array, Array, Array, Array, Array]

Retrieve a sample (x, y, t, w, rho, rho_grad) from the dataset at the given index.

Parameters:

idx (int) – The index of the sample to retrieve.

Returns:

A tuple containing:

• Input features (jnp.ndarray): Initial particle distribution.

• Target features (jnp.ndarray): Target particle distribution.

• Time label (jnp.ndarray): Time label.

• Weight of the coupling (jnp.ndarray): Weight of the coupling.

• Density value (jnp.ndarray): Density value.

• Gradient of densities (jnp.ndarray): Gradient of densities.

Return type:

Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]

__init__(dataset_name: str) → None

Initialize the CouplingsDataset by loading data from .npy files.

Parameters:

dataset_name (str) – The name of the dataset to load. The dataset is expected to be located in a directory named ‘data/{dataset_name}’ and consist of multiple .npy files.

__len__() → int

Returns the number of samples in the dataset.

Returns:

The number of samples.

Return type:

int

class dataset.LinearParametrizationDataset(dataset_name: str)

Bases: Dataset

This dataset class loads and organizes data necessary for linear parametrization solver tasks, for which all data is analyzed together.

__getitem__(_) → List[Tuple[ndarray, ndarray, ndarray, ndarray, ndarray, ndarray]]

Retrieve the entire dataset.

Since for the linear parametrization all data is used together, this method returns all data at once and the index parameter _ is ignored.

Parameters:

_ (any) – This parameter is ignored.

Returns:

• List[Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]] – A list of tuples, where each tuple contains:

• - Input features (jnp.ndarray) (Initial particle distribution.)

• - Target features (jnp.ndarray) (Target particle distribution.)

• - Time label (jnp.ndarray) (Time label associated with each sample.)

• - Weight of the coupling (jnp.ndarray) (Weight assigned to the coupling.)

• - Density values (jnp.ndarray) (Density values.)

• - Gradient of densities (jnp.ndarray) (Gradient of the density values.)

__init__(dataset_name: str) → None

Initialize the LinearParametrizationDataset.

Parameters:

dataset_name (str) – The name of the dataset to load.

__len__() → int

Return the number of elements in the dataset.

Returns:

The number of elements (always 1 for this dataset).

Return type:

int

class dataset.PopulationDataset(dataset_name: str, batch_size: int)

Bases: Dataset

Dataset class for loading and accessing particle trajectory data.

The dataset is expected to be located in a directory named ‘data/{dataset_name}’ and consist of a single .npy file named ‘data.npy’. The data contains particle trajectories over time, where each timestep has a set of particles.

If the number of particles in a timestep is less than the maximum number of particles in any timestep, the dataset wraps around to handle the imbalance.

trajectory

Array of shape (num_timesteps, num_particles, num_features) containing the particle trajectories. Each entry in the array represents a particle’s state at a given timestep.

Type:

np.ndarray

__getitem__(idx: int) → list

Retrieve particle data for each timestep at the given index.

Parameters:

idx (int) – The index of the particle to retrieve.

Returns:

A list where each element is an array representing the state of a particle at each timestep. The length of the list corresponds to the number of timesteps, and each array represents the particle state at a specific timestep.

Return type:

list of np.ndarray

__init__(dataset_name: str, batch_size: int) → None

Initialize the PopulationDataset by loading data from ‘data.npy’.

Parameters:

dataset_name (str) – The name of the dataset to load. The dataset should be located in ‘data/{dataset_name}’ and should contain a .npy file named ‘data.npy’.

__len__() → int

Returns the number of timesteps in the dataset.

Returns:

The number of timesteps in the dataset.

Return type:

int

class dataset.PopulationEvalDataset(key, dataset_name: str, solver: str, wasserstein_metric: int, label='test_data')

Bases: Dataset

This dataset class loads and organizes population trajectory data for evaluation.

trajectory

A dictionary where each key corresponds to a unique timestep in the dataset, and the value is an array of trajectory data associated with that timestep.

Type:

dict

label_mapping

A dictionary mapping the original sample labels to consecutive integer indices.

Type:

dict

T

The number of timesteps in the trajectories.

Type:

int

data_dim

The dimensionality of the data at each timestep.

Type:

int

no_ground_truth

Flag indicating if the dataset lacks a ground truth file.

Type:

bool

potential

The potential function used in the predictions.

Type:

str

internal

The internal dynamics setting used.

Type:

str

beta

The beta parameter used in the simulations.

Type:

float

interaction

The interaction function used in the predictions.

Type:

str

dt

The timestep size used in the simulation.

Type:

float

trajectory_only_potential

Trajectory predictions considering only the potential term.

Type:

np.ndarray

trajectory_only_interaction

Trajectory predictions considering only the interaction term.

Type:

np.ndarray

__getitem__(idx: int) → ndarray

Retrieves a particle’s features at the first timestep.

Parameters:

idx (int) – The index of the particle to retrieve.

Returns:

The features of the specified particle at the first timestep.

Return type:

np.ndarray

__init__(key, dataset_name: str, solver: str, wasserstein_metric: int, label='test_data')

Initialize the PopulationEvalDataset.

Parameters:

• key (Any) – A key used for random number generation or seeding.

• dataset_name (str) – The name of the dataset to load. The data should be located in the directory ‘data/{dataset_name}’ and consist of .npy files.

• solver (str) – The solver method used, primarily for plotting or prediction purposes.

• wasserstein_metric (int) – Specifies the order of the Wasserstein distance to be used for the error calculation.

• label (str, optional) – Specifies whether to load ‘test_data’ or ‘train_data’. Default is ‘test_data’.

__len__() → int

Get the number of particles at the first timestep.

Returns:

The number of particles at the first timestep.

Return type:

int

error_wasserstein(trajectory_predicted: ndarray) → float

Compute the Wasserstein loss between the predicted and true trajectories.

This method calculates the Wasserstein distance (a measure of distance between probability distributions) between the predicted trajectories and the true trajectories over all timesteps.

Parameters:

trajectory_predicted (np.ndarray) – The predicted trajectory with shape (T, n_particles, n_features).

Returns:

The cumulative Wasserstein error over all timesteps.

Return type:

float

error_wasserstein_cumulative(predictions: Array, model: str, plot_folder_name: str | None = None) → Array

Compute the cumulative Wasserstein error per timestep.

This method calculates the Wasserstein distance between the predicted and actual trajectories at each timestep and returns the cumulative error.

Parameters:

• predictions (jnp.ndarray) – Array of predicted trajectories with shape (T+1, n_particles, n_features). The predictions should cover the entire timespan from 0 to T.

• model (str) – Name of the solver model used. This is primarily used for plotting purposes.

• plot_folder_name (Optional[str], default=None) – Directory path where plots should be saved. If None, no plots will be saved.

Returns:

Array of cumulative Wasserstein errors, with each entry corresponding to the error at a specific timestep. The array has length T.

Return type:

jnp.ndarray

error_wasserstein_one_step_ahead(potential: Callable[[Array], float], beta: float, interaction: Callable[[Array], float], key_eval: Array, model: str, plot_folder_name: str | None = None) → Array

Compute the Wasserstein error for one-step-ahead predictions.

This method evaluates the prediction error by computing the Wasserstein distance between the predicted trajectory and the actual trajectory at each timestep, given the current true population.

Parameters:

• potential (Callable[[jnp.ndarray], float]) – Function that computes the potential based on a JAX array input.

• beta (float) – Beta parameter used in the predictions.

• interaction (Callable[[jnp.ndarray], float]) – Function that computes the interaction based on a JAX array input.

• key_eval (jnp.ndarray) – Random key for JAX-based random number generation.

• model (str) – Name of the solver model used. This is primarily used for plotting purposes.

• plot_folder_name (Optional[str], default=None) – Directory path where plots should be saved. If None, no plots will be saved.

Returns:

An array of Wasserstein errors for the one-step-ahead predictions over timesteps. The array has length T, where each entry corresponds to the error at a specific timestep.

Return type:

jnp.ndarray