Index of public API

EKS{NF} <: EnsembleInferenceAlgorithm

Represents a proxy for the Ensemble Kalman Sampler (Garbuno-Inigo et al. 2020) implementation provided by EnsembleKalmanProcesses.

ESMDA{NF} <: EnsembleInferenceAlgorithm

Implementation of the ensemble-smother multiple data assimilation algorithm of Emerick et al. 2013.

Emerick, Alexandre A., and Albert C. Reynolds. "Ensemble smoother with multiple data assimilation." Computers & Geosciences 55 (2013): 3-15.

EmpiricalGaussian <: GaussianApproximationMethod

Represents a simple Gaussian approximation method which calculates a sample mean and covariance.

EnIS

Basic ensemble importance sampling (EnIS) inference algorithm. Also sometimes referred to as the "particle batch smoother" (PBS) or "generalized likelihood uncertainty estimation" (GLUE) depending on the context.

EnsembleSolver{algType,probType,ensalgType,stateType<:EnsembleState,kwargTypes}

Generic implementation of an iterative solver for any ensemble-based algorithm. Uses the SciML EnsembleProblem interface to automatically parallelize forward runs over the ensemble.

GaussianApproximationMethod

Base type for Gaussian approximations of arbitrary prior distributions.

JointPrior{modelPriorType<:AbstractSimulatorPrior,likNames,likPriorTypes,axesType} <: AbstractSimulatorPrior

Represents the "joint" prior p(θₘ,θₗ) where θ = [θₘ θₗ] are the full set of parameters in the joint; distribution p(x,θ). θₘ are the model (simulator) parameters and θₗ are the noise/error model parameters.

Constructs a JointPrior from the given prior and likelihoods.

MCMC(alg, strat; kwargs...)

Generic consturctor for MCMC algorithms with implementation-specific type alg and "strategy" strat.

MCMC{algType,stratType,kwargType} <: SimulatorInferenceAlgorithm

Generic container type for Markov Chain Monte Carlo (MCMC) based inference algorithms.

NamedProductPrior{distTypes<:NamedTuple}

Simple diagonal prior that wraps a NamedTuple of distributions from the Distributions package.

Alias for EnIS.

Allocate and construct a new SimulationData from the given backend, or create a new InMemoryStorage backend if not specified.

SimulationData{B}

A view of the data for a single simulation stored in the given StorageBackend. The optional handle field maintains a persistent storage handle when needed (e.g., for observables that require scratch storage across multiple observe! calls).

SimulationDataSet{B<:StorageBackend}

A view of the whole collection of simulations held in a single StorageBackend — the simulations run during an inference procedure. SimulationDataSet() uses an in-memory backend; OnDiskSimulationDataSet(path) (provided by the SimulationBasedInferenceJLD2Ext extension) uses a disk-backed one.

Indexing returns a SimulationData view of the i-th simulation; iterating yields (input, outputs, metadata) triples.

See Simulator(simulator)

Simulator(simulator)
Simulator(::Type{SimulatorType})

Trait that characterizes the forward map defined by a simulator function or type. Options are ForwardMap, Iterative, or Dynamical. ForwardMap simulators are defined by a single function f(p) that should directly return the output of the simulation, while Iterative and Dynamical simulators should be defined using custom types that implement the following interface:

• CommonSolve.init(simulator, args...; p::AbstractVecOrMat, tspan, kwargs...)::SimulationType
• CommonSolve.step!(::SimulationType, [dt,] args...; kwargs...)
• CommonSolve.solve!(::SimulationType)::SolutionType
• current_state(::SimulationType) (optional)
• isdone(::SimulationType)::Boolean (optional)

where SimulationType and SolutionType are types defined by the user or downstream package. Dynamical simulators must additionally define the following:

• timespan(::SimulationType)::NTuple{2, timeType}
• current_time(::SimulationType)::timeType

The keyword arguments p and tspan are mandatory and respectively correspond to the choice of simulation parameters and time span. Simulators characterized by a SciML problem type are broadly assumed to be Iterative or Dynamical depending on the problem type.

