Fitting parameters of a sinusoid

In this example we have a model that produces a sinusoid $f(A, v) = A \sin(\phi + t) + v, \forall t \in [0,2\pi]$, with a random phase $\phi$. Given an initial guess of the parameters as $A^* \sim \mathcal{N}(2,1)$ and $v^* \sim \mathcal{N}(0,25)$, our goal is to estimate the parameters from a noisy observation of the maximum, minimum, and mean of the true model output.

First, we load the packages we need:

using LinearAlgebra, Random

using Distributions, Plots

using EnsembleKalmanProcesses
using EnsembleKalmanProcesses.ParameterDistributions
const EKP = EnsembleKalmanProcesses

# Setting up the model and data for our inverse problem

Now, we define a model which generates a sinusoid given parameters $\theta$: an amplitude and a vertical shift. We will estimate these parameters from data. The model adds a random phase shift upon evaluation.

dt = 0.01
trange = 0:dt:(2 * pi + dt)
function model(amplitude, vert_shift)
    phi = 2 * pi * rand(rng)
    return amplitude * sin.(trange .+ phi) .+ vert_shift
end

Seed for pseudo-random number generator.

rng_seed = 41
rng = Random.MersenneTwister(rng_seed)

We then define $G(\theta)$, which returns the observables of the sinusoid given a parameter vector. These observables should be defined such that they are informative about the parameters we wish to estimate. Here, the two observables are the $y$ range of the curve (which is informative about its amplitude), as well as its mean (which is informative about its vertical shift).

function G(u)
    theta, vert_shift = u
    sincurve = model(theta, vert_shift)
    return [maximum(sincurve) - minimum(sincurve), mean(sincurve)]
end

Suppose we have a noisy observation of the true system. Here, we create a pseudo-observation $y$ by running our model with the correct parameters and adding Gaussian noise to the output.

dim_output = 2

Γ = 0.1 * I
noise_dist = MvNormal(zeros(dim_output), Γ)

theta_true = [1.0, 7.0]
y = G(theta_true) .+ rand(noise_dist)

# Solving the inverse problem

We now define prior distributions on the two parameters. For the amplitude, we define a prior with mean 2 and standard deviation 1. It is additionally constrained to be nonnegative. For the vertical shift we define a Gaussian prior with mean 0 and standard deviation 5.

prior_u1 = constrained_gaussian("amplitude", 2, 1, 0, Inf)
prior_u2 = constrained_gaussian("vert_shift", 0, 5, -Inf, Inf)
prior = combine_distributions([prior_u1, prior_u2])

We now generate the initial ensemble and set up the ensemble Kalman inversion. We define an ensemble size and we define a maximum iteration for this experiment (depending on the timestepper, early termination criteria can be used). We highlight that a key benefit of many variants of the Kalman approach over other particle methods, is that typically the ensemble size does not need to scale (e.g., linearly) with the number of parameters.

N_ensemble = 20
N_iterations = 10

initial_ensemble = EKP.construct_initial_ensemble(rng, prior, N_ensemble)

ensemble_kalman_process = EKP.EnsembleKalmanProcess(initial_ensemble, y, Γ, Inversion(); rng = rng)

We are now ready to carry out the inversion. At each iteration, we get the ensemble from the last iteration, apply $G(\theta)$ to each ensemble member, and apply the Kalman update to the ensemble. We also check if termination criteria has been exceeded, and break the loop if it has - on termination, the update is not performed.

for i in 1:N_iterations
    params_i = get_ϕ_final(prior, ensemble_kalman_process)

    G_ens = hcat([G(params_i[:, i]) for i in 1:N_ensemble]...)

    terminate = EKP.update_ensemble!(ensemble_kalman_process, G_ens)
    if !(isnothing(terminate))
        @info "Termination at iteration $(i-1)"
        break
    end
end

[ Info: Termination condition of scheduler `DataMisfitController` will be exceeded during the next iteration.
┌ Warning: Termination condition of scheduler `DataMisfitController` has been exceeded, returning `true` from `update_ensemble!` and preventing futher updates
│  Set on_terminate="continue" in `DataMisfitController` to ignore termination
└ @ EnsembleKalmanProcesses ~/work/EnsembleKalmanProcesses.jl/EnsembleKalmanProcesses.jl/src/LearningRateSchedulers.jl:293
[ Info: Termination at iteration 5

Finally, we get the ensemble after the last iteration. This provides our estimate of the parameters.

final_ensemble = get_ϕ_final(prior, ensemble_kalman_process)

2×20 Matrix{Float64}:
 1.18627  1.18315  1.15203  1.18309  …  1.17191  1.15747  1.1845  1.174
 7.62926  7.62426  7.79582  7.4807      7.4908   7.82983  7.714   7.60742

To visualize the success of the inversion, we plot model with the true parameters, the initial ensemble, and the final ensemble.

p = plot(trange, model(theta_true...), c = :black, label = "Truth", legend = :bottomright, linewidth = 2)

plot!(p, trange, [model(get_ϕ(prior, ensemble_kalman_process, 1)[:, 1]...)], c = :red, label = "Initial ensemble")
plot!(p, trange, [model(final_ensemble[:, 1]...)], c = :blue, label = "Final ensemble")
plot!(p, trange, [model(get_ϕ(prior, ensemble_kalman_process, 1)[:, i]...) for i in 2:N_ensemble], c = :red, label = "")
plot!(p, trange, [model(final_ensemble[:, i]...) for i in 2:N_ensemble], c = :blue, label = "")

xlabel!("Time")

We see that the final ensemble is much closer to the truth. Note that the random phase shift is of no consequence.

savefig(p, "output.png")

"/home/runner/work/EnsembleKalmanProcesses.jl/EnsembleKalmanProcesses.jl/docs/build/literated/output.png"

This page was generated using Literate.jl.