Models

models.base module

Module for implementing and comparing different learning diffusion models using JAX.

This module provides an interface for writing code that seamlessly interface with the training and testing procedures.

class models.base.LearningDiffusionModel[source]

Bases: object

An abstract base class representing a learning diffusion model, designed for use with JAX.

This interface serves as a blueprint for implementing new models to learn the different diffusion terms.

abstract create_state(rng: PRNGKey) → Tuple[TrainState, TrainState, TrainState][source]

Abstract method to initialize and return the model’s state. Must be implemented by subclasses.

Parameters:

rng (PRNGKey) – The random key to use for initializing the state.

Returns:

A tuple containing the model’s state for the potential, internal, and interaction terms.

Note: You can use less than three states if your model does not require all three terms, or if it uses a single, shared, state.

Return type:

Tuple[TrainState, TrainState, TrainState]

get_beta(state: Tuple[TrainState, TrainState, TrainState]) → float[source]

Get the beta value from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

The beta value from the internal energy model.

Return type:

float

get_interaction(state: Tuple[TrainState, TrainState, TrainState]) → Callable[[Array], Array][source]

Get the interaction function from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

Function that computes the interaction.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

get_potential(state: Tuple[TrainState, TrainState, TrainState]) → Callable[[Array], Array][source]

Get the potential function from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

Function that computes the potential.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

load_dataset(dataset_name: str) → Dataset[source]

Loads and prepares the specified dataset.

Parameters:

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

Returns:

The loaded dataset.

Return type:

Dataset

abstract train_step(state: Tuple[TrainState, TrainState, TrainState], sample: Tuple[Array, Array, Array, Array, Array, Array]) → Tuple[Array, Tuple[TrainState, TrainState, TrainState]][source]

Perform a single training step.

Parameters:

  • state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

  • sample (Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]) – Training data sample consisting of xs, ys, t, ws, rho, and rho_grad.

Returns:

The loss value and the updated training states.

Return type:

Tuple[jnp.ndarray, Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]]

models.jkonet module

This module contains the implementation of the JKOnet model and of the Monge gap regularization.

Source: https://github.com/bunnech/jkonet

class models.jkonet.JKOnet(config: Dict, data_dim: int, tau: float)[source]

Bases: LearningDiffusionModel

This code ports https://github.com/bunnech/jkonet to our interface. See https://arxiv.org/abs/2106.06345

create_state(rng: PRNGKey) → Any[source]

Abstract method to initialize and return the model’s state. Must be implemented by subclasses.

Parameters:

rng (PRNGKey) – The random key to use for initializing the state.

Returns:

A tuple containing the model’s state for the potential, internal, and interaction terms.

Note: You can use less than three states if your model does not require all three terms, or if it uses a single, shared, state.

Return type:

Tuple[TrainState, TrainState, TrainState]

create_state_from_params(params: Dict) → TrainState[source]
get_beta(_) → float[source]

Get the beta value from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

The beta value from the internal energy model.

Return type:

float

get_interaction(_) → Callable[[Array], Array][source]

Get the interaction function from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

Function that computes the interaction.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

get_potential(state: TrainState) → Callable[[Array], Array][source]

Get the potential function from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

Function that computes the potential.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

load_dataset(dataset_name: str) → PopulationDataset[source]

Loads and prepares the specified dataset.

Parameters:

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

Returns:

The loaded dataset.

Return type:

Dataset

loss_fn_energy(params_energy: Dict, rng_psi: PRNGKey, batch: Array, t: int) → Tuple[float, Tuple[Array, Array]][source]
train_step(state: TrainState, sample: List[List[ndarray]]) → Tuple[Array, TrainState][source]

Perform a single training step.

Parameters:

  • state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

  • sample (Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]) – Training data sample consisting of xs, ys, t, ws, rho, and rho_grad.

Returns:

The loss value and the updated training states.

Return type:

Tuple[jnp.ndarray, Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]]

class models.jkonet.JKOnetMongeGap(config: Dict, data_dim: int, tau: float)[source]

Bases: JKOnetVanilla

A variant of the JKOnet model without the use of ICNN (Input Convex Neural Networks) but instead using the Monge gap regularization.

class models.jkonet.JKOnetVanilla(config: Dict, data_dim: int, tau: float)[source]