SimulatorForwardProblem{simType, paramType, seedType, obsTypes, names} <: SciMLBase.AbstractSciMLProblem

Represents a "forward" problem from parameters/initial conditions to output SimulatorObservables. This type wraps an underlying simulator, a set of Observables, and an optional RNG seed for stochastic forward models.

SimulatorForwardProblem(f, p::AbstractVector, observables::SimulatorObservable...)

Creates a forward problem from the given function or callable type f with parameters p and the given observables.

SimulatorForwardProblem(f, 0::AbstractVector)

Creates a forward problem from the given function or callable type f with parameters p and a default transient observable. Note that this constructor calls f(p) in order to determine the shape of the observable. If f is very costly and this is undesirable, it is recommended to use the explicit constructor.

SimulatorForwardProblem(prob::SciMLBase.AbstractSciMLProblem, observables::SimulatorObservable...)

Constructs a generic simulator forward problem from the given AbstractSciMLProblem; note that this could be any problem type, e.g. an optimization problem, nonlinear system, quadrature, etc.

SimulatorForwardProblem(simulator, p::AbstractVector, observables::SimulatorObservable...; rng_seed)

Constructs a SimulatorForwardProblem from the given simulator, parameters p, and observables.

SimulatorForwardProblem(prob::SciMLBase.AbstractSciMLProblem, observables::SimulatorObservable...; rng_seed)

Constructs a SimulatorForwardProblem from the given SciML problem and observables.

SimulatorForwardSolution{solType,probType}

Solution for a SimulatorForwardProblem that wraps the underlying forward solution.

Base type for all simulator-based inference algorithms.

SimulatorInferenceProblem{
    modelPriorType<:AbstractSimulatorPrior,
    uType,
    fwdProbType<:SimulatorForwardProblem,
    fwdSolverType,
    priorType<:JointPrior
} <: SciMLBase.AbstractSciMLProblem

Represents a generic simulation-based Bayesian inference problem for finding the posterior distribution over model parameters given some observed data.

SimulatorInferenceProblem(
    prob::SimulatorForwardProblem,
    forward_solver,
    prior::AbstractSimulatorPrior,
    likelihoods::AbstractLikelihood...;
    metadata::Dict=Dict(),
)

Constructs a SimulatorInferenceProblem from the given forward problem, prior, and likelihoods. Additional user-specified metadata may be included in the metadata dictionary.

SimulatorInferenceSolution{algType, probType, storageType}

Generic container for solutions to SimulatorInferenceProblems. The type of result is method dependent and should generally correspond to the final state or product of the inference algorithm (e.g. posterior samples). The field storage should be an instance of SimulationDataSet.

SimulatorLikelihood{distType,obsType,dataType,priorType}

Represents a simulator-based likelihood function. A SimulatorLikelihood consists of four basic components:

(1) A distribution type, e.g. Normal,

(2) A SimulatorObservable which represents the observation operator,

(3) A set of data, usually a Vector or Matrix, which matches the structure of the observable,

(4) A prior distribution governing one or more additional parameters required to compute the likelihood.

Due to the typically high cost of evaluating the parameter forward map, SimulatorLikelihood effectively decouples the computation of the likelihood from the simulator via the SimulatorObservable. When the SimulatorLikelihood is evaluated, the observable output is obtained from the simulation data via getvalue(data, obs) and the only additional parameters needed are those specified by prior.

SimulatorLikelihood(::Type{distType}, obs, data, prior, name=nameof(obs)) where {distType}

Creates a SimulatorLikelihood with the given distribution type, observable, data source, and prior distribtuion. A custom identifier may also be specified via the name argument; by default, the name of the observable is used.

SimulatorObservable{N, outputType<:SimulatorOutput, funcType, coordsType} <: Observable{outputType}

Represents a named "observable" that stores output from a simulator. obsfunc defines a mapping from the simulator state to the observed quantity. The type and implementation of output determines how the samples are stored. The simplest output type is Transient which simply maintains a pointer to the last observed output.

