High-level interfaces

This page contains a reference of some convenience classes for performing state estimation with stochastic variational inference (SVISE). Inheriting from the SDELearner class will handle most things common to all potential parametrizations of the prior and approximate posterior such as computing the loss, resampling, etc.

SDELearner

class svise.sde_learning.SDELearner(marginal_sde: MarginalSDE, likelihood: Likelihood, quad_scheme: QuadRule1D, sde_prior: SDEPrior, diffusion_prior: DiffusionPrior, n_reparam_samples: int)[source]

Base class for combining different priors and likelihoods for performing simultaneous state and governing equation discovery.

Parameters:

  • marginal_sde (MarginalSDE) – Parameterization for the approximate Markov GP

  • likelihood (Likelihood) – Likelihood for the observations

  • quad_scheme (QuadRule1D) – 1D quadrature rule used to approximate the residual

  • sde_prior (SDEPrior) – The prior over the drift function

  • diffusion_prior (DiffusionPrior) – The prior over the diffusion matrix

  • n_reparam_samples (int) – The number of samples to use when doing the reparametrization trick

diffusion(summary_type: str = 'sample') → Tensor[source]

Returns the diffusion matrix

Parameters:

summary_type (str, optional) – “sample” or “mean”

Returns:

the diffusion matrix

Return type:

Tensor

drift(t: Tensor, z: Tensor, summary_type: str = 'sample') → Tensor[source]

Returns the drift function for the approximate SDE

Parameters:

  • t (Tensor) – time stamps to evaluate drift function

  • z (Tensor) – states at time stamps

Returns:

drift evaluated at time stamps and states

Return type:

Tensor

elbo(t: Tensor, x: Tensor, beta: float, N: int, print_loss=False, compat_mode=True) → Tensor[source]

Computes the reweighted evidence lower bound (i.e. ELBO / N)

Parameters:

  • t (Tensor) – (bs,) time stamps of measurements

  • x (Tensor) – (bs, D) measurements

  • beta (float) – weight of kl-divergence

  • N (int) – total number of data points

Returns:

normalized evidence lower bound

Return type:

Tensor

resample_sde_params(n: int | None = None) → None[source]

Resample the sde_prior and diffusion_prior

Parameters:

n (int, optional) – number of samples to use when resampling

residual_loss(t: Tensor) → Tensor[source]

Returns the residual loss evaluated at a batch of time stamps

Parameters:

t (Tensor) – batch of time stamps

Returns:

residual evaulated at time stamps t

Return type:

Tensor

sde_params() → List[Parameter][source]

Returns all parameters not related to the state estimate

Returns:

parameters not related to the state estimate

Return type:

List[Parameter]

state_params() → List[Parameter][source]

Returns parameters related to only the state estimator

Returns:

parameters related to the state estimator

Return type:

List[Parameter]

Sparse polynomial SDEs

class svise.sde_learning.SparsePolynomialSDE(d: int, t_span: Tuple[float, float] | Tuple[Tensor, Tensor], degree: int, n_reparam_samples: int, G: Tensor, num_meas: int, measurement_noise: Tensor, tau: float = 1e-05, train_t: Tensor | None = None, train_x: Tensor | None = None, input_labels: List[str] | None = None, n_quad: int = 128, quad_percent: float = 0.8, n_tau: int = 200)[source]

Subclass of SDELearner for the case that the prior over the drift function is a sparse linear combination of polynomials, the observation matrix is linear, observation noise is a diagonal Gaussian, and the Markov GP is parameterized by RBF models with the Matern 5/2 kernel. Also performs some useful initialization to make training more stable.

Parameters:

  • d (int) – dimension of the state

  • t_span (Union[Tuple[float, float], Tuple[Tensor, Tensor]]) – min and max boundary for RBF centers

  • degree (int) – degree of polynomial in drift function

  • n_reparam_samples (int) – number of samples to use when using the reparametrization trick

  • G (Tensor) – observation matrix (i.e. y = Gx)

  • num_meas (int) – number of observations (should be G.shape[0])

  • measurement_noise (Tensor) – variance of observations

  • tau (float) – global half-cauchy scaling parameter

  • train_t (Tensor, optional) – observation time stamps (ns,)

  • train_x (Tensor, optional) – observations (ns, num_meas)

  • input_labels (List[str], optional) – name of each variable (i.e. [“x”, “y”, …])

  • n_quad (int, optional) – number of quadrature nodes to use

  • quad_percent (float, optional) – 1 - quad_percent of the quad_nodes will be sampled uniformly in the time window

  • n_tau (int, optional) – number of centers for the RBF models

elbo(t: Tensor, x: Tensor, beta: float, N: int, print_loss=False) → Tensor[source]

Computes the reweighted evidence lower bound (i.e. ELBO / N)

Parameters:

  • t (Tensor) – (bs,) time stamps of measurements

  • x (Tensor) – (bs, D) measurements

  • beta (float) – weight of kl-divergence

  • N (int) – total number of data points

Returns:

normalized evidence lower bound

Return type:

Tensor

Neural SDE