Bases: JKOnet

A variant of the JKOnet model without the use of ICNN (Input Convex Neural Networks).

models.jkonet.get_optimize_psi_fn(loss_fn_psi: Callable[[Array, Array, Array], Tuple[Array, Array]], optimizer_psi: GradientTransformation, n_iter: int = 100, min_iter: int = 50, max_iter: int = 200, inner_iter: int = 10, threshold: float = 1e-05, fploop: bool = False) → Callable[source]

models.jkonet_star module

Module that implements the JKOnet* model based on the base interface.

The models are implemented using JAX and the FLAX library, following a functional paradigm to support efficient differentiation and optimization. The core classes include:

  • JKOnetStar: The full JKOnet* method, used for learning all the energy terms.

  • JKOnetStarPotentialInternal: A variant focusing on potential and internal energies.

  • JKOnetStarPotential: A variant focusing solely on the potential energy term.

  • JKOnetStarTimePotential: A time-extended variant of JKOnetStarPotential.

  • JKOnetStarLinear: A model using linear parametrizations with various feature functions.

class models.jkonet_star.JKOnetStar(config: dict, data_dim: int, tau: float)[source]

Bases: LearningDiffusionModel

The full JKOnet* model for learning all energy terms.

__init__(config: dict, data_dim: int, tau: float) → None[source]

Initialize the JKOnetStar model.

Parameters:

  • config (dict) – Configuration dictionary containing model and optimizer settings.

  • data_dim (int) – Dimension of the input data.

  • tau (float) – Represents the time scale over which the diffusion process described by the Fokker-Planck equation is considered.

create_state(rng: PRNGKey) → Tuple[TrainState, TrainState, TrainState][source]

Create initial training states for the potential, internal, and interaction models.

Parameters:

rng (jax.random.PRNGKey) – Random key for initialization.

Returns:

Tuple containing the training states for the potential, internal, and interaction models.

Return type:

Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]

create_state_from_params(potential_params: dict, internal_params: dict, interaction_params: dict) → Tuple[TrainState, TrainState, TrainState][source]

Create training states from the provided parameters.

Parameters:

  • potential_params (dict) – Parameters for the potential model.

  • internal_params (dict) – Parameters for the internal model.

  • interaction_params (dict) – Parameters for the interaction model.

Returns:

Tuple containing the training states for the potential, internal, and interaction models.

Return type:

Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]

get_beta(state: Tuple[TrainState, TrainState, TrainState]) → float[source]

Get the beta value from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

The beta value from the internal energy model.

Return type:

float

get_interaction(state: Tuple[TrainState, TrainState, TrainState]) → Callable[[Array], Array][source]

Get the interaction function from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

Function that computes the interaction.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

get_params(state: Tuple[TrainState, TrainState, TrainState]) → Tuple[FrozenDict[str, Any], FrozenDict[str, Any], FrozenDict[str, Any]][source]

Get parameters from the training state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

Tuple containing the parameters for the potential, internal, and interaction models.

Return type:

Tuple[dict, dict, dict]

get_potential(state: Tuple[TrainState, TrainState, TrainState]) → Callable[[Array], Array][source]

Get the potential function from the model state.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

Returns:

Function that computes the potential.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

load_dataset(dataset_name: str) → CouplingsDataset[source]

Load and return a dataset based on the given dataset name.

This method creates an instance of the CouplingsDataset class using the specified dataset name.

Parameters:

dataset_name (str) – The name of the dataset to load. This name is used to locate and initialize the dataset.

Returns:

An instance of the CouplingsDataset class, which contains the loaded dataset.

Return type:

CouplingsDataset

loss(potential_params: dict, internal_params: dict, interaction_params: dict, xs: Array, ys: Array, ws: Array, rho: Array, rho_grad: Array) → Array[source]

Compute the total loss for the model by combining potential, internal, and interaction terms.

Parameters:

  • potential_params (dict) – Parameters for the potential model.

  • internal_params (dict) – Parameters for the internal model.

  • interaction_params (dict) – Parameters for the interaction model.

  • xs (jnp.ndarray) – Initial particle distribution.

  • ys (jnp.ndarray) – Target particle distribution.

  • ws (jnp.ndarray) – Weights of the couplings.

  • rho (jnp.ndarray) – Density values.

  • rho_grad (jnp.ndarray) – Gradient of density values.