SimulatorObservable(func, coords::Tuple; output::SimulatorOutput = Transient(), name::Symbol = :obs)

Constructs an observable based on the given function func(state)::T and output type. Defaults to Transient output which simply saves the last observed value of func. The coordinates coords describe the shape of the output.

TimeAggregated{timeType, outputType, reducerType, converterType} <: SimulatorOutput

SimulatorOutput representing a coarser-scale temporal aggregation of another (source) observable's already-saved outputs. Unlike TimeSampled, it does not sample the raw simulator state and is not observed during the forward solve; its value is computed by a single streaming reduction over the source's stored outputs and materialized after the solve completes (see TimeAggregatedObservable). Peak memory is bounded by one aggregation window, not the full source series.

TimeAggregatedObservable(
    source::TimeSampledObservable,
    tsave::AbstractVector;
    transform = identity,
    reducer = mean,
    coords = source.coords,
    name = Symbol(nameof(source), :_agg),
    output_type = Any,
)

Constructs an observable that aggregates the already-saved outputs of source at the coarser save times tsave (which must be a subset of savetimes(source)). For each aggregation window, the per-slice transform is applied to every source time slice and the results are combined with reducer. The value is computed once by a streaming reduction over the source's stored outputs (materialized after the forward solve, or lazily on first getvalue), so the source observable can still be retained for diagnostics while this coarser observable is compared to observations via a likelihood.

transform defaults to identity; when it changes the per-slice shape (e.g. a spatial reduction), pass the resulting spatial coords explicitly.

TimeSampled{timeType, storageType, reducerType, converterType} <: SimulatorOutput

SimulatorOutput which buffers samples taken from the simulator at preset times and applies a reduction operation at (lower frequency) save times. A simple example would be a windowed average or resampling operation that saves averages over higher frequency samples.

TimeSampled(
    t0::tType,
    tsave::AbstractVector{tType};
    reducer=mean,
    samplerate=default_sample_rate(tsave),
    handle=nothing,  # Optional handle parameter for efficient scratch storage
) where {tType}

Constructs a TimeSampled simulator output which iteratively samples and stores outputs on each call to observe!.

Transient{T} <: SimulatorOutput

Simple output type that retains only the last observed value of the observable function. The value itself is stored in the SimulationData.

DiagonalGaussianLikelihood(obs, data, prior, name)

Alias for SimulatorLikelihood(DiagNormal, obs, data, prior, name). Represents a multivaraite i.i.d Guassian likelihood with heteroscedastic noise.

DiracLikelihood(
    obs,
    data,
    name=nameof(obs),
)

Represents a Dirac delta likelihood function which assigns full probability mass to a single point, i.e. the predicted value of the simulator observable.

GaussianLikelihood(obs, data, prior, name)

Alias for SimulatorLikelihood(Normal, obs, data, prior, name). Represents a univariate Guassian likelihood.

ImplicitLikelihood(
    obs,
    data,
    name=nameof(obs),
)

Represents a Dirac delta likelihood function which assigns full probability mass to a single point, i.e. the predicted value of the simulator observable.

IsotropicGaussianLikelihood(obs, data, prior, name)

Alias for SimulatorLikelihood(IsoNormal, obs, data, prior, name). Represents a multivaraite i.i.d Guassian likelihood with homoscedastic noise.

allocate!(storage::SimulationDataSet, inputs=nothing; metadata...)

Allocate simulation data storage in the underlying backend (with the given input) and return a SimulationData view of it.

allocate!(storage::SimulationDataSet, inputs::AbstractMatrix; metadata...)

Allocate a new SimulationData on the storage backend for each column of inputs. Returns a vector of SimulationData of length size(inputs, 2).

autoprior(mean, stddev; lower=-Inf, upper=Inf)

Helper function that automatically constructs a prior distribution with the given mean, standard deviation, and support. Note that lower < mean <= upper must be satisfied. For unbounded variables, Normal is returned. For variabels bounded from below (i.e. either lower or upper bound is finite), a transformed LogNormal distribution is constructed. For double bounded variables, (i.e. both lower and upper are finite) a tranformed Beta distribution is constructed.

