Data generator

Module for generating and processing population trajectory data.

This module provides tools to:

Simulate particle trajectories with different potential, internal and interaction energy configurations, or load pre-existing data.
Fit Gaussian Mixture Models (GMMs) on trajectory data.
Compute couplings between particle distributions at consecutive timesteps.
Generate density and gradient data from the simulated or loaded trajectories.
Plot couplings, density levels, and particle trajectories.

Steps 2-4 are the preprocessing steps required to train a JKOnet* model.

Functions

filename_from_args
Generates a descriptive filename based on the provided command-line arguments.
train_test_split
Splits a dataset into training and testing subsets, ensuring the label distribution is preserved.
generate_data_from_trajectory
Processes trajectory data by fitting a GMM, computing couplings, saving data, and plotting particle densities and couplings.
main
Main entry point for the data generation pipeline. It handles argument parsing, SDE simulation, train-test splitting, and calls functions to generate and save the processed data.

Example

To generate synthetic trajectory data with 1000 particles, a chosen potential, and internal Wiener energy:

python data_generator.py --n-particles 1000 --potential styblinski_tang --n-timesteps 5

To load previously generated data and compute couplings:

python data_generator.py –load-from-file my_trajectory_data –test-ratio 0.2 –n-gmm-components 5

Command-line Arguments

The script accepts the following command-line arguments:

–load-from-file (str):
Load trajectory data from a file instead of generating it. Must be a NumPy array of shape (n_timesteps + 1, n_particles, dimension).
–potential (str):
Specify the potential energy to use for the SDE simulation.
–n-timesteps (int):
Number of timesteps for the SDE simulation.
–dt (float):
Timestep size for the SDE simulation.
–internal (str):
Type of internal energy (e.g., ‘wiener’) to use in the simulation.
–beta (float):
Standard deviation of the Wiener process for internal energy.
–interaction (str):
Specify the interaction energy between particles.
–dimension (int):
Dimensionality of the simulated system.
–n-particles (int):
Number of particles in the simulation.
–batch-size (int):
Batch size for computing couplings during the data processing phase.
–n-gmm-components (int):
Number of components for the Gaussian Mixture Model.
–seed (int):
Random seed for reproducibility.
–test-ratio (float):
Proportion of data to be used as test data during splitting.
–split-population (bool):
If set, data is split at every timestep; otherwise, it is split along the trajectories.
–leave-one-out (int):
If non-negative, leaves one time point out from the training set.
–sinkhorn (float):
Regularization parameter for the Sinkhorn algorithm. If < 1e-12, no regularization is applied.
–dataset-name (str):
Specifies the name of the output dataset. If not provided, a directory name will be automatically generated based on the simulation parameters. This option is only used if data is generated. If data is loaded from a file (using –load-from-file), the output dataset will retain the name of the input file.

data_generator.filename_from_args(args)[source]

Generates a filename based on the arguments given.

Parameters:: args (argparse.Namespace) – Arguments parsed from the command line. See main() for the arguments.
Returns:: Generated filename based on the provided arguments.
Return type:: str

data_generator.generate_data_from_trajectory(folder: str, values: Array, sample_labels: Array, n_gmm_components: int = 10, batch_size: int = 1000, leave_one_out: int = -1, sinkhorn: float = 0.0) → None[source]

Preprocesses the trajectory data for JKOnet*.

Fits Gaussian Mixture Models (GMM) to the trajectory data, computes couplings, and saves the results to disk. This function also plots the data and saves the plots.

Parameters:

folder (str) – Directory where the data and plots will be saved.
values (jnp.ndarray) – Array of trajectory data points.
sample_labels (jnp.ndarray) – Array of sample labels corresponding to each data point.
n_gmm_components (int, optional) – Number of components for the Gaussian Mixture Model (default is 10).
batch_size (int, optional) – Batch size for computing couplings (default is 1000).
leave_one_out (int, optional) – If non-negative, leaves one time point out from the training set (default is -1).
sinkhorn (float, optional) – Regularization parameter for the Sinkhorn algorithm. If < 1e-12, no regularization is applied (default is 0.0).

Return type:

None

data_generator.main(args: Namespace) → None[source]

Main function to run the data generation and processing pipeline.

Parameters:: args (argparse.Namespace) – Command-line arguments (see the module docstring for details).
Return type:: None

data_generator.train_test_split(values: Array, sample_labels: Array, test_ratio: float = 0.4, split_trajectories: bool = True, seed: int = 0) → Tuple[Array, Array, Array, Array][source]

Splits the dataset into training and testing sets while preserving the distribution of labels.

This function ensures that the proportion of each label in the dataset is preserved in both the training and testing subsets.

Parameters:

values (jnp.ndarray) – The data array to be split.
sample_labels (jnp.ndarray) – The corresponding labels for the data. Contains the timestep linked to each value.
test_ratio (float, optional) – The proportion of the dataset to include in the test split. Defaults to 0.4.
split_trajectories (bool, optional) – If True, the data is split by trajectories. Defaults to True. If False, the data is split by individual data points.
seed (int, optional) – Random seed for reproducibility. Defaults to 0.

Returns:

A tuple containing:

Train values: Subset of the data for training.
Train labels: Corresponding labels for the training data.
Test values: Subset of the data for testing.
Test labels: Corresponding labels for the testing data.

Return type:

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