Returns:

Total loss value.

Return type:

jnp.ndarray

train_step(state: Tuple[TrainState, TrainState, TrainState], sample: Tuple[Array, Array, Array, Array, Array, Array]) → Tuple[Array, Tuple[TrainState, TrainState, TrainState]][source]

Perform a single training step.

Parameters:

  • state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – Training state containing potential, internal, and interaction models.

  • sample (Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]) – Training data sample consisting of xs, ys, t, ws, rho, and rho_grad.

Returns:

The loss value and the updated training states.

Return type:

Tuple[jnp.ndarray, Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]]

class models.jkonet_star.JKOnetStarLinear(config: Dict[str, Dict[str, int | float | bool] | float], data_dim: int, tau: float)[source]

Bases: LearningDiffusionModel

The linear parametrization of the JKOnet* model.

__init__(config: Dict[str, Dict[str, int | float | bool] | float], data_dim: int, tau: float) → None[source]

Initializes the JKOnetStarLinear model with configuration and data dimensions.

Parameters:

  • config (dict) – Configuration dictionary specifying model parameters and feature settings.

  • data_dim (int) – Dimensionality of the input data.

  • tau (float) – Represents the time scale over which the diffusion process described by the Fokker-Planck equation is considered.

create_state(_) → Tuple[Array, Array, Array][source]

Creates the initial state for the model.

This method returns a tuple of zero-initialized arrays corresponding to the potential, interaction, and internal parameters.

Parameters:

_ (Ignored) – Placeholder for compatibility, not used.

Returns:

The initial state consisting of zero-initialized potential, interaction, and internal parameters.

Return type:

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

create_state_from_params(potential_params: Array, interaction_params: Array, internal_params: Array) → Tuple[Array, Array, Array][source]

Creates the state from given parameters.

Parameters:

  • potential_params (jnp.ndarray) – The parameters for the potential term.

  • interaction_params (jnp.ndarray) – The parameters for the interaction term.

  • internal_params (jnp.ndarray) – The parameters for the internal term.

Returns:

A tuple containing the potential, interaction, and internal parameters.

Return type:

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

features(x: Array) → Array[source]

Computes the feature functions for the input data x.

Parameters:

x (jnp.ndarray) – Input data for which to compute feature functions.

Returns:

The computed feature functions for the input data.

Return type:

jnp.ndarray

property features_dim: int

The dimension of the feature space.

Computes and caches the dimension of the feature space based on the feature functions.

Returns:

The dimension of the feature space.

Return type:

int

features_grad(x: Array) → Array[source]

Computes the gradients of the feature functions with respect to x.

Parameters:

x (jnp.ndarray) – Input data for which to compute gradients of feature functions.

Returns:

The gradients of the feature functions with respect to the input data.

Return type:

jnp.ndarray

get_beta(state: Tuple[Array, Array, Array]) → float[source]

Returns the value of the internal parameter (beta) based on the current state.

Parameters:

state (Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]) – The current state containing the potential, interaction, and internal parameters.

Returns:

The beta value from the internal energy model.

Return type:

float

get_interaction(state: Tuple[Array, Array, Array]) → Callable[[Array], Array][source]

Returns the interaction function based on the current state.

Parameters:

state (Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]) – The current state containing the potential, interaction, and internal parameters.

Returns:

A function that computes the interaction for a given input x.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

get_params(state: Tuple[Array, Array, Array]) → Tuple[Array, Array, Array][source]

Retrieves the parameters from the given state.

Parameters:

state (Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]) – The current state containing the potential, interaction, and internal parameters.

Returns:

The potential, interaction, and internal parameters extracted from the state.

Return type:

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

get_potential(state: Tuple[Array, Array, Array]) → Callable[[Array], Array][source]

Returns the potential function based on the current state.

Parameters:

state (Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]) – The current state containing the potential, interaction, and internal parameters.

Returns:

A function that computes the potential for a given input x.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

load_dataset(dataset_name: str) → LinearParametrizationDataset[source]

Loads and returns the dataset for linear parametrizations.

Parameters:

dataset_name (str) – Name of the dataset to load.

Returns:

The dataset object for linear parametrizations.