coordinates(dims...)

Converts arguments dims to a tuple of coordinate Dimensions according to the following rules:

- Integers `n` are converted to simple step indices `1:n`
- Vectors are converted to `Dim`s

coordinates(obs::Observable)

Retrieves coordinates for each dimension of the observables as a Tuple with length matching the number of dimensions.

ensemble_kalman_analysis(
    prior::AbstractVecOrMat,
    obs::AbstractVector,
    pred::AbstractMatrix,
    alpha,
    R_cov;
    ρ_AB=1.0,
    ρ_BB=1.0,
    stochastic=true,
    dosvd=true,
    svd_thresh=0.90,
    rng::AbstractRNG=Random.GLOBAL_RNG,
)

Performs a single ensemble Kalman analysis step. Adapted from the python implementation by Kristoffer Aalstad:

https://github.com/ealonsogzl/MuSA/blob/0c02b8dc25a0f902bf63099de68174f4738705f0/modules/filters.py#L33

ensemblestep!(solver::EnsembleSolver{algType}) where {algType}

Executes a single ensemble step (forward solve + update) for the given algorithm type. Must be implemented by all ensemble algorithm implementations.

finalize!(solver::EnsembleSolver)

Finalizes the solver state after iteration has completed. Default implementation invokes ensemble_forward! on the current ensemble state.

finalize!(data::SimulationData, observables)

Materialize any TimeAggregatedObservables from their sources' stored outputs. Called at the end of a forward solve so that aggregated observables are persisted like any other output (and thus available to storage/likelihood retrieval paths). A no-op for non-aggregated observables.

forward_map(prior::AbstractSimulatorPrior, ζ)

Applies the forward map from the sample space of the prior to the parameter space of the forward model (simulator). Note that this mapping need not necessarily be bijective in the case of hierarchical or reparamterized formulations of the model parameter prior. Defaults to returning identity(ζ).

from_moments(::Type{T}, mean, stddev) where {T<:Distribution}

Construct a distribution of type T using the method of moments. It is assumed that the given mean lies within the untransformed support of the distribution.

gaussian_gapprox(prior::AbstractSimulatorPrior; num_prior_samples::Int=10_000, rng::Random.AbstractRNG=Random.default_rng())

Builds an empirical multivariate Gaussian approximation of the given prior distribution by computing the moments of the transformed samples.

get_ensemble(sol::EnsembleInferenceSolution, iter::Int=iterations(sol.storage))

Fetches the state of the ensemble from the given solution object by reconstructing the parameter matrix from the per-member simulation records at iteration iter (columns are ensemble members). For iterative algorithms, the optional argument iter may be provided; iter < 1 selects the final iteration.

get_ensemble(state::EnsembleState)

Retrieves the current ensemble matrix from the given EnsembleState. Must be implemented for each ensemble algorithm state type.

get_observables(sol::EnsembleInferenceSolution, iter::Int=iterations(sol.storage))

Returns a NamedTuple of the ensemble simulated observables at iteration iter, assembled by concatenating each member's observable value along the ensemble dimension.

get_transformed_ensemble(sol::EnsembleInferenceSolution, iter::Int=iterations(sol.storage))

Fetches the transformed ensemble from the given solution object. For iterative algorithms, the optional argument iter may be provided, which then retrieves the ensemble at the given iteration.

getoutput(data::SimulationData, name::Symbol)

Return the collected output sequence (a Vector) for observable name.

getoutputs(data::SimulationData)

Return a NamedTuple mapping each output name to its corresponding data series.

getvalue(data::SimulationData, ::Observable)

Retrieve the observable value (at all coordinates) from the given SimulationData.

importance_weights(obs::AbstractVector, pred::AbstractMatrix, R::AbstractMatrix)

Computes the importance weights from the given observations, predictions, and noise covariance. When used for data assimilation, this algorithm is also known as the "particle batch smoother". This implementation is adapted from the python implementation by Kristoffer Aalstad:

