From 0095bdc4f9ccf2fe87564504e98e62def2c10f2d Mon Sep 17 00:00:00 2001 From: Dan Giles Date: Sun, 29 Oct 2023 13:27:56 +0000 Subject: [PATCH 1/2] feat: Added support for the Lorenz 96 equations --- test/models/lorenz96.jl | 220 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 220 insertions(+) create mode 100644 test/models/lorenz96.jl diff --git a/test/models/lorenz96.jl b/test/models/lorenz96.jl new file mode 100644 index 00000000..d01586d0 --- /dev/null +++ b/test/models/lorenz96.jl @@ -0,0 +1,220 @@ +module Lorenz96 + +using Base.Threads +using Distributions +using FillArrays +using HDF5 +using Random +using PDMats +using OrdinaryDiffEq +using SciMLBase +using ParticleDA + + +Base.@kwdef struct Lorenz96ModelParameters{S <: Real, T <: Real} + F::S = 8.0 + time_step::S = 0.05 + N::Int = 40 + num_observed_indices::S = 10 + state_dimension::Int = 40 + observed_indices::Vector{Int} = collect(1:(state_dimension/num_observed_indices):state_dimension) + initial_state_std::Union{S, Vector{S}} = 0.1 + state_noise_std::Union{S, Vector{S}} = 0.1 + observation_noise_std::Union{T, Vector{T}} = 0.1 + operator_type::String = "linear" +end + +function get_params( + P::Type{Lorenz96ModelParameters{S, T}}, model_params_dict::Dict +) where {S <: Real, T <: Real} + return P(; (; (Symbol(k) => v for (k, v) in model_params_dict)...)...) +end + +struct Lorenz96Model{S <: Real, T <: Real} + parameters::Lorenz96ModelParameters{S, T} + integrators::Vector{<:SciMLBase.AbstractODEIntegrator} + initial_state_distribution::MvNormal{S} + state_noise_distribution::MvNormal{S} + observation_noise_distribution::MvNormal{T} +end + +function update_time_derivative!( + dx_dt::Vector{S}, x::Vector{S}, parameters::Lorenz96ModelParameters{S, T}, t::U +) where {S <: Real, T <: Real, U <: Real} + + for i in 1:parameters.N + if i == 1 + dx_dt[i] = (x[i+1] - x[parameters.N-1])*x[parameters.N] - x[i] + parameters.F + elseif i == 2 + dx_dt[i] = (x[i+1] - x[parameters.N])*x[i-1] - x[i] + parameters.F + elseif i == parameters.N + dx_dt[i] = (x[1] - x[i-2])*x[i-1] - x[i] + parameters.F + else + dx_dt[i] = (x[i+1] - x[i-2])x[i-1] - x[i] + parameters.F + + end + end +end + +function init(parameters_dict::Dict; S::Type{<:Real}=Float64, T::Type{<:Real}=Float64) + parameters = get_params(Lorenz96ModelParameters{S, T}, parameters_dict) + time_span = (0, parameters.time_step) + integrators = [ + OrdinaryDiffEq.init( + ODEProblem(update_time_derivative!, x, time_span, parameters), + Tsit5(); + save_everystep=false + ) + for x in eachcol(Matrix{S}(undef, 40, nthreads())) + ] + state_dimension = parameters.state_dimension + observation_dimension = length(parameters.observed_indices) + + return Lorenz96Model( + parameters, + integrators, + ( + MvNormal(m, isa(s, Vector) ? PDiagMat(s.^2) : ScalMat(length(m), s.^2)) + for (m, s) in ( + (Ones{S}(state_dimension), parameters.initial_state_std), + (Zeros{S}(state_dimension), parameters.state_noise_std), + (Zeros{T}(observation_dimension), parameters.observation_noise_std), + ) + )... + ) +end + +ParticleDA.get_state_dimension(::Lorenz96Model) = 40 +ParticleDA.get_observation_dimension(model::Lorenz96Model) = length( + model.parameters.observed_indices +) +ParticleDA.get_state_eltype(::Lorenz96Model{S, T}) where {S, T} = S +ParticleDA.get_observation_eltype(::Lorenz96Model{S, T}) where {S, T} = T + +function ParticleDA.sample_initial_state!( + state::AbstractVector{T}, + model::Lorenz96Model{S, T}, + rng::Random.AbstractRNG, +) where {S, T} + rand!(rng, model.initial_state_distribution, state) +end + +function ParticleDA.update_state_deterministic!( + state::AbstractVector{T}, + model::Lorenz96Model{S, T}, + time_index::Int, +) where {S, T} + reinit!(model.integrators[threadid()], state) + step!(model.integrators[threadid()], model.parameters.time_step, true) + state .= model.integrators[threadid()].u +end + +function ParticleDA.update_state_stochastic!( + state::AbstractVector{T}, + model::Lorenz96Model{S, T}, + rng::Random.AbstractRNG, +) where {S, T} + rand!(rng, state + model.state_noise_distribution, state) +end + + +function observation_operator!( + observation::AbstractVector{T}, + operator_type::String +) where {T <: Real} + if operator_type == "log" + observation .= log.(abs.(observation)) + elseif operator_type == "square" + observation .= (observation).^2 + else + observation .= observation + end +end + + +function ParticleDA.sample_observation_given_state!( + observation::AbstractVector{T}, + state::AbstractVector{S}, + model::Lorenz96Model{S, T}, + rng::Random.AbstractRNG, +) where {S <: Real, T <: Real} + + observation .= view(state, model.parameters.observed_indices) + + rand!( + rng, + observation_operator!(observation, model.parameters.operator_type) + + model.observation_noise_distribution, + observation + ) + +end + +function ParticleDA.get_log_density_observation_given_state( + observation::AbstractVector{T}, state::AbstractVector{S}, model::Lorenz96Model{S, T} +) where {S <: Real, T <: Real} + + obs_given_state = (view(state, model.parameters.observed_indices) + model.observation_noise_distribution) + observation_operator!(obs_given_state.μ, model.parameters.operator_type) + + return logpdf( + obs_given_state, + observation + ) +end + +function ParticleDA.write_model_metadata(file::HDF5.File, model::Lorenz96Model) + group_name = "parameters" + if !haskey(file, group_name) + group = create_group(file, group_name) + for field in fieldnames(typeof(model.parameters)) + value = getfield(model.parameters, field) + attributes(group)[string(field)] = ( + isa(value, AbstractVector) ? collect(value) : value + ) + end + else + @warn "Write failed, group $group_name already exists in $(file.filename)!" + end +end + +function ParticleDA.get_covariance_state_noise(model::Lorenz96Model) + return model.state_noise_distribution.Σ +end + +function ParticleDA.get_covariance_state_noise(model::Lorenz96Model, i::Int, j::Int) + return model.state_noise_distribution.Σ[i, j] +end + +function ParticleDA.get_covariance_observation_noise(model::Lorenz96Model) + return model.observation_noise_distribution.Σ +end + +function ParticleDA.get_covariance_observation_noise( + model::Lorenz96Model, i::Int, j::Int +) + return model.observation_noise_distribution.Σ[i, j] +end + +function ParticleDA.get_covariance_state_observation_given_previous_state( + model::Lorenz96Model, i::Int, j::Int +) + return model.state_noise_distribution.Σ[i, model.parameters.observed_indices[j]] +end + +function ParticleDA.get_covariance_observation_observation_given_previous_state( + model::Lorenz96Model, i::Int, j::Int +) + return ( + model.state_noise_distribution.Σ[ + model.parameters.observed_indices[i], model.parameters.observed_indices[j] + ] + + model.observation_noise_distribution.Σ[i, j] + ) +end + +function ParticleDA.get_state_indices_correlated_to_observations(model::Lorenz96Model) + return model.parameters.observed_indices +end + +end From 748e6f95bff0ce01985230f59c18b86edcbe886a Mon Sep 17 00:00:00 2001 From: Dan Giles Date: Fri, 3 Nov 2023 11:22:50 +0000 Subject: [PATCH 2/2] feat: Nonlinear observation operators --- test/models/lorenz63.jl | 22 +++++++++++++++++++--- test/models/lorenz96.jl | 2 -- 2 files changed, 19 insertions(+), 5 deletions(-) diff --git a/test/models/lorenz63.jl b/test/models/lorenz63.jl index 9cfd1e9e..ebebed04 100644 --- a/test/models/lorenz63.jl +++ b/test/models/lorenz63.jl @@ -19,6 +19,7 @@ Base.@kwdef struct Lorenz63ModelParameters{S <: Real, T <: Real} initial_state_std::Union{S, Vector{S}} = 0.05 state_noise_std::Union{S, Vector{S}} = 0.05 observation_noise_std::Union{T, Vector{T}} = 2. + operator_type::String = "linear" end function get_params( @@ -103,15 +104,28 @@ function ParticleDA.update_state_stochastic!( rand!(rng, state + model.state_noise_distribution, state) end +function observation_operator!( + observation::AbstractVector{T}, + operator_type::String +) where {T <: Real} + if operator_type == "log" + observation .= log.(abs.(observation)) + else + observation .= observation + end +end + function ParticleDA.sample_observation_given_state!( observation::AbstractVector{T}, state::AbstractVector{S}, model::Lorenz63Model{S, T}, rng::Random.AbstractRNG, ) where {S <: Real, T <: Real} + + observation .= view(state, model.parameters.observed_indices) rand!( rng, - view(state, model.parameters.observed_indices) + observation_operator!(observation, model.parameters.operator_type) + model.observation_noise_distribution, observation ) @@ -120,9 +134,11 @@ end function ParticleDA.get_log_density_observation_given_state( observation::AbstractVector{T}, state::AbstractVector{S}, model::Lorenz63Model{S, T} ) where {S <: Real, T <: Real} + + obs_given_state = (view(state, model.parameters.observed_indices) + model.observation_noise_distribution) + observation_operator!(obs_given_state.μ, model.parameters.operator_type) return logpdf( - view(state, model.parameters.observed_indices) - + model.observation_noise_distribution, + obs_given_state, observation ) end diff --git a/test/models/lorenz96.jl b/test/models/lorenz96.jl index d01586d0..884db0dc 100644 --- a/test/models/lorenz96.jl +++ b/test/models/lorenz96.jl @@ -124,8 +124,6 @@ function observation_operator!( ) where {T <: Real} if operator_type == "log" observation .= log.(abs.(observation)) - elseif operator_type == "square" - observation .= (observation).^2 else observation .= observation end