Implement a DataUpdateCallback (#287)

* Implement a data update likelihood callback

* Make the likelihood accumulator optional

* Make some DataUpdateCallback args into kwargs

* WIP

* Fix the partial observability implementation

* Better type signatures

* Misc

* Add Fenrir to ProbNumDiffEq.jl

* Add proper data likelihood tests

* Split the `data_likelihoods.jl` file into separate files per method

* Fix a bug

* Remove the smooth=false suggestion if dense=false

* Fix another broken test (again?)

* Revert the ManifoldUpdate renaming

* Actually remove the filtering likelihood from the dalton file

* Try out DocStringExtensions.jl

* Add compat entry to DocStringExtensions

* Create a DataLikelihoods submodule

* Remove Fenrir from the docs

* Write docstrings for the fenrir and dalton likelihood and doc them

* Make marginalize a bit more flexible

* Update the Probabilistic Exponential Integrator citation

* Remove the underscores from the likelihoods again

* Add a parameter inference tutorial again

* Update the doc index

* Make the parameter inference example even nicer

* Improve tests

* Make Fenrir compatible with matrix-valued observation noise

* Test matrix-valued observation noise

* JuliaFormatter.jl

* Add support for PSDMatrices as observation noise

* Fix a MarkovKernel docstring

* Implement and test update equations with non-zero observation noise

* Shorten and streamline the `update!` functionality a bit

* Simplify Fenrir quite a bit

* Misc updates to the data update callback

* Make the Fenrir code yet a bit more compact

* JuliaFormatter.jl

* Change the parameter inference example to partial observations

* Add DiffEqCallbacks compat entry

* Make the parameter inference doc code a bit nicer

* Add docstrings for DataUpdateLogLikelihood and DataUpdateCallback

* Polish the docs

* Make sure that Fenrir ll is not inf, even if it would technically be

* JuliaFormatter.jl

* Faster DataUpdateCallback

* Slight fenrir speed improvement

* Remove some changes that I somehow introduced earlier

* Add the data likelihood tests to runtests.jl

* Fix the failing data likelihood tests

* JuliaFormatter.jl

* Test for more data likelihood observation noise types

Author: nathanaelbosch

Project.toml

@@ -6,7 +6,9 @@ version = "0.14.0"
 [deps]
 ArrayAllocators = "c9d4266f-a5cb-439d-837c-c97b191379f5"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 ExponentialUtilities = "d4d017d3-3776-5f7e-afef-a10c40355c18"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
@@ -47,7 +49,9 @@ RecipesBaseExt = "RecipesBase"
 [compat]
 ArrayAllocators = "0.3"
 DiffEqBase = "6.122"
+DiffEqCallbacks = "2.36"
 DiffEqDevTools = "2"
+DocStringExtensions = "0.9"
 ExponentialUtilities = "1"
 FastBroadcast = "0.2"
 FastGaussQuadrature = "0.5, 1"

docs/Project.toml

@@ -3,7 +3,6 @@ Bibliography = "f1be7e48-bf82-45af-a471-ae754a193061"
 DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
-Fenrir = "e9b4b195-f5cd-427c-8076-5358c553c37f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"

docs/make.jl

@@ -34,18 +34,19 @@ makedocs(
     pages=[
         "Home" => "index.md",
         "Tutorials" => [
-            "Getting Started" => "tutorials/getting_started.md"
-            "Second Order ODEs and Energy Preservation" => "tutorials/dynamical_odes.md"
-            "Differential Algebraic Equations" => "tutorials/dae.md"
-            "Probabilistic Exponential Integrators" => "tutorials/exponential_integrators.md"
-            "Parameter Inference" => "tutorials/fenrir.md"
+            "Getting Started" => "tutorials/getting_started.md",
+            "Second Order ODEs and Energy Preservation" => "tutorials/dynamical_odes.md",
+            "Differential Algebraic Equations" => "tutorials/dae.md",
+            "Probabilistic Exponential Integrators" => "tutorials/exponential_integrators.md",
+            "Parameter Inference" => "tutorials/ode_parameter_inference.md",
         ],
         "Solvers and Options" => [
             "solvers.md",
             "priors.md",
             "initialization.md",
             "diffusions.md",
         ],
+        "Data Likelihoods" => "likelihoods.md",
         "Benchmarks" => [
             "Multi-Language Wrapper Benchmark" => "benchmarks/multi-language-wrappers.md",
             "Non-stiff ODEs" => [
@@ -64,8 +65,8 @@ makedocs(
             ],
         ],
         "Internals" => [
-            "Filtering and Smoothing" => "filtering.md"
-            "Implementation via OrdinaryDiffEq.jl" => "implementation.md"
+            "Filtering and Smoothing" => "filtering.md",
+            "Implementation via OrdinaryDiffEq.jl" => "implementation.md",
         ],
         "References" => "references.md",
     ],

docs/src/index.md

@@ -14,7 +14,7 @@ Run Julia, enter `]` to bring up Julia's package manager, and add the ProbNumDif
 
 ```
 julia> ]
-(v1.9) pkg> add ProbNumDiffEq
+(v1.10) pkg> add ProbNumDiffEq
 ```
 
 ## Getting Started
@@ -35,10 +35,10 @@ For a quick introduction check out the "[Solving ODEs with Probabilistic Numeric
 - Arbitrary precision via Julia's built-in [arbitrary precision arithmetic](https://docs.julialang.org/en/v1/manual/integers-and-floating-point-numbers/#Arbitrary-Precision-Arithmetic)
 - Specialized solvers for second-order ODEs (see [Second Order ODEs and Energy Preservation](@ref))
 - Compatible with DAEs in mass-matrix ODE form (see [Solving DAEs with Probabilistic Numerics](@ref))
+- Data likelihoods for parameter-inference in ODEs (see [Parameter Inference with ProbNumDiffEq.jl](@ref))
 
 
 ## Related packages
 
 - [probdiffeq](https://pnkraemer.github.io/probdiffeq/): Fast and feature-rich filtering-based probabilistic ODE solvers in JAX.
 - [ProbNum](https://probnum.readthedocs.io/en/latest/): Probabilistic numerics in Python. It has not only probabilistic ODE solvers, but also probabilistic linear solvers, Bayesian quadrature, and many filtering and smoothing implementations.
-- [Fenrir.jl](https://github.com/nathanaelbosch/Fenrir.jl): Parameter-inference in ODEs with probabilistic ODE solvers. This package builds on ProbNumDiffEq.jl to provide a negative marginal log-likelihood function, which can then be used with an optimizer or with MCMC for parameter inference.

docs/src/likelihoods.md

@@ -0,0 +1,7 @@
+# Data Likelihoods
+
+
+```@docs
+ProbNumDiffEq.DataLikelihoods.fenrir_data_loglik
+ProbNumDiffEq.DataLikelihoods.dalton_data_loglik
+```

docs/src/refs.bib

@@ -180,4 +180,16 @@ @book{sarkka19appliedsde
   author =       {Särkkä, Simo and Solin, Arno},
   year =         2019,
   collection =   {Institute of Mathematical Statistics Textbooks}
-}
+}
+
+@article{wu23dalton,
+  author =       {Mohan Wu and Martin Lysy},
+  title =        {Data-Adaptive Probabilistic Likelihood Approximation for
+                  Ordinary Differential Equations},
+  journal =      {CoRR},
+  year =         2023,
+  url =          {http://arxiv.org/abs/2306.05566},
+  archivePrefix ={arXiv},
+  eprint =       {2306.05566},
+  primaryClass = {stat.ML}
+}

docs/src/tutorials/fenrir.md …src/tutorials/ode_parameter_inference.mddocs/src/tutorials/fenrir.md renamed to docs/src/tutorials/ode_parameter_inference.md docs/src/tutorials/fenrir.md renamed to docs/src/tutorials/ode_parameter_inference.md

@@ -1,19 +1,7 @@
-# Parameter Inference with ProbNumDiffEq.jl and Fenrir.jl
+# Parameter Inference with ProbNumDiffEq.jl
 
-!!! note
-    This is mostly just a copy from [the tutorial included in the Fenrir.jl documentation](https://nathanaelbosch.github.io/Fenrir.jl/stable/gettingstarted/), so have a look there too!
 
 
-```@example fenrir
-using LinearAlgebra
-using OrdinaryDiffEq, ProbNumDiffEq, Plots
-using Fenrir
-using Optimization, OptimizationOptimJL
-stack(x) = copy(reduce(hcat, x)') # convenient
-nothing # hide
-```
-
-## The parameter inference problem in general
 Let's assume we have an initial value problem (IVP)
 ```math
 \begin{aligned}
@@ -29,12 +17,15 @@ u(t_n) = H y(t_n) + v_n, \qquad v_n \sim \mathcal{N}(0, R).
 The question of interest is: How can we compute the marginal likelihood ``p(\mathcal{D} \mid \theta)``?
 Short answer: We can't. It's intractable, because computing the true IVP solution exactly ``y(t)`` is intractable.
 What we can do however is compute an approximate marginal likelihood.
-This is what Fenrir.jl provides.
-For details, check out the [paper](https://proceedings.mlr.press/v162/tronarp22a.html).
+This is what `ProbNumDiffEq.DataLikelihoods` provides.
 
 ## The specific problem, in code
 Let's assume that the true underlying dynamics are given by a FitzHugh-Nagumo model
-```@example fenrir
+
+```@example parameterinference
+using ProbNumDiffEq, LinearAlgebra, OrdinaryDiffEq, Plots
+Plots.theme(:default; markersize=2, markerstrokewidth=0.1)
+
 function f(du, u, p, t)
     a, b, c = p
     du[1] = c*(u[1] - u[1]^3/3 + u[2])
@@ -46,30 +37,43 @@ p = (0.2, 0.2, 3.0)
 true_prob = ODEProblem(f, u0, tspan, p)
 ```
 from which we generate some artificial noisy data
-```@example fenrir
+```@example parameterinference
 true_sol = solve(true_prob, Vern9(), abstol=1e-10, reltol=1e-10)
 
 times = 1:0.5:20
-observation_noise_var = 1e-1
-odedata = [true_sol(t) .+ sqrt(observation_noise_var) * randn(length(u0)) for t in times]
+σ = 1e-1
+H = [1 0;]
+odedata = [H*true_sol(t) .+ σ * randn() for t in times]
 
 plot(true_sol, color=:black, linestyle=:dash, label=["True Solution" ""])
-scatter!(times, stack(odedata), markersize=2, markerstrokewidth=0.1, color=1, label=["Noisy Data" ""])
+scatter!(times, stack(odedata)',  color=1, label=["Noisy Data" ""])
 ```
 Our goal is then to recover the true parameter `p` (and thus also the true trajectory plotted above) the noisy data.
 
 ## Computing the negative log-likelihood
-To do parameter inference - be it maximum-likelihod, maximum a posteriori, or full Bayesian inference with MCMC - we need to evaluate the likelihood of given a parameter estimate ``\theta_\text{est}``.
-This is exactly what Fenrir.jl's [`fenrir_nll`](https://nathanaelbosch.github.io/Fenrir.jl/stable/#Fenrir.fenrir_nll) provides:
-```@example fenrir
-p_est = (0.1, 0.1, 2.0)
-prob = remake(true_prob, p=p_est)
+To do parameter inference - be it maximum-likelihod, maximum a posteriori, or full Bayesian inference with MCMC - we need to evaluate the likelihood of given a parameter estimate ``\theta_\text{est}``, which corresponds to the probability of the data under the trajectory returned by the ODE solver
+```@example parameterinference
+θ_est = (0.1, 0.1, 2.0)
+prob = remake(true_prob, p=θ_est)
+plot(true_sol, color=:black, linestyle=:dash, label=["True Solution" ""])
+scatter!(times, stack(odedata)', color=1, label=["Noisy Data" ""])
+sol = solve(prob, EK1(), adaptive=false, dt=1e-1)
+plot!(sol, color=2, label=["Numerical solution for θ_est" ""])
+```
+This quantity can be computed in multiple ways; see 
+[Data Likelihoods](@ref).
+Here we use 
+[`ProbNumDiffEq.DataLikelihoods.fenrir_data_loglik`](@ref):
+```@example parameterinference
+using ProbNumDiffEq.DataLikelihoods
+
 data = (t=times, u=odedata)
-κ² = 1e10
-nll, _, _ = fenrir_nll(prob, data, observation_noise_var, κ²; dt=1e-1)
-nll
+nll = -fenrir_data_loglik(
+    prob, EK1(smooth=true); 
+    data, observation_noise_cov=σ^2, observation_matrix=H, 
+    adaptive=false, dt=1e-1)
 ```
-This is the negative marginal log-likelihood of the parameter `p_est`.
+This is the negative marginal log-likelihood of the parameter `θ_est`.
 You can use it as any other NLL: Optimize it to compute maximum-likelihood estimates or MAPs, or plug it into MCMC to sample from the posterior.
 In our paper [tronarp22fenrir](@cite) we compute MLEs by pairing Fenrir with [Optimization.jl](http://optimization.sciml.ai/stable/) and [ForwardDiff.jl](https://juliadiff.org/ForwardDiff.jl/stable/).
 Let's quickly explore how to do this next.
@@ -80,23 +84,29 @@ Let's quickly explore how to do this next.
 To compute a maximum-likelihood estimate (MLE), we just need to maximize ``\theta \to p(\mathcal{D} \mid \theta)`` - that is, minimize the `nll` from above.
 We use [Optimization.jl](https://docs.sciml.ai/Optimization/stable/) for this.
 First, define a loss function and create an `OptimizationProblem`
-```@example fenrir
+```@example parameterinference
+using Optimization, OptimizationOptimJL
+
 function loss(x, _)
     ode_params = x[begin:end-1]
     prob = remake(true_prob, p=ode_params)
-    κ² = exp(x[end]) # the diffusion parameter of the EK1
-    return fenrir_nll(prob, data, observation_noise_var, κ²; dt=1e-1)
+    κ² = exp(x[end]) # we also optimize the diffusion parameter of the EK1
+    return -fenrir_data_loglik(
+        prob, EK1(smooth=true, diffusionmodel=FixedDiffusion(κ², false));
+        data, observation_noise_cov=σ^2, observation_matrix=H,
+        adaptive=false, dt=1e-1
+    )
 end
 
 fun = OptimizationFunction(loss, Optimization.AutoForwardDiff())
 optprob = OptimizationProblem(
-    fun, [p_est..., 1e0];
+    fun, [θ_est..., 1e0];
     lb=[0.0, 0.0, 0.0, -10], ub=[1.0, 1.0, 5.0, 20] # lower and upper bounds
 )
 ```
 
 Then, just `solve` it! Here we use LBFGS:
-```@example fenrir
+```@example parameterinference
 optsol = solve(optprob, LBFGS())
 p_mle = optsol.u[1:3]
 p_mle # hide
@@ -105,21 +115,27 @@ p_mle # hide
 Success! The computed MLE is quite close to the true parameter which we used to generate the data.
 As a final step, let's plot the true solution, the data, and the result of the MLE:
 
-```@example fenrir
+```@example parameterinference
 plot(true_sol, color=:black, linestyle=:dash, label=["True Solution" ""])
-scatter!(times, stack(odedata), markersize=2, markerstrokewidth=0.1, color=1, label=["Noisy Data" ""])
+scatter!(times, stack(odedata)', color=1, label=["Noisy Data" ""])
 mle_sol = solve(remake(true_prob, p=p_mle), EK1())
 plot!(mle_sol, color=3, label=["MLE-parameter Solution" ""])
 ```
 
 Looks good!
 
 
-### Reference
+## API Documentation
+
+For more details, see the API documentation of `ProbNumDiffEq.DataLikelihoods` at [Data Likelihoods](@ref).
+
+
+### References
 
 ```@bibliography
 Pages = []
 Canonical = false
 
 tronarp22fenrir
+wu23dalton
 ```

src/ProbNumDiffEq.jl

@@ -8,6 +8,7 @@ using LinearAlgebra
 import LinearAlgebra: mul!
 import Statistics: mean, var, std, cov
 using Printf
+using DocStringExtensions
 
 using Reexport
 @reexport using DiffEqBase
@@ -31,6 +32,7 @@ using ArrayAllocators
 using FiniteHorizonGramians
 using FillArrays
 using MatrixEquations
+using DiffEqCallbacks
 
 @reexport using GaussianDistributions
 
@@ -95,8 +97,18 @@ if !isdefined(Base, :get_extension)
     include("../ext/DiffEqDevToolsExt.jl")
 end
 
-include("callbacks.jl")
+include("callbacks/manifoldupdate.jl")
 export ManifoldUpdate
+include("callbacks/dataupdate.jl")
+export DataUpdateLogLikelihood, DataUpdateCallback
+
+include("data_likelihoods/dalton.jl")
+include("data_likelihoods/filtering.jl")
+include("data_likelihoods/fenrir.jl")
+module DataLikelihoods
+import ..ProbNumDiffEq: dalton_data_loglik, filtering_data_loglik, fenrir_data_loglik
+export dalton_data_loglik, filtering_data_loglik, fenrir_data_loglik
+end
 
 include("precompile.jl")
 

src/algorithms.jl

@@ -187,7 +187,7 @@ julia> solve(prob, RosenbrockExpEK())
 ```
 
 # Reference
-* [bosch23expint](@cite) Bosch et al, "Probabilistic Exponential Integrators", arXiv (2021)
+* [bosch23expint](@cite) Bosch et al, "Probabilistic Exponential Integrators", NeurIPS (2023)
 """
 RosenbrockExpEK(; order=3, kwargs...) =
     EK1(; prior=IOUP(order, update_rate_parameter=true), kwargs...)

src/caches.jl

@@ -168,7 +168,7 @@ function OrdinaryDiffEq.alg_cache(
     copy!(x0.Σ, apply_diffusion(x0.Σ, initdiff))
 
     # Measurement model related things
-    R = factorized_similar(FAC, d, d)
+    R = nothing # factorized_similar(FAC, d, d)
     H = factorized_similar(FAC, d, D)
     v = similar(Array{uElType}, d)
     S = PSDMatrix(factorized_zeros(FAC, D, d))