API

Manifold

The core functionality is found in the trajectory_manifold.manifold module.

Computing geometric transformations onto the trajectory manifold.

This module contains the core functions requires for geometric transformations of quantities onto the trajectory manifold through the use of an ODE solver and a known vector field.

class trajectory_manifold.manifold.SolverParameters(NamedTuple)

Stores Information for ODE Solvers.

Records the parameters for solving an ODE using Diffrax, including the solver, tolerances, output grid size, and time horizon

Parameters:

  • relative_tolerance – Relative tolerance for the ODE solution

  • absolute_tolerance – Absolute tolerance for the ODE solution

  • step_size – Output mesh size. Note: Does not impact internal computations.

  • time_horizon – Length of the solution in seconds.

  • solver – The particular ODE solver to use.

trajectory_manifold.manifold.system_pushforward_weight(vector_field: Callable[[Float, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], time_interval: tuple[float, float], initial_condition: Float[Array, 'dim'], system_parameters: Float[Array, 'dim2']) → Float

Computes the pushforward weight for a given initial condition.

Given a differential equation, initial condition, and desired time horizon, computes the weight required to push densities onto the Riemannian manifold of valid trajectories of the system.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • time_interval – Time interval for the trajectory manifold in the form (initial time, final time).

  • initial_condition – The position in the statespace to be pushed onto the manifold.

  • system_parameters – The parameters for the vector field.

Returns:

The weight required to push a density onto the trajectory manifold.

trajectory_manifold.manifold.system_pushforward_weight_reweighted(vector_field: Callable[[any, Float[Array, 'dim'], any], Float[Array, 'dim']], time_interval: Float, initial_condition: Float[Array, 'dim'], system_parameters: Float[Array, 'dim2'], step_size: Float, kernel: Float[Array, 'timesteps dim dim']) → Float

Computes the pushforward weight for a given initial condition.

Given a differential equation, initial condition, and desired time horizon, computes the weight required to push densities onto the Riemannian manifold of valid trajectories of the system.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • time_interval – Time interval for the trajectory manifold.

  • initial_condition – The position in the statespace to be pushed onto the manifold.

  • system_parameters – The parameters for the vector field.

  • step_size – The step size for the numerical solution of the ODE.

  • kernel – An N by K by K array of N timesteps of an integral kernel to apply to a K dimensional space.

Returns:

The weight required to push a density onto the trajectory manifold.

trajectory_manifold.manifold.system_sensitivity(vector_field: Callable[[Float, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], initial_condition: Float[Array, 'dim'], system_parameters: PyTree, solver_parameters: SolverParameters) → Float[Array, 'dim timesteps dim']

Computes the differential equation sensitivity to the initial conditions.

Given a differential equation, initial condition, and desired time horizon, computes the Jacobian of the transformation from the initial condition to the Riemannian manifold of valid trajectories. The Jacobian is expressed in the ambient space of square integrable functions.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • initial_condition – The position in the statespace to be pushed onto the manifold.

  • system_parameters – The parameters for the vector field.

  • solver_parameters – The set of parameters for the ODE solver.

Returns:

A matrix where each row represents the sensitivity of the system solution to a perturbation along some element of an orthonormal basis of the state space.

Estimation

The module trajectory_manifold.estimation contains functions related to probability.

This module includes functions for statistical estimation

trajectory_manifold.estimation.assign_center(point: Float[Array, 'dim'], centers: Float[Array, 'k dim']) → int

Returns the center for kmeans clustering.

Parameters:

  • point – The vector being assigned.

  • centers – Set of centers to compare.

Returns:

The index of the nearest center to the given point.

trajectory_manifold.estimation.compute_bounding_box(points: Float[Array, 'sample ...'])

Computes a bounding box around a set of points

trajectory_manifold.estimation.compute_centers(points: Float[Array, 'sample dim'], labels: Int[Array, 'sample'], old_centers: Float[Array, 'k dim']) → Float[Array, 'k dim']

Given a labeled dataset, compute the centers.

Parameters:

  • points – A set of sample points.

  • labels – labels for the sample points.

  • old_centers – An array containing the previous centers.

Returns:

An array of averages of the labeled points.

Notes

old_centers is only used for the shape due to the requirement in Jax for type concretization.

trajectory_manifold.estimation.credible_trajectories(state_and_param, vector_field, posterior, parameters) → Float[Array, 'dim']

Computes the set of credible trajectories

Parameters:

  • state_and_param –

  • vector_field –

  • posterior –

  • parameters –

Returns:

The list of credible trajectories in the set

trajectory_manifold.estimation.kmeans(points: Float[Array, 'point dim'], initial_centers: Float[Array, 'k dim'], max_steps) → tuple[Float[Array, 'k dim'], Int[Array, 'point']]

Performs K-means clustering on an array.

Parameters:

  • points – An N by D array of N datapoints

  • initial_centers – A K by D array of K initial centers

  • max_steps – Maximum number of steps for k-means

Returns:

(centers, labels) where centers is a K by D array containing the means in k-means and labels is an N dimensional array mapping points onto their nearest center.

trajectory_manifold.estimation.quantized_credible_regions(points: Float[Array, 'points ...'], region_count: int, max_steps: int) → Float[Array, 'dim1 2 region_count']

Computes quantized credible regions based on a set of credible points.

Applies k-means clustering to quantize the space of points. Computes the bounding boxes of clusters representing the quantized credible regions.

Parameters:

  • points – A set of points

  • region_count – Number of credible regions

Returns:

A set of bounding boxes of the credible regions.

trajectory_manifold.estimation.state_log_posterior(vector_field: Callable[[any, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], observations: Float[Array, 'timesteps dim2'], observation_times: Float[Array, 'timesteps'], observation_likelihood: Callable[[Float[Array, 'dim2'], Float[Array, 'dim']], float], log_prior: Callable[[Float[Array, 'dim'], PyTree], float], parameters: SolverParameters) → Callable[[Float[Array, 'dim'], PyTree], float]

Constructs a posterior distribution for the initial state of the system.

Given a representation of a differential equation, a set of observation times and an observation likelihood function dependent on the state, construct a posterior distribution jointly over all observations.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • observations – An N by K dimensional array of observations.

  • observation_times – An N dimensional array of observation times.

  • observation_likelihood – A function mapping pairs of states and observations to the likelihood.

  • log_prior – A function representing the prior distribution of the initial state and parameters of the system.

  • parameters – Parameters for the ODE solver

Returns:

A function mapping the state and parameters at the initial observation time to the log posterior of the trajectory.

trajectory_manifold.estimation.state_posterior(vector_field: Callable[[any, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], observations: Float[Array, 'timesteps dim2'], observation_times: Float[Array, 'timesteps'], observation_likelihood: Callable[[Float[Array, 'dim2'], Float[Array, 'dim']], float], prior: Callable[[Float[Array, 'dim'], PyTree], float], parameters: SolverParameters) → Callable[[Float[Array, 'dim'], PyTree], Float]

Constructs a posterior distribution for the initial state of the system.

Given a representation of a differential equation, a set of observation times and an observation likelihood function dependent on the state, construct a posterior distribution jointly over all observations.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • observations – An N by K dimensional array of observations.

  • observation_times – An N dimensional array of observation times.

  • observation_likelihood – A function mapping pairs of states and observations to the likelihood.

  • prior – A function representing the prior distribution of the initial state and parameters of the system.

  • parameters – Parameters for the ODE solver

Returns:

A function mapping the state and initial condition at the initial observation time to the posterior of the initial condition and parameters.

trajectory_manifold.estimation.trajectory_likelihood(vector_field: Callable[[any, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], observations: Float[Array, 'timesteps dim2'], observation_times: Float[Array, 'timesteps'], observation_likelihood: Callable[[Float[Array, 'dim2'], Float[Array, 'dim']], float], parameters: SolverParameters) → Callable[[Float[Array, 'dim'], PyTree], Float]

Constructs a likelihood function for a set of observations of a system.

Given a representation of a differential equation, a set of observation times and an observation likelihood function dependent on the state, construct a likelihood function jointly over all observations.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • observations – An N by K dimensional array of observations.

  • observation_times – An N dimensional array of observation times.

  • observation_likelihood – A function mapping pairs of states and observations to the likelihood.

  • parameters – Parameters for the ODE solver

Returns:

A function mapping the state and parameters at the initial observation time to the likelihood of the observation.

trajectory_manifold.estimation.trajectory_log_likelihood(vector_field: Callable[[any, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], observations: Float[Array, 'timesteps dim2'], observation_times: Float[Array, 'timesteps'], observation_log_likelihood: Callable[[Float[Array, 'dim2'], Float[Array, 'dim']], float], parameters: SolverParameters) → Callable[[Float[Array, 'dim'], PyTree], Float]

Constructs a likelihood function for a set of observations of a system.

Given a representation of a differential equation, a set of observation times and an observation likelihood function dependent on the state, construct a likelihood function jointly over all observations.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • observations – An N by K dimensional array of observations.

  • observation_times – An N dimensional array of observation times.

  • observation_likelihood – A function mapping pairs of states and observations to the likelihood.

  • parameters – Parameters for the ODE solver

Returns:

A function mapping the state and parameters at the initial observation time to the log likelihood of the observation.

trajectory_manifold.estimation.trajectory_log_posterior(vector_field: Callable[[any, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], observations: Float[Array, 'timesteps dim2'], observation_times: Float[Array, 'timesteps'], observation_likelihood: Callable[[Float[Array, 'dim2'], Float[Array, 'dim']], float], log_prior: Callable[[Float[Array, 'dim'], PyTree], float], time_interval: tuple[float, float], parameters: SolverParameters) → Callable[[Float[Array, 'dim'], PyTree], Float]

Constructs a posterior distribution for system trajectories.

Given a representation of a differential equation, a set of observation times and an observation likelihood function dependent on the state, construct a posterior distribution jointly over all observations.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • observations – An N by K dimensional array of observations.

  • observation_times – An N dimensional array of observation times.

  • observation_likelihood – A function mapping pairs of states and observations to the likelihood.

  • prior – A function representing the prior distribution of the initial state and parameters of the system.

  • time_interval – Time interval for the trajectory manifold in the form (initial time, final time).

  • parameters – Parameters for the ODE solver

Returns:

A function mapping the state and system parameters at the initial observation time to the posterior distribution of the trajectory.

trajectory_manifold.estimation.trajectory_posterior(vector_field: Callable[[any, Float[Array, 'dim'], PyTree], Float[Array, 'dim']], observations: Float[Array, 'timesteps dim2'], observation_times: Float[Array, 'timesteps'], observation_likelihood: Callable[[Float[Array, 'dim2'], Float[Array, 'dim']], float], prior: Callable[[Float[Array, 'dim'], PyTree], float], time_interval: tuple[float, float], parameters: SolverParameters) → Callable[[Float[Array, 'dim'], PyTree], Float]

Constructs a posterior distribution for system trajectories.

Given a representation of a differential equation, a set of observation times and an observation likelihood function dependent on the state, construct a posterior distribution jointly over all observations.

Parameters:

  • vector_field – Governing differential equation mapping the current state to the derivative.

  • observations – An N by K dimensional array of observations.

  • observation_times – An N dimensional array of observation times.

  • observation_likelihood – A function mapping pairs of states and observations to the likelihood.

  • prior – A function representing the prior distribution of the initial state and parameters of the system.

  • time_interval – Time interval for the trajectory manifold in the form (initial time, final time).

  • parameters – Parameters for the ODE solver

Returns:

A function mapping the state at the initial observation time to the posterior distribution of the observation.

Optimization

The module trajectory_manifold.optimize contains tools for optimization on the manifold of feasible trajectories. The main approach is the computation of pullbacks of gradients into the state space.

This module contains functions useful for optimization.

This module contains useful functions for optimization on the trajectory manifold. It enables the computation of the pullback of gradients into the ambient space, as well as helper functions for grid-based methods.

trajectory_manifold.optimize.distance_gradient(initial_condition: Float[Array, 'dim'], system_params: Float[Array, 'dim2'], vector_field: Callable[[Float, Float[Array, 'dim'], Float[Array, 'dim2']], Float[Array, 'dim']], trajectory: Float[Array, 'dim dim3'], solver_params: SolverParameters) → tuple[Float[Array, 'dim'], Float[Array, 'dim2']]

Computes the gradient of the squared distance to a chosen trajectory.

Computes the pullback along the transformation from initial conditions to system trajectories of the gradient of the squared distance from a given sampled trajectory in the ambient space.

The gradient in the trajectory manifold is the projection of the gradient in the ambient space onto the tangent space of the manifold.

Parameters:

  • initial_condition – The initial condition around which to linearize.

  • vector_field – Vector field defining the differential equation.

  • trajectory – A sampled function from which the distance is computed.

  • solver_params – Parameters for the ODE solvers.

Returns:

The gradient of the distance to a function

trajectory_manifold.optimize.grid_optimum(grid: Float[Array, 'dim1 dim2'], x: Float[Array, 'dim1'], y: Float[Array, 'dim2'], upper: bool = True, reverse: bool = False) → Float[Array, '2']

Compute the optimal value in a grid.

Given a matrix of values representing some objective function on a grid, compute the point with the maximum or minimum value.

Parameters:

  • grid – Sampled function evaluations

  • x – Values in the first index

  • y – Values for the second index

  • upper – Maximize if True, Minimize if False

  • reverse – Reverse x and y coordinates. Common due to transposes.

Returns:

An array containing the x, y coordininate of the optimal grid point.

trajectory_manifold.optimize.zero_order_gradient_estimate(f: Callable[[Float[Array, 'dim batch_size']], Float[Array, 'batch_size']], x: Float[Array, 'dim'], smoothing: float, batch_size: int, key: PRNGKeyArray) → Float[Array, 'dim']

Construct a zero-order estimate of the gradient of a function.

Constructs an approximation of the gradient of a given function through a Monte Carlo method. Generates batch_size samples of Gaussian perturbations from the chosen point, and averages the estimates constructed by finite-difference approximations. Includes a smoothing factor which adjusts the variance of the Gaussian samples.

Parameters:

  • f – A function of which the gradient is to be computed.

  • x – The point at which the gradient is computed

  • smoothing – Smoothing factor in the gradient computation, determines the ball size for the random samples

  • batch_size – The number of random points used in the gradient estimate.

  • key – PRNG key for Jax RNG functioanlity

Returns:

An approximation of the gradient of f at x

Notes

For more information, see

S. Liu, P. -Y. Chen, B. Kailkhura, G. Zhang, A. O. Hero III and P. K. Varshney, “A Primer on Zeroth-Order Optimization in Signal Processing and Machine Learning: Principals, Recent Advances, and Applications,” in IEEE Signal Processing Magazine, vol. 37, no. 5, pp. 43-54, Sept. 2020, doi: 10.1109/MSP.2020.3003837.

Helpers

The module trajectory_manifold.helpers contains a collection of helper functions for the small modifications to linear algebra operations required in the project.

This module contains linear algebra helper functions.

The purpose of this module is to adapt standard linear algebra fuctions to the needs of this project. These include items such as computation of an inner product through the trapezoidal rule.

trajectory_manifold.helpers.apply_kernel(x: Float[Array, 'timesteps dim'], kernel: Float[Array, 'timesteps dim dim']) → Float[Array, 'timesteps dim']

Helper Function to apply an integration kernel to an input.

Applies the kernel to an input at each timestep.

Parameters:

  • x – An N by K array of N samples of each of a function taking values in a K-dimensional space.

  • kernel – An N by K by K array of N samples of an integral kernel applying transformations to each associated timestep of x

Returns:

An N by K array represented the result of the transformation.

trajectory_manifold.helpers.apply_kernel_vec(x: Float[Array, 'functions timesteps dim'], kernel: Float[Array, 'timesteps dim dim']) → Float[Array, 'timesteps dim']

Helper Function to apply an integration kernel to an input.

Applies the kernel to an input at each timestep. Vectorized along a set of functions.

Parameters:

  • x – An M by N by K array of N samples of each of M functions taking values in a K-dimensional space.

  • kernel – An N by K by K array of N samples of an integral kernel applying transformations to each associated timestep of x

Returns:

An M by N by K array represented the result of the transformation.

trajectory_manifold.helpers.frobenius_inner_product(x: Float[Array, '*dim'], y: Float[Array, '*dim']) → Float

Computes the Frobenius inner product of two matrices.

Given two multidimensional arrays, computes the sum of the elementwise product of the arrays.

Parameters:

  • x – A multidimensional array.

  • y – A multidimensional array.

Returns:

The sum of the elementwise product of x and y.

trajectory_manifold.helpers.safe_det(A: Float[Array, 'dim dim'])

Return the determinant with an empty matrix set to 1

Used for computing the determinants of block diagonal matrices in the case where a block size may be zero

Parameters:

A – A matrix to compute the determinant

Returns:

The determinant of the matrix, or 1 if the matrix is empty

trajectory_manifold.helpers.trapezoidal_correlation(U: Float[Array, 'functions timesteps dim'], step_size: Float) → Float[Array, 'functions functions']

Computes the inner products between rows of the given matrix.

Constructs an M by M matrix of approximate inner products between M multi-variate functions computed using N evenly spaced samples in a trapezoidal integration scheme.

Parameters:

  • U – M by N by K array representing N samples each of M functions that take values in a K-dimensional space.

  • step_size – Spacing between sample points in the functions.

Returns:

An M by M matrix where the (i,j)’th element is the approximate inner product between rows i and j of the input matrix.

trajectory_manifold.helpers.trapezoidal_correlation_weighted(U: Float[Array, 'functions timesteps dim'], step_size: Float, kernel: Float[Array, 'timesteps dim dim']) → Float[Array, 'functions functions']

Computes the inner products between rows of the given matrix.

Constructs an M by M matrix of approximate inner products between M multi-variate functions computed using N evenly spaced samples in a trapezoidal integration scheme. Uses an integral kernel to define the inner product.

Parameters:

  • U – M by N by K array representing N samples each of M functions that take values in a K-dimensional space.

  • step_size – Spacing between sample points in the functions.

  • kernel – N by K by K array representing an integral kernel defining the inner product in the underlying space.

Returns:

An M by M matrix where the (i,j)’th element is the approximate inner product between rows i and j of the input matrix.

trajectory_manifold.helpers.trapezoidal_inner_product(x: Float[Array, 'timesteps dim'], y: Float[Array, 'timesteps dim'], step_size: Float) → Float

Approximate the inner product by the trapezoidal rule.

Computes an approximate inner product between two functions represented by a finite-grid approximation with a fixed step size. Computed through the trapezoidal integration scheme.

Parameters:

  • x – Grid approximation of the first function. Each row represents the value of the multivariate function at a given timestep.

  • y – Grid approximation of the second function. Each row represents the value of the multivariate function at a given timestep.

  • step_size – Spacing between sample points in the functions.

Returns:

An approximation of the L2 inner product.

trajectory_manifold.helpers.trapezoidal_inner_product_weighted(x: Float[Array, 'timesteps dim'], y: Float[Array, 'timesteps dim'], step_size: Float, kernel: Float[Array, 'timesteps dim dim']) → Float

Approximate the inner product by the trapezoidal rule.

Computes an approximate inner product between two functions represented by a finite-grid approximation with a fixed step size. Computed through the trapezoidal integration scheme.

Parameters:

  • x – Grid approximation of the first function. Each row represents the value of the multivariate function at a given timestep.

  • y – Grid approximation of the second function. Each row represents the value of the multivariate function at a given timestep.

  • step_size – Spacing between sample points in the functions.

Returns:

An approximation of the L2 inner product.

trajectory_manifold.helpers.trapezoidal_matrix_product(X: Float[Array, 'functions1 timesteps dim'], Y: Float[Array, 'functions2 timesteps dim'], step_size: Float) → Float[Array, 'functions1 functions2']

Computes the product of two matrices where rows represent functions.

Given two matrices, X and Y, such that the first index indicates a chosen function, computes X @ Y.T based on the trapezoidal rule for integration.

Parameters:

  • X – An M by N by K array of N samples of each of M functions taking values in a K-dimensional space.

  • Y – An M by N by K array of N samples of each of M functions taking values in a K-dimensional space.

  • step_size – Spacing between sample points in the functions.

Returns:

The result of the generalized matrix-matrix product.

trajectory_manifold.helpers.trapezoidal_pytree_vector_product(y: Float[Array, 'timesteps dim'], X: PyTree, step_size: Float) → PyTree

Computes the pytree-vector product where leaves are sampled functions.

Given a PyTree where leaves are 2-dimensional arrays representing sampled function and another 2-dimensional array reepresenting a single multivariate function, returns a PyTree of the same shape as the original input with leaves resulting from the L2 inner products with the function.

Parameters:

  • y – A 2D array with indices (timestep, dim)

  • X – A PyTree where leaves are 2D arrays with indices (timestep, dim)

  • step_size – Step size required for numerical integration scaling

Returns:

A PyTree of the same shape a X, but with results of inner products as leaves.

Examples

The module trajectory_manifold.examples contains a collection of example systems to be used with the trajectory forecasting work. It contains functions which generate vector fields for a linear system, a periodic system, and a chaotic system.

trajectory_manifold.examples.linear_vector_field(A: Float[Array, 'dim dim']) → Callable[[Float, Float[Array, 'dim'], PyTree], Float[Array, 'dim']]

Returns a function representing a linear system.

Parameters:

A – A matrix defining the system dynamics

Returns:

A function representing the vector field for a linear system.

trajectory_manifold.examples.lorenz96_vector_field(F: float) → Callable[[Float, Float[Array, 'dim'], PyTree], Float[Array, 'dim']]

Returns a function representing the Lorenz 96 system.

Parameters:

F – The constant forcing term

Returns:

A function representing the vector field for the Lorenz 96 system.

trajectory_manifold.examples.lorenz_vector_field(sigma: float = 10, rho: float = 28, beta: float = 2.6666666666666665) → Callable[[Float, Float[Array, 'dim'], PyTree], Float[Array, 'dim']]

Returns a function representing the Lorenz system.

Parameters:

  • sigma – Lorenz Parameter

  • rho – Lorenz Parameter

  • beta – Lorenz Parameter

Returns:

A function representing the vector field for the Lorenz 96 system.

trajectory_manifold.examples.lotka_volterra_vector_field(prey_growth: float, predation_rate: float, predator_growth: float, predator_decay: float) → Callable[[Float, Float[Array, 'dim'], PyTree], Float[Array, 'dim']]

Returns a function representing a Lotka-Volterra system.

Parameters:

  • prey_growth – Exponential growth rate of prey without predation.

  • predation_rate – Decay rate for prey in contact with predator.

  • predator_growth – Growth rate for predators in contact with prey.

  • predator_decay – Exponential decay rate of predators without prey.

Returns:

A function representing the vector field for a Lotka-Volterra system.