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']) 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_conditon – The position in the statespace to be pushed onto the manifold.

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'], 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_conditon – The position in the statespace to be pushed onto the manifold.

  • 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'], 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.

  • 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.state_log_posterior(vector_field: Callable[[any, Float[Array, ' dim'], any], Float[Array, ' dim']], observations: Float[Array, ' timesteps dim2'], observation_times: Float[Array, ' timesteps'], observation_likelihood: Callable[[Float[Array, ' dim2'], Float[Array, ' dim']], float], state_log_prior: Callable[[Float[Array, ' dim']], float], parameters: SolverParameters) Callable[[Float[Array, ' dim']], 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.

  • state_prior – A function representing the prior distribution of the state of the system at the initial observation time.

Returns:

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

trajectory_manifold.estimation.state_posterior(vector_field: Callable[[any, Float[Array, ' dim'], any], Float[Array, ' dim']], observations: Float[Array, ' timesteps dim2'], observation_times: Float[Array, ' timesteps'], observation_likelihood: Callable[[Float[Array, ' dim2'], Float[Array, ' dim']], float], state_prior: Callable[[Float[Array, ' dim']], float]) Callable[[Float[Array, ' dim']], 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.

  • state_prior – A function representing the prior distribution of the state of the system at the initial observation time.

Returns:

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

trajectory_manifold.estimation.trajectory_likelihood(vector_field: Callable[[any, Float[Array, ' dim'], any], 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']], 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.

Returns:

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

trajectory_manifold.estimation.trajectory_log_likelihood(vector_field: Callable[[any, Float[Array, ' dim'], any], 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']], 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.

Returns:

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

trajectory_manifold.estimation.trajectory_log_posterior(vector_field: Callable[[any, Float[Array, ' dim'], any], Float[Array, ' dim']], observations: Float[Array, ' timesteps dim2'], observation_times: Float[Array, ' timesteps'], observation_likelihood: Callable[[Float[Array, ' dim2'], Float[Array, ' dim']], float], state_log_prior: Callable[[Float[Array, ' dim']], float], time_interval: tuple[float, float], parameters: SolverParameters) Callable[[Float[Array, ' dim']], 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.

  • state_prior – A function representing the prior distribution of the state of the system at the initial observation time.

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

Returns:

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

trajectory_manifold.estimation.trajectory_posterior(vector_field: Callable[[any, Float[Array, ' dim'], any], Float[Array, ' dim']], observations: Float[Array, ' timesteps dim2'], observation_times: Float[Array, ' timesteps'], observation_likelihood: Callable[[Float[Array, ' dim2'], Float[Array, ' dim']], float], state_prior: Callable[[Float[Array, ' dim']], float], time_interval: tuple[float, float]) Callable[[Float[Array, ' dim']], 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.

  • state_prior – A function representing the prior distribution of the state of the system at the initial observation time.

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

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'], vector_field: Callable[[Float, Float[Array, ' dim'], PyTree], Float[Array, ' dim']], trajectory: Float[Array, ' dim dim2'], params: SolverParameters) Float[Array, ' dim']

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.

  • 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.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.

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.