Return type:

LinearParametrizationDataset

train_step(_, all_samples: Tuple) → Tuple[Array, Tuple[Array, Array, Array]][source]

Performs a single training step, updating the model’s parameters.

This method solves the least squares problem to update the model parameters based on the provided samples.

Parameters:

  • _ (Ignored) – Placeholder for compatibility, not used.

  • all_samples (Tuple) – A tuple of all the samples, where each sample contains (xs, ys, t, ws, rho, rho_grad).

Returns:

The error and the updated parameters (potential, interaction, internal).

Return type:

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

class models.jkonet_star.JKOnetStarPotential(config: dict, data_dim: int, tau: float)[source]

Bases: JKOnetStarPotentialInternal

A variant of the JKOnetStar model to learn only the potential term.

get_beta(state: Tuple[TrainState, TrainState, TrainState]) → float[source]

Return a constant zero value for the beta term of the internal energy model.

Parameters:

state (Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]) – The current states of the potential, internal, and interaction models.

Returns:

The constant zero value for the beta term of the internal energy model.

Return type:

float

loss(potential_params: dict, xs: Array, ys: Array, ws: Array, rho: Array, rho_grad: Array) → Array[source]

Compute the loss for the potential term of the model.

Parameters:

  • potential_params (dict) – Parameters of the potential model.

  • xs (jnp.ndarray) – Initial particle distribution.

  • ys (jnp.ndarray) – Target particle distribution.

  • ws (jnp.ndarray) – Weights of the couplings.

  • rho (jnp.ndarray) – Density values for the data samples.

  • rho_grad (jnp.ndarray) – Gradient of the density values.

Returns:

The computed loss value.

Return type:

jnp.ndarray

class models.jkonet_star.JKOnetStarPotentialInternal(config: dict, data_dim: int, tau: float)[source]

Bases: JKOnetStar

A specialized variant of the JKOnetStar model that only considers potential and internal terms.

get_interaction(_) → Callable[[Array], Array][source]

Returns a function representing the interaction term.

This implementation returns a constant zero function, as the interaction is not used in this variant.

Parameters:

_ (Any) – Unused parameter in this context.

Returns:

A function that always returns 0.

Return type:

Callable[[jnp.ndarray], jnp.ndarray]

loss(potential_params: dict, internal_params: dict, xs: Array, ys: Array, ws: Array, rho: Array, rho_grad: Array) → Array[source]

Compute the loss for the potential and internal terms of the model.

Parameters:

  • potential_params (dict) – Parameters of the potential model.

  • internal_params (dict) – Parameters of the internal model.

  • xs (jnp.ndarray) – Initial particle distribution.

  • ys (jnp.ndarray) – Target particle distribution.

  • ws (jnp.ndarray) – Weights of the couplings.

  • rho (jnp.ndarray) – Density values for the data samples.

  • rho_grad (jnp.ndarray) – Gradient of the density values.

Returns:

The computed loss value.

Return type:

jnp.ndarray

class models.jkonet_star.JKOnetStarTimePotential(config: dict, data_dim: int, tau: float)[source]

Bases: JKOnetStarPotential

A variant of the JKOnetStarPotential model that incorporates time information in the potential term.

create_state(rng: PRNGKey) → Tuple[TrainState, TrainState, TrainState][source]

Creates initial training states for the potential, internal, and interaction models.

Parameters:

rng (jax.random.PRNGKey) – Random key for JAX-based random number generation.

Returns:

The initial states for the potential, internal, and interaction models.

Return type:

Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]

loss(potential_params: dict, xs: Array, ys: Array, t: Array, ws: Array, rho: Array, rho_grad: Array) → Array[source]

Computes the total loss for the model, considering the potential term with time information.

Parameters:

  • potential_params (dict) – Parameters of the potential model.

  • xs (jnp.ndarray) – Initial particle distribution.

  • ys (jnp.ndarray) – Target particle distribution.

  • t (jnp.ndarray) – Timestep of the target particle distribution.

  • ws (jnp.ndarray) – Weights of the couplings.

  • rho (jnp.ndarray) – Density values for the data samples.

  • rho_grad (jnp.ndarray) – Gradient of the density values.

Returns:

The computed loss value.

Return type:

jnp.ndarray