class svise.sde_learning.NeuralSDE(d: int, t_span: Tuple[float, float] | Tuple[Tensor, Tensor], n_reparam_samples: int, G: Tensor, drift_layer_description: List[int], nonlinearity: Module, measurement_noise: Tensor, tau: float = 1e-05, train_t: Tensor | None = None, train_x: Tensor | None = None, n_quad: int = 128, quad_percent: float = 0.8, n_tau: int = 200, Q_diag: Tensor | None = None)[source]

Subclass of SDELearner when the drift function is represented by a fully connected neural network.

Parameters:

  • d (int) – dimension of the state

  • t_span (Union[Tuple[float, float], Tuple[Tensor, Tensor]]) – min and max boundary for RBF centers

  • n_reparam_samples (int) – number of samples to use when using the reparametrization trick

  • G (Tensor) – observation matrix (i.e. y = Gx)

  • drift_layer_description (List[int]) – number of neurons in each layer of the drift function

  • nonlinearity (nn.Module) – nonlinearity to use in the drift function

  • measurement_noise (Tensor) – variance of observations

  • tau (float) – global half-cauchy scaling parameter

  • train_t (Tensor, optional) – observation time stamps (ns,)

  • train_x (Tensor, optional) – observations (ns, num_meas)

  • n_quad (int, optional) – number of quadrature nodes to use

  • quad_percent (float, optional) – 1 - quad_percent of the quad_nodes will be sampled uniformly in the time window

  • n_tau (int, optional) – number of centers for the RBF models

  • Q_diag (Tensor, optional) – diagonal component of the diffusion matrix

Sparse polynomial for SDEs with Newtonian drift

class svise.sde_learning.SparsePolynomialIntegratorSDE(d: int, t_span: Tuple[float, float] | Tuple[Tensor, Tensor], degree: int, n_reparam_samples: int, G: Tensor, num_meas: int, measurement_noise: Tensor, tau: float = 1e-05, train_t: Tensor | None = None, train_x: Tensor | None = None, input_labels: List[str] | None = None, n_tau: int = 200)[source]

Assumes that states are governed by a drift function which is a function of both the position and it’s velocity but we can only observe the states.

Like the SparsePolynomialSDE, this is a Subclass of SDELearner for the case that the prior over the drift function is a sparse linear combination of polynomials, the observation matrix is linear, observation noise is a diagonal Gaussian, the Markov GP is parameterized by RBF models with the Matern 5/2 kernel. Also performs some useful initialization to make training more stable.

Parameters:

  • d (int) – dimension of the state

  • t_span (Union[Tuple[float, float], Tuple[Tensor, Tensor]]) – min and max boundary for RBF centers

  • degree (int) – degree of polynomial in drift function

  • n_reparam_samples (int) – number of samples to use when using the reparametrization trick

  • G (Tensor) – observation matrix (i.e. y = Gx)

  • num_meas (int) – number of observations (should be G.shape[0])

  • measurement_noise (Tensor) – variance of observations

  • tau (float) – global half-cauchy scaling parameter

  • train_t (Tensor, optional) – observation time stamps (ns,)

  • train_x (Tensor, optional) – observations (ns, num_meas)

  • input_labels (List[str], optional) – name of each variable (i.e. [“x”, “y”, …])

  • n_tau (int, optional) – number of centers for the RBF models

State estimation

class svise.sde_learning.StateEstimator(d: int, t_span: Tuple[float, float] | Tuple[Tensor, Tensor], n_reparam_samples: int, G: Tensor, drift: Callable[[Tensor, Tensor], Tensor], num_meas: int, measurement_noise: Tensor, tau: float = 1e-05, train_t: Tensor | None = None, train_x: Tensor | None = None, n_quad: int = 128, quad_percent: float = 0.8, n_tau: int = 200, Q_diag: Tensor | None = None)[source]

Subclass of SDELearner for the case the drift and diagional diffusion matrix is known, the observation matrix is linear, observation noise is a diagonal Gaussian, the Markov GP is parametrized by RBF models with the Matern 5/2 kernel. Also performs some useful initialization to make training more stable.

Parameters:

  • d (int) – dimension of the state

  • t_span (Union[Tuple[float, float], Tuple[Tensor, Tensor]]) – min and max boundary for RBF centers

  • n_reparam_samples (int) – number of samples to use when using the reparametrization trick

  • G (Tensor) – observation matrix (i.e. y = Gx)

  • drift (Callable[[Tensor, Tensor], Tensor]) – known drift function

  • num_meas (int) – number of observations (should be G.shape[0])

  • measurement_noise (Tensor) – variance of observations

  • tau (float) – global half-cauchy scaling parameter

  • train_t (Tensor, optional) – observation time stamps (ns,)

  • train_x (Tensor, optional) – observations (ns, num_meas)

  • input_labels (List[str], optional) – name of each variable (i.e. [“x”, “y”, …])

  • n_tau (int, optional) – number of centers for the RBF models

  • Q_diag (Tensor, optional) – diagonal component of the diffusion matrix,

elbo(t: Tensor, x: Tensor, beta: float, N: int, print_loss=False) → Tensor[source]

Computes the reweighted evidence lower bound (i.e. ELBO / N)

Parameters:

  • t (Tensor) – (bs,) time stamps of measurements

  • x (Tensor) – (bs, D) measurements

  • beta (float) – weight of kl-divergence

  • N (int) – total number of data points

Returns:

normalized evidence lower bound

Return type:

Tensor