https://github.com/ealonsogzl/MuSA/blob/0c02b8dc25a0f902bf63099de68174f4738705f0/modules/filters.py#L169

Original docstring:

Inputs: obs: Observation vector (m x 1 array) pred: Predicted observation ensemble matrix (m x N array) R: Observation error covariance 'matrix' (m x 1 array, or scalar) Outputs: w: Posterior weights (N x 1 array) Neff: Effective sample size Dimensions: N is the number of ensemble members and m is the number of observations.

Here we have implemented the particle batch smoother, which is a batch-smoother version of the particle filter (i.e. a particle filter without resampling), described in Margulis et al. (2015, doi: 10.1175/JHM-D-14-0177.1). As such, this routine can also be used for particle filtering with sequential data assimilation. This scheme is obtained by using a particle (mixture of Dirac delta functions) representation of the prior and applying this directly to Bayes theorem. In other words, it is just an application of importance sampling with the prior as the proposal density (importance distribution). It is also the same as the Generalized Likelihood Uncertainty Estimation (GLUE) technique (with a formal Gaussian likelihood)which is widely used in hydrology.

This particular scheme assumes that the observation errors are additive Gaussian white noise (uncorrelated in time and space). The "logsumexp" trick is used to deal with potential numerical issues with floating point operations that occur when dealing with ratios of likelihoods that vary by many orders of magnitude.

initialstate(
    alg::EnsembleInferenceAlgorithm,
    ens::AbstractMatrix,
    obs::AbstractVector,
    obscov::AbstractMatrix;
    rng::AbstractRNG=Random.GLOBAL_RNG
)

Constructs the initial ensemble state for the given algorithm and observations.

iterations(storage::SimulationDataSet)

Return the number of distinct iter values recorded in the simulations' metadata, i.e. the number of inference iterations. Returns 0 for an empty storage and 1 when no iter metadata is present.

loglikelihood(data::SimulationData, lik::SimulatorLikelihood, args...)

Evaluates the log-likelihood of lik on the observable value stored in data by constructing the predictive_distribution and evaluating the logpdf of the data.

logprob(d::UnivariateDistribution, x::Number)

Alias for logpdf(d,x) on Distribution types.

obscov(::AbstractMatrix)
obscov(::AbstractVector)
obscov(::Number)

Builds a standard form multivariate normal covariance matrix from the given matrix, vector (diagonal), or scalar (isotropic).

obscov(likelihoods::SimulatorLikelihood...)

Implementations should build a combined covariance matrix from the given likelihood types. The default implementation simply throws an error.

observe!(data::SimulationData, ::Observable, state)

Computes the relevant state variables from the given simulator state and stores them in the given SimulationData.

observe!(data::SimulationData, obs::TimeSampledObservable, sim)

Observe the given time-sampled observable from the current simulator state. Called at each sample time, it extracts the observable value into the scratch buffer and, when the current time (current_time(sim)) coincides with a save time, reduces the accumulated window into the output buffer. The simulator must implement current_time.

prior(args...; kwargs...)

Generic constructor for prior distribution types that can be implemented by subtypes of AbstractSimulatorPrior.

prior(name::Symbol, dist::Distribution)

Alias for NamedProductPrior((name = dist)).

prior(; dists...)

Alias for NamedProductPrior((; dists...)).

store!(data::SimulationData, name::Symbol, value)

Append value to the persistent output series for observable name.

store!(storage::SimulationDataSet, data::SimulationData; metadata...)

Append an existing SimulationData (possibly held in a different backend, e.g. a per-member forward solve) to storage, copying its input and output series into the backing store and merging any extra metadata.

unconstrained_forward_map(prior::AbstractSimulatorPrior, ζ)

Applies the forward map from unconstrained to forward model parameter (g ∘ f): Ζ ↦ Θ ↦ Φ where f⁻¹ is the inverse bijector of prior (i.e. mapping from unconstrained Ζ to constrained Θ space) and g is defined by forward_map which maps from Θ to the parameter space of the forward model, Φ.