From 9cc94eb49aeaee13f65299e204514cc99fc743b4 Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Thu, 24 Apr 2025 16:36:30 -0400 Subject: [PATCH 1/9] update to most recent SSMProblems interface --- .gitignore | 2 +- Project.toml | 1 - examples/gaussian-process/script.jl | 105 +++++++++++++++------------- examples/gaussian-ssm/script.jl | 88 ++++++++--------------- examples/particle-gibbs/script.jl | 105 ++++++++++------------------ src/model.jl | 26 +++++++ src/pgas.jl | 27 +++---- test/Project.toml | 1 - test/container.jl | 28 ++++---- test/linear-gaussian.jl | 55 +++++++-------- test/pgas.jl | 51 ++++++++------ 11 files changed, 232 insertions(+), 257 deletions(-) diff --git a/.gitignore b/.gitignore index c5ff299b..482e8044 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,5 @@ *.jl.cov *.jl.*.cov *.jl.mem -/Manifest.toml +Manifest.toml /test/Manifest.toml diff --git a/Project.toml b/Project.toml index 7641cce5..2813c961 100644 --- a/Project.toml +++ b/Project.toml @@ -26,7 +26,6 @@ Random = "<0.0.1, 1" Random123 = "1.3" Requires = "1.0" StatsFuns = "0.9, 1" -SSMProblems = "0.1" julia = "1.7" [extras] diff --git a/examples/gaussian-process/script.jl b/examples/gaussian-process/script.jl index 1a58bde3..34db890c 100644 --- a/examples/gaussian-process/script.jl +++ b/examples/gaussian-process/script.jl @@ -8,74 +8,83 @@ using Distributions using Libtask using SSMProblems -Parameters = @NamedTuple begin - a::Float64 - q::Float64 - kernel +# Gaussian process encoded transition dynamics +mutable struct GaussianProcessDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T} + proc::AbstractGPs.AbstractGP + q::T + function GaussianProcessDynamics(q::T, kernel::KT) where {T<:Real,KT<:Kernel} + return new{T}(GP(ZeroMean{T}(), kernel), q) + end end -mutable struct GPSSM <: SSMProblems.AbstractStateSpaceModel - X::Vector{Float64} - observations::Vector{Float64} - θ::Parameters - - GPSSM(params::Parameters) = new(Vector{Float64}(), params) - GPSSM(y::Vector{Float64}, params::Parameters) = new(Vector{Float64}(), y, params) +function SSMProblems.distribution(dyn::GaussianProcessDynamics{T}) where {T<:Real} + return Normal(zero(T), dyn.q) end -seed = 1 -T = 100 -Nₚ = 20 -Nₛ = 250 -a = 0.9 -q = 0.5 - -params = Parameters((a, q, SqExponentialKernel())) +# TODO: broken... +function SSMProblems.simulate( + rng::AbstractRNG, dyn::GaussianProcessDynamics, step::Int, state +) + dyn.proc = posterior(dyn.proc(step:step), [state]) + μ, σ = mean_and_cov(dyn.proc, [step]) + return rand(rng, Normal(μ[1], sqrt(σ[1]))) +end -f(θ::Parameters, x, t) = Normal(θ.a * x, θ.q) -h(θ::Parameters) = Normal(0, θ.q) -g(θ::Parameters, x, t) = Normal(0, exp(0.5 * x)^2) +function SSMProblems.logdensity(dyn::GaussianProcessDynamics, step::Int, state, prev_state) + μ, σ = mean_and_cov(dyn.proc, [step]) + return logpdf(Normal(μ, sqrt(σ)), state) +end -rng = Random.MersenneTwister(seed) +# Linear Gaussian dynamics used for simulation +struct LinearGaussianDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T} + a::T + q::T +end -x = zeros(T) -y = similar(x) -x[1] = rand(rng, h(params)) -for t in 1:T - if t < T - x[t + 1] = rand(rng, f(params, x[t], t)) - end - y[t] = rand(rng, g(params, x[t], t)) +function SSMProblems.distribution(dyn::LinearGaussianDynamics{T}) where {T<:Real} + return Normal(zero(T), dyn.q) end -function gp_update(model::GPSSM, state, step) - gp = GP(model.θ.kernel) - prior = gp(1:(step - 1)) - post = posterior(prior, model.X[1:(step - 1)]) - μ, σ = mean_and_cov(post, [step]) - return Normal(μ[1], σ[1]) +function SSMProblems.distribution(dyn::LinearGaussianDynamics, ::Int, state) + return Normal(dyn.a * state, dyn.q) end -SSMProblems.transition!!(rng::AbstractRNG, model::GPSSM) = rand(rng, h(model.θ)) -function SSMProblems.transition!!(rng::AbstractRNG, model::GPSSM, state, step) - return rand(rng, gp_update(model, state, step)) +# Observation process used in both variants of the model +struct StochasticVolatility{T<:Real} <: SSMProblems.ObservationProcess{T,T} end + +function SSMProblems.distribution(::StochasticVolatility{T}, ::Int, state) where {T<:Real} + return Normal(zero(T), exp((1 / 2) * state)) end -function SSMProblems.emission_logdensity(model::GPSSM, state, step) - return logpdf(g(model.θ, state, step), model.observations[step]) +# Baseline model (for simulation) +function LinearGaussianStochasticVolatilityModel(a::T, q::T) where {T<:Real} + dyn = LinearGaussianDynamics(a, q) + obs = StochasticVolatility{T}() + return SSMProblems.StateSpaceModel(dyn, obs) end -function SSMProblems.transition_logdensity(model::GPSSM, prev_state, current_state, step) - return logpdf(gp_update(model, prev_state, step), current_state) + +# Gaussian process model (for sampling) +function GaussianProcessStateSpaceModel(q::T, kernel::KT) where {T<:Real,KT<:Kernel} + dyn = GaussianProcessDynamics(q, kernel) + obs = StochasticVolatility{T}() + return SSMProblems.StateSpaceModel(dyn, obs) end -AdvancedPS.isdone(::GPSSM, step) = step > T +# Everything is now ready to simulate some data. +rng = Random.MersenneTwister(1234) +true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5) +_, x, y = sample(rng, true_model, 100); -model = GPSSM(y, params) -pg = AdvancedPS.PGAS(Nₚ) -chains = sample(rng, model, pg, Nₛ) +# Create the model and run the sampler +gpssm = GaussianProcessStateSpaceModel(0.5, SqExponentialKernel()) +model = gpssm(y) +pg = AdvancedPS.PGAS(20) +chains = sample(rng, model, pg, 250) +#md nothing #hide particles = hcat([chain.trajectory.model.X for chain in chains]...) mean_trajectory = mean(particles; dims=2); +#md nothing #hide scatter(particles; label=false, opacity=0.01, color=:black, xlabel="t", ylabel="state") plot!(x; color=:darkorange, label="Original Trajectory") diff --git a/examples/gaussian-ssm/script.jl b/examples/gaussian-ssm/script.jl index d7a5889c..e8fd9ad3 100644 --- a/examples/gaussian-ssm/script.jl +++ b/examples/gaussian-ssm/script.jl @@ -28,71 +28,46 @@ using SSMProblems # as well as the initial distribution $f_0(x) = \mathcal{N}(0, q^2/(1-a^2))$. # To use `AdvancedPS` we first need to define a model type that subtypes `AdvancedPS.AbstractStateSpaceModel`. -Parameters = @NamedTuple begin - a::Float64 - q::Float64 - r::Float64 +mutable struct Parameters{T<:Real} + a::T + q::T + r::T end -mutable struct LinearSSM <: SSMProblems.AbstractStateSpaceModel - X::Vector{Float64} - observations::Vector{Float64} - θ::Parameters - LinearSSM(θ::Parameters) = new(Vector{Float64}(), θ) - LinearSSM(y::Vector, θ::Parameters) = new(Vector{Float64}(), y, θ) +struct LinearGaussianDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T} + a::T + q::T end -# and the densities defined above. -f(θ::Parameters, state, t) = Normal(θ.a * state, θ.q) # Transition density -g(θ::Parameters, state, t) = Normal(state, θ.r) # Observation density -f₀(θ::Parameters) = Normal(0, θ.q^2 / (1 - θ.a^2)) # Initial state density -#md nothing #hide +function SSMProblems.distribution(dyn::LinearGaussianDynamics{T}; kwargs...) where {T<:Real} + return Normal(zero(T), sqrt(dyn.q^2 / (1 - dyn.a^2))) +end -# We also need to specify the dynamics of the system through the transition equations: -# - `AdvancedPS.initialization`: the initial state density -# - `AdvancedPS.transition`: the state transition density -# - `AdvancedPS.observation`: the observation score given the observed data -# - `AdvancedPS.isdone`: signals the end of the execution for the model -SSMProblems.transition!!(rng::AbstractRNG, model::LinearSSM) = rand(rng, f₀(model.θ)) -function SSMProblems.transition!!( - rng::AbstractRNG, model::LinearSSM, state::Float64, step::Int -) - return rand(rng, f(model.θ, state, step)) +function SSMProblems.distribution(dyn::LinearGaussianDynamics, step::Int, state; kwargs...) + return Normal(dyn.a * state, dyn.q) end -function SSMProblems.emission_logdensity(modeL::LinearSSM, state::Float64, step::Int) - return logpdf(g(model.θ, state, step), model.observations[step]) +struct LinearGaussianObservation{T<:Real} <: SSMProblems.ObservationProcess{T,T} + r::T end -function SSMProblems.transition_logdensity( - model::LinearSSM, prev_state, current_state, step + +function SSMProblems.distribution( + obs::LinearGaussianObservation, step::Int, state; kwargs... ) - return logpdf(f(model.θ, prev_state, step), current_state) + return Normal(state, obs.r) end -# We need to think seriously about how the data is handled -AdvancedPS.isdone(::LinearSSM, step) = step > Tₘ +function LinearGaussianStateSpaceModel(θ::Parameters) + dyn = LinearGaussianDynamics(θ.a, θ.q) + obs = LinearGaussianObservation(θ.r) + return SSMProblems.StateSpaceModel(dyn, obs) +end # Everything is now ready to simulate some data. -a = 0.9 # Scale -q = 0.32 # State variance -r = 1 # Observation variance -Tₘ = 200 # Number of observation -Nₚ = 20 # Number of particles -Nₛ = 500 # Number of samples -seed = 1 # Reproduce everything - -θ₀ = Parameters((a, q, r)) -rng = Random.MersenneTwister(seed) - -x = zeros(Tₘ) -y = zeros(Tₘ) -x[1] = rand(rng, f₀(θ₀)) -for t in 1:Tₘ - if t < Tₘ - x[t + 1] = rand(rng, f(θ₀, x[t], t)) - end - y[t] = rand(rng, g(θ₀, x[t], t)) -end +rng = Random.MersenneTwister(1234) +θ = Parameters(0.9, 0.32, 1.0) +true_model = LinearGaussianStateSpaceModel(θ) +_, x, y = sample(rng, true_model, 200); # Here are the latent and obseravation timeseries plot(x; label="x", xlabel="t") @@ -100,9 +75,8 @@ plot!(y; seriestype=:scatter, label="y", xlabel="t", mc=:red, ms=2, ma=0.5) # `AdvancedPS` subscribes to the `AbstractMCMC` API. To sample we just need to define a Particle Gibbs kernel # and a model interface. -model = LinearSSM(y, θ₀) -pgas = AdvancedPS.PGAS(Nₚ) -chains = sample(rng, model, pgas, Nₛ; progress=false); +pgas = AdvancedPS.PGAS(20) +chains = sample(rng, true_model(y), pgas, 500; progress=false); #md nothing #hide # @@ -118,7 +92,7 @@ plot!(mean_trajectory; color=:dodgerblue, label="Mean trajectory", opacity=0.9) # We used a particle gibbs kernel with the ancestor updating step which should help with the particle # degeneracy problem and improve the mixing. # We can compute the update rate of $x_t$ vs $t$ defined as the proportion of times $t$ where $x_t$ gets updated: -update_rate = sum(abs.(diff(particles; dims=2)) .> 0; dims=2) / Nₛ +update_rate = sum(abs.(diff(particles; dims=2)) .> 0; dims=2) / length(chains) #md nothing #hide # and compare it to the theoretical value of $1 - 1/Nₚ$. @@ -130,4 +104,4 @@ plot( xlabel="Iteration", ylabel="Update rate", ) -hline!([1 - 1 / Nₚ]; label="N: $(Nₚ)") +hline!([1 - 1 / length(chains)]; label="N: $(length(chains))") diff --git a/examples/particle-gibbs/script.jl b/examples/particle-gibbs/script.jl index f7a9d2aa..1eb782dc 100644 --- a/examples/particle-gibbs/script.jl +++ b/examples/particle-gibbs/script.jl @@ -37,7 +37,7 @@ end # We consider the following stochastic volatility model: # # ```math -# x_{t+1} = a x_t + v_t \quad v_{t} \sim \mathcal{N}(0, r^2) +# x_{t+1} = a x_t + v_t \quad v_{t} \sim \mathcal{N}(0, q^2) # ``` # ```math # y_{t} = e_t \exp(\frac{1}{2}x_t) \quad e_t \sim \mathcal{N}(0, 1) @@ -45,89 +45,54 @@ end # # We can reformulate the above in terms of transition and observation densities: # ```math -# x_{t+1} \sim f_{\theta}(x_{t+1}|x_t) = \mathcal{N}(a x_t, r^2) +# x_{t+1} \sim f_{\theta}(x_{t+1}|x_t) = \mathcal{N}(a x_t, q^2) # ``` # ```math # y_t \sim g_{\theta}(y_t|x_t) = \mathcal{N}(0, \exp(\frac{1}{2}x_t)^2) # ``` # with the initial distribution $f_0(x) = \mathcal{N}(0, q^2)$. -# Here we assume the static parameters $\theta = (q^2, r^2)$ are known and we are only interested in sampling from the latent state $x_t$. -Parameters = @NamedTuple begin - a::Float64 - q::Float64 - T::Int +# Here we assume the static parameters $\theta = (a^2, q^2)$ are known and we are only interested in sampling from the latent state $x_t$. +struct LinearGaussianDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T} + a::T + q::T end -f(θ::Parameters, state, t) = Normal(θ.a * state, θ.q) -g(θ::Parameters, state, t) = Normal(0, exp(0.5 * state)) -f₀(θ::Parameters) = Normal(0, θ.q) -#md nothing #hide +function SSMProblems.distribution(dyn::LinearGaussianDynamics{T}) where {T<:Real} + return Normal(zero(T), dyn.q) +end -# Let's simulate some data -a = 0.9 # State Variance -q = 0.5 # Observation variance -Tₘ = 200 # Number of observation -Nₚ = 20 # Number of particles -Nₛ = 200 # Number of samples -seed = 1 # Reproduce everything - -θ₀ = Parameters((a, q, Tₘ)) -rng = Random.MersenneTwister(seed) - -x = zeros(Tₘ) -y = zeros(Tₘ) -x[1] = 0 -for t in 1:Tₘ - if t < Tₘ - x[t + 1] = rand(rng, f(θ₀, x[t], t)) - end - y[t] = rand(rng, g(θ₀, x[t], t)) +function SSMProblems.distribution(dyn::LinearGaussianDynamics, ::Int, state) + return Normal(dyn.a * state, dyn.q) end -# Here are the latent and observation series: -plot(x; label="x", xlabel="t") +struct StochasticVolatility{T<:Real} <: SSMProblems.ObservationProcess{T,T} end -# -plot(y; label="y", xlabel="t") - -# Each model takes an `AbstractRNG` as input and generates the logpdf of the current transition: -mutable struct NonLinearTimeSeries <: SSMProblems.AbstractStateSpaceModel - X::Vector{Float64} - observations::Vector{Float64} - θ::Parameters - NonLinearTimeSeries(θ::Parameters) = new(Float64[], θ) - NonLinearTimeSeries(y::Vector{Float64}, θ::Parameters) = new(Float64[], y, θ) +function SSMProblems.distribution(::StochasticVolatility{T}, ::Int, state) where {T<:Real} + return Normal(zero(T), exp((1 / 2) * state)) end -# The dynamics of the model is defined through the `AbstractStateSpaceModel` interface: -function SSMProblems.transition!!(rng::AbstractRNG, model::NonLinearTimeSeries) - return rand(rng, f₀(model.θ)) -end -function SSMProblems.transition!!( - rng::AbstractRNG, model::NonLinearTimeSeries, state::Float64, step::Int -) - return rand(rng, f(model.θ, state, step)) +function LinearGaussianStochasticVolatilityModel(a::T, q::T) where {T<:Real} + dyn = LinearGaussianDynamics(a, q) + obs = StochasticVolatility{T}() + return SSMProblems.StateSpaceModel(dyn, obs) end +#md nothing #hide -function SSMProblems.emission_logdensity( - modeL::NonLinearTimeSeries, state::Float64, step::Int -) - return logpdf(g(model.θ, state, step), model.observations[step]) -end -function SSMProblems.transition_logdensity( - model::NonLinearTimeSeries, prev_state, current_state, step -) - return logpdf(f(model.θ, prev_state, step), current_state) -end +# Let's simulate some data +rng = Random.MersenneTwister(1234) +true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5) +_, x, y = sample(rng, true_model, 200); -# We need to tell AdvancedPS when to stop the execution of the model -# TODO -AdvancedPS.isdone(::NonLinearTimeSeries, step) = step > Tₘ +# Here are the latent and observation series: +plot(x; label="x", xlabel="t") + +# +plot(y; label="y", xlabel="t") # Here we use the particle gibbs kernel without adaptive resampling. -model = NonLinearTimeSeries(y, θ₀) -pg = AdvancedPS.PG(Nₚ, 1.0) -chains = sample(rng, model, pg, Nₛ; progress=false); +model = true_model(y) +pg = AdvancedPS.PG(20, 1.0) +chains = sample(rng, model, pg, 200; progress=false); #md nothing #hide particles = hcat([chain.trajectory.model.X for chain in chains]...) # Concat all sampled states @@ -145,11 +110,11 @@ plot!(mean_trajectory; color=:dodgerblue, label="Mean trajectory", opacity=0.9) # We can also check the mixing as defined in the Gaussian State Space model example. As seen on the # scatter plot above, we are mostly left with a single trajectory before timestep 150. The orange # bar is the optimal mixing rate for the number of particles we use. -plot_update_rate(update_rate(particles, Nₛ)[:, 1], Nₚ) +plot_update_rate(update_rate(particles, 200)[:, 1], 20) # Let's see if ancestor sampling can help with the degeneracy problem. We use the same number of particles, but replace the sampler with PGAS. -pgas = AdvancedPS.PGAS(Nₚ) -chains = sample(rng, model, pgas, Nₛ; progress=false); +pgas = AdvancedPS.PGAS(20) +chains = sample(rng, model, pgas, 200; progress=false); particles = hcat([chain.trajectory.model.X for chain in chains]...); mean_trajectory = mean(particles; dims=2); @@ -161,4 +126,4 @@ plot!(x; color=:darkorange, label="Original Trajectory") plot!(mean_trajectory; color=:dodgerblue, label="Mean trajectory", opacity=0.9) # The update rate is now much higher throughout time. -plot_update_rate(update_rate(particles, Nₛ)[:, 1], Nₚ) +plot_update_rate(update_rate(particles, 200)[:, 1], 20) diff --git a/src/model.jl b/src/model.jl index b7b72db3..d7fb811c 100644 --- a/src/model.jl +++ b/src/model.jl @@ -10,6 +10,32 @@ const Particle = Trace const SSMTrace{R} = Trace{<:SSMProblems.AbstractStateSpaceModel,R} const GenericTrace{R} = Trace{<:AbstractGenericModel,R} +mutable struct TracedSSM{SSM,XT,YT} <: SSMProblems.AbstractStateSpaceModel + model::SSM + X::Vector{XT} + Y::Vector{YT} + + function TracedSSM( + model::SSMProblems.StateSpaceModel{T,LD,OP}, Y::Vector{YT} + ) where {T,LD,OP,YT} + XT = eltype(LD) + @assert eltype(OP) == YT + return new{SSMProblems.StateSpaceModel{T,LD,OP},XT,YT}(model, Vector{XT}(), Y) + end + + function TracedSSM(model::SSMProblems.StateSpaceModel{T,LD,OP}) where {T,LD,OP} + XT = eltype(LD) + YT = eltype(OP) + return new{SSMProblems.StateSpaceModel{T,LD,OP},XT,YT}(model, Vector{YT}()) + end +end + +(model::SSMProblems.StateSpaceModel)(Y::AbstractVector) = TracedSSM(model, Y) +dynamics(ssm::TracedSSM) = ssm.model.dyn +observation(ssm::TracedSSM) = ssm.model.obs + +isdone(ssm::TracedSSM, step::Int) = step > length(ssm.Y) + reset_logprob!(::AdvancedPS.Particle) = nothing reset_model(f) = deepcopy(f) delete_retained!(f) = nothing diff --git a/src/pgas.jl b/src/pgas.jl index 41bcfea7..eb82a2dc 100644 --- a/src/pgas.jl +++ b/src/pgas.jl @@ -23,12 +23,10 @@ current_step(trace::SSMTrace) = trace.rng.count Get the log weight of the transition from previous state of `model` to `x` """ -function transition_logweight(particle::SSMTrace, x) - score = SSMProblems.transition_logdensity( - particle.model, - particle.model.X[current_step(particle) - 2], - x, - current_step(particle) - 1, +function transition_logweight(particle::SSMTrace, x; kwargs...) + iter = current_step(particle) - 1 + score = SSMProblems.logdensity( + dynamics(particle.model), iter, particle.model.X[iter - 1], x, kwargs... ) return score end @@ -61,22 +59,25 @@ function advance!(particle::SSMTrace, isref::Bool=false) if !isref if running_step == 1 - new_state = SSMProblems.transition!!(particle.rng, model) + new_state = SSMProblems.simulate(particle.rng, dynamics(model)) else current_state = model.X[running_step - 1] - new_state = SSMProblems.transition!!( - particle.rng, model, current_state, running_step + new_state = SSMProblems.simulate( + particle.rng, dynamics(model), running_step, current_state ) end else - new_state = model.X[running_step] # We need the current state from the reference particle + # We need the current state from the reference particle + new_state = model.X[running_step] end - score = SSMProblems.emission_logdensity(model, new_state, running_step) + score = SSMProblems.logdensity( + observation(model), running_step, new_state, model.Y[running_step] + ) - # accept transition + # Accept transition and move the time index/rng counter !isref && push!(model.X, new_state) - inc_counter!(particle.rng) # Increase rng counter, we use it as the model `time` index instead of handling two distinct counters + inc_counter!(particle.rng) return score end diff --git a/test/Project.toml b/test/Project.toml index c5edae50..cac881c2 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -21,6 +21,5 @@ HypothesisTests = "0.11" DynamicIterators = "0.4" Libtask = "0.8" Random123 = "1.3" -SSMProblems = "0.1" StableRNGs = "1" julia = "1.3" diff --git a/test/container.jl b/test/container.jl index ba266499..f5434d61 100644 --- a/test/container.jl +++ b/test/container.jl @@ -1,18 +1,20 @@ @testset "container.jl" begin # Since the extension would hide the low level function call API - mutable struct LogPModel{T} <: SSMProblems.AbstractStateSpaceModel + struct LogPDynamics{T} <: LatentDynamics{T,T} end + struct LogPObservation{T} <: ObservationProcess{T,T} logp::T - X::Array{T} end - SSMProblems.transition!!(rng::AbstractRNG, model::LogPModel) = rand(rng, Uniform()) - function SSMProblems.transition!!(rng::AbstractRNG, model::LogPModel, state, step) - return rand(rng, Uniform()) - end - SSMProblems.emission_logdensity(model::LogPModel, state, step) = model.logp + SSMProblems.logdensity(proc::LogPObservation, ::Int, state, observation) = proc.logp + SSMProblems.distribution(proc::LogPDynamics, ::Int, state) = Uniform() + SSMProblems.distribution(::LogPDynamics) = Uniform() - AdvancedPS.isdone(model::LogPModel, step) = false + function LogPModel(logp::T) where {T<:Real} + ssm = StateSpaceModel(LogPDynamics{T}(), LogPObservation(logp)) + # pick some arbitrarily large observables + return ssm(ones(T, 10)) + end @testset "copy particle container" begin pc = AdvancedPS.ParticleContainer(AdvancedPS.Trace[]) @@ -27,13 +29,13 @@ logps = [0.0, -1.0, -2.0] particles = map(logps) do logp trng = AdvancedPS.TracedRNG() - AdvancedPS.Trace(LogPModel(logp, []), trng) + AdvancedPS.Trace(LogPModel(logp), trng) end pc = AdvancedPS.ParticleContainer(particles) # Push! pc2 = deepcopy(pc) - particle = AdvancedPS.Trace(LogPModel(0, []), AdvancedPS.TracedRNG()) + particle = AdvancedPS.Trace(LogPModel(0.0), AdvancedPS.TracedRNG()) push!(pc2, particle) @test length(pc2.vals) == 4 @test pc2.logWs ≈ zeros(4) @@ -67,7 +69,7 @@ # Resample and propagate particles with reference particle particles_ref = map(logps) do logp trng = AdvancedPS.TracedRNG() - AdvancedPS.Trace(LogPModel(logp, []), trng) + AdvancedPS.Trace(LogPModel(logp), trng) end pc_ref = AdvancedPS.ParticleContainer(particles_ref) @@ -168,7 +170,7 @@ particles = map(1:n) do _ trng = AdvancedPS.TracedRNG() - AdvancedPS.Trace(LogPModel(1.0, []), trng) + AdvancedPS.Trace(LogPModel(1.0), trng) end pc = AdvancedPS.ParticleContainer(particles, AdvancedPS.TracedRNG()) @@ -193,7 +195,7 @@ @testset "seed history" begin # Test task copy version of trace trng = AdvancedPS.TracedRNG() - tr = AdvancedPS.Trace(LogPModel(1.0, []), trng) + tr = AdvancedPS.Trace(LogPModel(1.0), trng) AdvancedPS.advance!(tr) AdvancedPS.advance!(tr) diff --git a/test/linear-gaussian.jl b/test/linear-gaussian.jl index 23b0a802..deccc042 100644 --- a/test/linear-gaussian.jl +++ b/test/linear-gaussian.jl @@ -56,46 +56,41 @@ end Xf, ll = kalmanfilter(M, 1 => G0, y_pairs) # Define AdvancedPS model - mutable struct LinearGaussianParams - a::Float64 - b::Float64 - q::Float64 - h::Float64 - r::Float64 - x0::Float64 - p0::Float64 + struct LinearGaussianDynamics{T<:Real} <: LatentDynamics{T,T} + a::T + b::T + q::T end - mutable struct LinearGaussianModel <: SSMProblems.AbstractStateSpaceModel - X::Vector{Float64} - observations::Vector{Float64} - θ::LinearGaussianParams - function LinearGaussianModel(y::Vector{Float64}, θ::LinearGaussianParams) - return new(Vector{Float64}(), y, θ) - end + function SSMProblems.distribution(proc::LinearGaussianDynamics; kwargs...) + return Normal(convert(T, x0), convert(T, p0)) end - function SSMProblems.transition!!(rng::AbstractRNG, model::LinearGaussianModel) - return rand(rng, Normal(model.θ.x0, model.θ.p0)) - end - function SSMProblems.transition!!( - rng::AbstractRNG, model::LinearGaussianModel, state, step + function SSMProblems.distribution( + proc::LinearGaussianDynamics, step::Int, state; kwargs... ) - return rand(rng, Normal(model.θ.a * state + model.θ.b, model.θ.q)) + return Normal(proc.a * state + proc.b, proc.q) end - function SSMProblems.transition_logdensity( - model::LinearGaussianModel, prev_state, current_state, step - ) - return logpdf(Normal(model.θ.a * prev_state + model.θ.b, model.θ.q), current_state) + + struct LinearGaussianObservation{T<:Real} <: ObservationProcess{T,T} + h::T + r::T end - function SSMProblems.emission_logdensity(model::LinearGaussianModel, state, step) - return logpdf(Normal(model.θ.h * state, model.θ.r), model.observations[step]) + + function SSMProblems.distribution( + proc::LinearGaussianObservation, step::Int, state; kwargs... + ) + return Normal(proc.h * state, proc.r) end - AdvancedPS.isdone(::LinearGaussianModel, step) = step > T + function LinearGaussianStateSpaceModel(a, b, q, h, r) + dyn = LinearGaussianDynamics(a, b, q) + obs = LinearGaussianObservation(h, r) + return StateSpaceModel(dyn, obs) + end - params = LinearGaussianParams(a, b, q, H, R, x0, P0) - model = LinearGaussianModel(ys, params) + lgssm = LinearGaussianStateSpaceModel(a, b, q, H, R) + model = lgssm(ys) @testset "PGAS" begin pgas = AdvancedPS.PGAS(N_PARTICLES) diff --git a/test/pgas.jl b/test/pgas.jl index 3a701db7..6a909ac1 100644 --- a/test/pgas.jl +++ b/test/pgas.jl @@ -1,35 +1,40 @@ @testset "pgas.jl" begin - mutable struct Params - a::Float64 - q::Float64 - r::Float64 + mutable struct Params{T<:Real} + a::T + q::T + r::T end - mutable struct BaseModel <: SSMProblems.AbstractStateSpaceModel - X::Vector{Float64} - θ::Params - BaseModel(params::Params) = new(Vector{Float64}(), params) + struct BaseModelDynamics{T<:Real} <: LatentDynamics{T,T} + a::T + q::T end - function SSMProblems.transition!!(rng::AbstractRNG, model::BaseModel) - return rand(rng, Normal(0, model.θ.q)) + function SSMProblems.distribution(dyn::BaseModelDynamics{T}) where {T<:Real} + return Normal(zero(T), dyn.q) end - function SSMProblems.transition!!(rng::AbstractRNG, model::BaseModel, state, step) - return rand(rng, Normal(model.θ.a * state, model.θ.q)) + + function SSMProblems.distribution(dyn::BaseModelDynamics, step::Int, state) + return Normal(dyn.a * state, dyn.q) end - function SSMProblems.emission_logdensity(model::BaseModel, state, step) - return logpdf(Distributions.Normal(state, model.θ.r), 0) + + struct BaseModelObservation{T<:Real} <: ObservationProcess{T,T} + r::T end - function SSMProblems.transition_logdensity( - model::BaseModel, prev_state, current_state, step - ) - return logpdf(Normal(model.θ.a * prev_state, model.θ.q), current_state) + + function SSMProblems.distribution(obs::BaseModelObservation, step::Int, state) + return Normal(state, obs.r) end - AdvancedPS.isdone(::BaseModel, step) = step > 3 + function BaseModel(θ::Params{T}) where {T<:Real} + dyn = BaseModelDynamics(θ.a, θ.q) + obs = BaseModelObservation(θ.r) + ssm = StateSpaceModel(dyn, obs) + return ssm(zeros(T, 3)) + end @testset "fork reference" begin - model = BaseModel(Params(0.9, 0.32, 1)) + model = BaseModel(Params(0.9, 0.32, 1.0)) part = AdvancedPS.Trace(model, AdvancedPS.TracedRNG()) AdvancedPS.advance!(part) @@ -50,7 +55,7 @@ @testset "update reference" begin base_rng = Random.MersenneTwister(31) particles = [ - AdvancedPS.Trace(BaseModel(Params(0.9, 0.31, 1)), AdvancedPS.TracedRNG()) for + AdvancedPS.Trace(BaseModel(Params(0.9, 0.31, 1.0)), AdvancedPS.TracedRNG()) for _ in 1:3 ] sampler = AdvancedPS.PGAS(3) @@ -86,7 +91,7 @@ end @testset "rng stability" begin - model = BaseModel(Params(0.9, 0.32, 1)) + model = BaseModel(Params(0.9, 0.32, 1.0)) seed = 10 rng = Random.MersenneTwister(seed) @@ -117,7 +122,7 @@ # Smoke test mostly @testset "smc sampler" begin - model = BaseModel(Params(0.9, 0.32, 1)) + model = BaseModel(Params(0.9, 0.32, 1.0)) npart = 10 sampler = AdvancedPS.SMC(npart) From 1c253aaea257f0975d5d79303d3fb20dc0649707 Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Thu, 24 Apr 2025 17:02:16 -0400 Subject: [PATCH 2/9] fixed linear Gaussian uinit test --- test/linear-gaussian.jl | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/test/linear-gaussian.jl b/test/linear-gaussian.jl index deccc042..6409ab63 100644 --- a/test/linear-gaussian.jl +++ b/test/linear-gaussian.jl @@ -29,18 +29,18 @@ end N_SAMPLES = 200 # Model dynamics - a = 0.5 - b = 0.2 - q = 0.1 - E = LinearEvolution(a, Gaussian(b, q)) + A = 0.5 + B = 0.2 + Q = 0.1 + E = LinearEvolution(A, Gaussian(B, Q)) H = 1.0 R = 0.1 Obs = LinearObservationModel(H, R) - x0 = 0.0 + X0 = 0.0 P0 = 1.0 - G0 = Gaussian(x0, P0) + G0 = Gaussian(X0, P0) M = LinearStateSpaceModel(E, Obs) O = LinearObservation(E, H, R) @@ -62,8 +62,8 @@ end q::T end - function SSMProblems.distribution(proc::LinearGaussianDynamics; kwargs...) - return Normal(convert(T, x0), convert(T, p0)) + function SSMProblems.distribution(proc::LinearGaussianDynamics{T}; kwargs...) where {T} + return Normal(convert(T, X0), convert(T, P0)) end function SSMProblems.distribution( @@ -89,7 +89,7 @@ end return StateSpaceModel(dyn, obs) end - lgssm = LinearGaussianStateSpaceModel(a, b, q, H, R) + lgssm = LinearGaussianStateSpaceModel(A, B, Q, H, R) model = lgssm(ys) @testset "PGAS" begin From a04e243a0edf132a8df91f20f8ba8faf0a077e3f Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Fri, 25 Apr 2025 11:04:49 -0400 Subject: [PATCH 3/9] fixed Levy SSM --- examples/levy-ssm/script.jl | 223 ++++++++++++++++-------------------- 1 file changed, 97 insertions(+), 126 deletions(-) diff --git a/examples/levy-ssm/script.jl b/examples/levy-ssm/script.jl index 2b35858f..7c2465a3 100644 --- a/examples/levy-ssm/script.jl +++ b/examples/levy-ssm/script.jl @@ -1,6 +1,4 @@ # # Levy-SSM latent state inference -using AdvancedPS: SSMProblems -using AdvancedPS using Random using Plots using Distributions @@ -8,10 +6,11 @@ using AdvancedPS using LinearAlgebra using SSMProblems -struct GammaProcess - C::Float64 - β::Float64 - tol::Float64 +struct GammaProcess{T} + C::T + β::T + tol::T + GammaProcess(C::T, β::T; ϵ::T=1e-10) where {T<:Real} = new{T}(C, β, ϵ) end struct GammaPath{T} @@ -19,36 +18,18 @@ struct GammaPath{T} times::Vector{T} end -struct LangevinDynamics{T} - A::Matrix{T} - L::Vector{T} - θ::T - H::Vector{T} - σe::T -end - -struct NormalMeanVariance{T} - μ::T - σ::T -end - function simulate( - rng::AbstractRNG, - process::GammaProcess, - rate::Float64, - start::Float64, - finish::Float64, - t0::Float64=0.0, -) + rng::AbstractRNG, process::GammaProcess{T}, rate::T, start::T, finish::T, t0::T=zero(T) +) where {T<:Real} let β = process.β, C = process.C, tolerance = process.tol - jumps = Float64[] + jumps = T[] last_jump = Inf t = t0 truncated = last_jump < tolerance while !truncated - t += rand(rng, Exponential(1.0 / rate)) - xi = 1.0 / (β * (exp(t / C) - 1)) - prob = (1.0 + β * xi) * exp(-β * xi) + t += rand(rng, Exponential(one(T) / rate)) + xi = one(T) / (β * (exp(t / C) - one(T))) + prob = (one(T) + β * xi) * exp(-β * xi) if rand(rng) < prob push!(jumps, xi) last_jump = xi @@ -60,17 +41,54 @@ function simulate( end end -function integral(times::Array{Float64}, path::GammaPath) +function integral(times::Array{<:Real}, path::GammaPath) let jumps = path.jumps, jump_times = path.times return [sum(jumps[jump_times .<= t]) for t in times] end end +struct LangevinDynamics{T} + A::Matrix{T} + L::Vector{T} + θ::T + H::Vector{T} + σe::T +end + +struct NormalMeanVariance{T} + μ::T + σ::T +end + +f(dt, θ) = exp(θ * dt) +function Base.exp(dyn::LangevinDynamics{T}, dt::T) where {T<:Real} + let θ = dyn.θ + f_val = f(dt, θ) + return [one(T) (f_val - 1)/θ; zero(T) f_val] + end +end + +function meancov( + t::T, dyn::LangevinDynamics, path::GammaPath, nvm::NormalMeanVariance +) where {T<:Real} + μ = zeros(T, 2) + Σ = zeros(T, (2, 2)) + let times = path.times, jumps = path.jumps, μw = nvm.μ, σw = nvm.σ + for (v, z) in zip(times, jumps) + ft = exp(dyn, (t - v)) * dyn.L + μ += ft .* μw .* z + Σ += ft * transpose(ft) .* σw^2 .* z + end + + # Guarantees positive semi-definiteness + return μ, Σ + T(1e-6) * I + end +end + # Gamma Process C = 1.0 β = 1.0 -ϵ = 1e-10 -process = GammaProcess(C, β, ϵ) +process = GammaProcess(C, β) # Normal Mean-Variance representation μw = 0.0 @@ -78,8 +96,12 @@ process = GammaProcess(C, β, ϵ) nvm = NormalMeanVariance(μw, σw) # Levy SSM with Langevin dynamics -# dx(t) = A x(t) dt + L dW(t) -# y(t) = H x(t) + ϵ(t) +# ```math +# dx_{t} = A x_{t} dt + L dW_{t} +# ``` +# ```math +# y_{t} = H x_{t} + ϵ{t} +# ``` θ = -0.5 A = [ 0.0 1.0 @@ -91,44 +113,17 @@ H = [1.0, 0] dyn = LangevinDynamics(A, L, θ, H, σe) # Simulation parameters -start, finish = 0, 100 N = 200 -ts = range(start, finish; length=N) -seed = 4 -rng = Random.MersenneTwister(seed) -Np = 50 -Ns = 100 - -f(dt, θ) = exp(θ * dt) -function Base.exp(dyn::LangevinDynamics, dt::Real) - let θ = dyn.θ - f_val = f(dt, θ) - return [1.0 (f_val - 1)/θ; 0 f_val] - end -end - -function meancov( - t::T, dyn::LangevinDynamics, path::GammaPath, nvm::NormalMeanVariance -) where {T<:Real} - μ = zeros(T, 2) - Σ = zeros(T, (2, 2)) - let times = path.times, jumps = path.jumps, μw = nvm.μ, σw = nvm.σ - for (v, z) in zip(times, jumps) - ft = exp(dyn, (t - v)) * dyn.L - μ += ft .* μw .* z - Σ += ft * transpose(ft) .* σw^2 .* z - end - return μ, Σ - end -end +ts = range(0, 100; length=N) +rng = Random.MersenneTwister(seed) X = zeros(Float64, (N, 2)) Y = zeros(Float64, N) for (i, t) in enumerate(ts) if i > 1 s = ts[i - 1] dt = t - s - path = simulate(rng, process, dt, s, t, ϵ) + path = simulate(rng, process, dt, s, t) μ, Σ = meancov(t, dyn, path, nvm) X[i, :] .= rand(rng, MultivariateNormal(exp(dyn, dt) * X[i - 1, :] + μ, Σ)) end @@ -138,77 +133,57 @@ for (i, t) in enumerate(ts) end end -# AdvancedPS -Parameters = @NamedTuple begin - dyn::LangevinDynamics - process::GammaProcess - nvm::NormalMeanVariance - times::Vector{Float64} -end +# NOTE: doesn't match 1:1, but I think that's okay +rng = Random.MersenneTwister(seed) +_, x, y = sample(rng, levyssm, N) -struct MixedState{T} - x::Vector{T} - path::GammaPath{T} +# TODO: this can surely be optimized +struct LevyLangevin{T} <: LatentDynamics{T,Vector{T}} + dt::T + dyn::LangevinDynamics{T} + process::GammaProcess{T} + nvm::NormalMeanVariance{T} end -mutable struct LevyLangevin <: SSMProblems.AbstractStateSpaceModel - X::Vector{MixedState{Float64}} - observations::Vector{Float64} - θ::Parameters - LevyLangevin(θ::Parameters) = new(Vector{MixedState{Float64}}(), θ) - function LevyLangevin(y::Vector{Float64}, θ::Parameters) - return new(Vector{MixedState{Float64}}(), y, θ) - end +function SSMProblems.distribution(proc::LevyLangevin{T}) where {T<:Real} + return MultivariateNormal(zeros(T, 2), I) end -function SSMProblems.transition!!(rng::AbstractRNG, model::LevyLangevin) - return MixedState( - rand(rng, MultivariateNormal([0, 0], I)), GammaPath(Float64[], Float64[]) - ) +function SSMProblems.distribution(proc::LevyLangevin{T}, step::Int, state) where {T<:Real} + dt = proc.dt + path = simulate(rng, proc.process, dt, (step - 1) * dt, step * dt) + μ, Σ = meancov(step * dt, proc.dyn, path, proc.nvm) + return MultivariateNormal(exp(proc.dyn, dt) * state + μ, Σ) end -function SSMProblems.transition!!( - rng::AbstractRNG, model::LevyLangevin, state::MixedState, step -) - times = model.θ.times - s = times[step - 1] - t = times[step] - dt = t - s - path = simulate(rng, model.θ.process, dt, s, t) - μ, Σ = meancov(t, model.θ.dyn, path, model.θ.nvm) - Σ += 1e-6 * I - return MixedState(rand(rng, MultivariateNormal(exp(dyn, dt) * state.x + μ, Σ)), path) +struct LinearGaussianObservation{T<:Real} <: ObservationProcess{T,T} + H::Vector{T} + R::T end -function SSMProblems.transition_logdensity( - model::LevyLangevin, prev_state::MixedState, current_state::MixedState, step -) - times = model.θ.times - s = times[step - 1] - t = times[step] - dt = t - s - path = simulate(rng, model.θ.process, dt, s, t) - μ, Σ = meancov(t, model.θ.dyn, path, model.θ.nvm) - Σ += 1e-6 * I - return logpdf(MultivariateNormal(exp(dyn, dt) * prev_state.x + μ, Σ), current_state.x) +function SSMProblems.distribution(proc::LinearGaussianObservation, step::Int, state) + return Normal(transpose(proc.H) * state, proc.R) end -function SSMProblems.emission_logdensity(model::LevyLangevin, state::MixedState, step) - return logpdf(Normal(transpose(H) * state.x, σe), model.observations[step]) +function LevyModel(dt, A, L, θ, H, σe, C, β, ϵ, μw, σw) + dyn = LevyLangevin( + dt, + LangevinDynamics(A, L, θ, H, σe), + GammaProcess(C, β; ϵ), + NormalMeanVariance(μw, σw), + ) + + obs = LinearGaussianObservation(H, σe) + return StateSpaceModel(dyn, obs) end -AdvancedPS.isdone(model::LevyLangevin, step) = step > length(model.θ.times) +levyssm = LevyModel(0.5025125628140756, A, L, θ, H, σe, C, β, ϵ, μw, σw); -θ₀ = Parameters((dyn, process, nvm, ts)) -model = LevyLangevin(Y, θ₀) -pg = AdvancedPS.PGAS(Np) -chains = sample(rng, model, pg, Ns; progress=false); +pg = AdvancedPS.PGAS(50); +chains = sample(rng, levyssm(Y), pg, 100); # Concat all sampled states -particles = hcat([chain.trajectory.model.X for chain in chains]...) -marginal_states = map(s -> s.x, particles); -jump_times = map(s -> s.path.times, particles); -jump_intensities = map(s -> s.path.jumps, particles); +marginal_states = hcat([chain.trajectory.model.X for chain in chains]...) # Plot marginal state and jump intensities for one trajectory p1 = plot( @@ -224,12 +199,8 @@ plot!( label="Marginal State (x2)", ) -p2 = scatter( - vcat([t for t in jump_times[:, end]]...), - vcat([j for j in jump_intensities[:, end]]...); - color=:darkorange, - label="Jumps", -) +# TODO: collect jumps from the model +p2 = scatter([], []; color=:darkorange, label="Jumps") plot( p1, p2; plot_title="Marginal State and Jump Intensities", layout=(2, 1), size=(600, 600) From 9e2f73fbcf5a37d52d8836d3cac588505b6367f2 Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Fri, 25 Apr 2025 11:58:50 -0400 Subject: [PATCH 4/9] temporary fix, that yields inefficient sampler --- examples/gaussian-process/script.jl | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/examples/gaussian-process/script.jl b/examples/gaussian-process/script.jl index 34db890c..9b0b3627 100644 --- a/examples/gaussian-process/script.jl +++ b/examples/gaussian-process/script.jl @@ -1,3 +1,6 @@ +using Pkg +Pkg.activate("/sftwr/user-pkg/m1cak00/julia/dev/AdvancedPS/examples/gaussian-process") + # # Gaussian Process State-Space Model (GP-SSM) using LinearAlgebra using Random @@ -25,14 +28,14 @@ end function SSMProblems.simulate( rng::AbstractRNG, dyn::GaussianProcessDynamics, step::Int, state ) - dyn.proc = posterior(dyn.proc(step:step), [state]) + dyn.proc = posterior(dyn.proc(step:step, 0.1), [state]) μ, σ = mean_and_cov(dyn.proc, [step]) return rand(rng, Normal(μ[1], sqrt(σ[1]))) end function SSMProblems.logdensity(dyn::GaussianProcessDynamics, step::Int, state, prev_state) μ, σ = mean_and_cov(dyn.proc, [step]) - return logpdf(Normal(μ, sqrt(σ)), state) + return logpdf(Normal(μ[1], sqrt(σ[1])), state) end # Linear Gaussian dynamics used for simulation @@ -79,7 +82,7 @@ _, x, y = sample(rng, true_model, 100); gpssm = GaussianProcessStateSpaceModel(0.5, SqExponentialKernel()) model = gpssm(y) pg = AdvancedPS.PGAS(20) -chains = sample(rng, model, pg, 250) +chains = sample(rng, model, pg, 50) #md nothing #hide particles = hcat([chain.trajectory.model.X for chain in chains]...) From 5632c921740eca66912dcebf1d990f1458b31fcc Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Fri, 25 Apr 2025 13:54:55 -0400 Subject: [PATCH 5/9] improved GP fix --- examples/gaussian-process/script.jl | 81 +++++++++++++---------------- src/model.jl | 4 +- src/pgas.jl | 8 +-- 3 files changed, 43 insertions(+), 50 deletions(-) diff --git a/examples/gaussian-process/script.jl b/examples/gaussian-process/script.jl index 9b0b3627..a971cf28 100644 --- a/examples/gaussian-process/script.jl +++ b/examples/gaussian-process/script.jl @@ -11,81 +11,74 @@ using Distributions using Libtask using SSMProblems -# Gaussian process encoded transition dynamics -mutable struct GaussianProcessDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T} - proc::AbstractGPs.AbstractGP - q::T - function GaussianProcessDynamics(q::T, kernel::KT) where {T<:Real,KT<:Kernel} - return new{T}(GP(ZeroMean{T}(), kernel), q) +struct GaussianProcessDynamics{T<:Real,KT<:Kernel} <: LatentDynamics{T,T} + proc::GP{ZeroMean{T},KT} + function GaussianProcessDynamics(::Type{T}, kernel::KT) where {T<:Real,KT<:Kernel} + return new{T,KT}(GP(ZeroMean{T}(), kernel)) end end -function SSMProblems.distribution(dyn::GaussianProcessDynamics{T}) where {T<:Real} - return Normal(zero(T), dyn.q) -end - -# TODO: broken... -function SSMProblems.simulate( - rng::AbstractRNG, dyn::GaussianProcessDynamics, step::Int, state -) - dyn.proc = posterior(dyn.proc(step:step, 0.1), [state]) - μ, σ = mean_and_cov(dyn.proc, [step]) - return rand(rng, Normal(μ[1], sqrt(σ[1]))) -end - -function SSMProblems.logdensity(dyn::GaussianProcessDynamics, step::Int, state, prev_state) - μ, σ = mean_and_cov(dyn.proc, [step]) - return logpdf(Normal(μ[1], sqrt(σ[1])), state) -end - -# Linear Gaussian dynamics used for simulation -struct LinearGaussianDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T} +struct LinearGaussianDynamics{T<:Real} <: LatentDynamics{T,T} a::T + b::T q::T end -function SSMProblems.distribution(dyn::LinearGaussianDynamics{T}) where {T<:Real} - return Normal(zero(T), dyn.q) +function SSMProblems.distribution(proc::LinearGaussianDynamics{T}) where {T<:Real} + return Normal(zero(T), proc.q) end -function SSMProblems.distribution(dyn::LinearGaussianDynamics, ::Int, state) - return Normal(dyn.a * state, dyn.q) +function SSMProblems.distribution(proc::LinearGaussianDynamics, ::Int, state) + return Normal(proc.a * state + proc.b, proc.q) end -# Observation process used in both variants of the model -struct StochasticVolatility{T<:Real} <: SSMProblems.ObservationProcess{T,T} end +struct StochasticVolatility{T<:Real} <: ObservationProcess{T,T} end function SSMProblems.distribution(::StochasticVolatility{T}, ::Int, state) where {T<:Real} return Normal(zero(T), exp((1 / 2) * state)) end -# Baseline model (for simulation) function LinearGaussianStochasticVolatilityModel(a::T, q::T) where {T<:Real} - dyn = LinearGaussianDynamics(a, q) + dyn = LinearGaussianDynamics(a, zero(T), q) obs = StochasticVolatility{T}() return SSMProblems.StateSpaceModel(dyn, obs) end -# Gaussian process model (for sampling) -function GaussianProcessStateSpaceModel(q::T, kernel::KT) where {T<:Real,KT<:Kernel} - dyn = GaussianProcessDynamics(q, kernel) +function GaussianProcessStateSpaceModel(::Type{T}, kernel::KT) where {T<:Real,KT<:Kernel} + dyn = GaussianProcessDynamics(T, kernel) obs = StochasticVolatility{T}() return SSMProblems.StateSpaceModel(dyn, obs) end +const GPSSM{T,KT<:Kernel} = SSMProblems.StateSpaceModel{ + T, + GaussianProcessDynamics{T,KT}, + StochasticVolatility{T} +}; + +# for non-markovian models, we can redefine dynamics to reference the trajectory +function AdvancedPS.dynamics( + ssm::AdvancedPS.TracedSSM{<:GPSSM{T},T,T}, step::Int +) where {T<:Real} + prior = ssm.model.dyn.proc(1:(step - 1)) + post = posterior(prior, ssm.X[1:(step - 1)]) + μ, σ = mean_and_cov(post, [step]) + return LinearGaussianDynamics(zero(T), μ[1], sqrt(σ[1])) +end + # Everything is now ready to simulate some data. -rng = Random.MersenneTwister(1234) -true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5) +rng = MersenneTwister(1234); +true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5); _, x, y = sample(rng, true_model, 100); # Create the model and run the sampler -gpssm = GaussianProcessStateSpaceModel(0.5, SqExponentialKernel()) -model = gpssm(y) -pg = AdvancedPS.PGAS(20) -chains = sample(rng, model, pg, 50) +gpssm = GaussianProcessStateSpaceModel(Float64, SqExponentialKernel()); +model = gpssm(y); +pg = AdvancedPS.PGAS(20); +chains = sample(rng, model, pg, 250; progress=false); #md nothing #hide -particles = hcat([chain.trajectory.model.X for chain in chains]...) +particles = hcat([chain.trajectory.model.X for chain in chains]...); mean_trajectory = mean(particles; dims=2); #md nothing #hide diff --git a/src/model.jl b/src/model.jl index d7fb811c..bb212ed9 100644 --- a/src/model.jl +++ b/src/model.jl @@ -31,8 +31,8 @@ mutable struct TracedSSM{SSM,XT,YT} <: SSMProblems.AbstractStateSpaceModel end (model::SSMProblems.StateSpaceModel)(Y::AbstractVector) = TracedSSM(model, Y) -dynamics(ssm::TracedSSM) = ssm.model.dyn -observation(ssm::TracedSSM) = ssm.model.obs +dynamics(ssm::TracedSSM, step::Int) = ssm.model.dyn +observation(ssm::TracedSSM, step::Int) = ssm.model.obs isdone(ssm::TracedSSM, step::Int) = step > length(ssm.Y) diff --git a/src/pgas.jl b/src/pgas.jl index eb82a2dc..c14bc736 100644 --- a/src/pgas.jl +++ b/src/pgas.jl @@ -26,7 +26,7 @@ Get the log weight of the transition from previous state of `model` to `x` function transition_logweight(particle::SSMTrace, x; kwargs...) iter = current_step(particle) - 1 score = SSMProblems.logdensity( - dynamics(particle.model), iter, particle.model.X[iter - 1], x, kwargs... + dynamics(particle.model, iter), iter, particle.model.X[iter - 1], x, kwargs... ) return score end @@ -59,11 +59,11 @@ function advance!(particle::SSMTrace, isref::Bool=false) if !isref if running_step == 1 - new_state = SSMProblems.simulate(particle.rng, dynamics(model)) + new_state = SSMProblems.simulate(particle.rng, dynamics(model, running_step)) else current_state = model.X[running_step - 1] new_state = SSMProblems.simulate( - particle.rng, dynamics(model), running_step, current_state + particle.rng, dynamics(model, running_step), running_step, current_state ) end else @@ -72,7 +72,7 @@ function advance!(particle::SSMTrace, isref::Bool=false) end score = SSMProblems.logdensity( - observation(model), running_step, new_state, model.Y[running_step] + observation(model, running_step), running_step, new_state, model.Y[running_step] ) # Accept transition and move the time index/rng counter From 87cc0035cde9f83c5e7f8bffa2e2304b9816af37 Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Fri, 25 Apr 2025 14:00:58 -0400 Subject: [PATCH 6/9] remove environment path --- examples/gaussian-process/script.jl | 3 --- 1 file changed, 3 deletions(-) diff --git a/examples/gaussian-process/script.jl b/examples/gaussian-process/script.jl index a971cf28..40957858 100644 --- a/examples/gaussian-process/script.jl +++ b/examples/gaussian-process/script.jl @@ -1,6 +1,3 @@ -using Pkg -Pkg.activate("/sftwr/user-pkg/m1cak00/julia/dev/AdvancedPS/examples/gaussian-process") - # # Gaussian Process State-Space Model (GP-SSM) using LinearAlgebra using Random From 519d101c414c87cf1f501e412b5f620adaa8365b Mon Sep 17 00:00:00 2001 From: Hong Ge <3279477+yebai@users.noreply.github.com> Date: Mon, 28 Apr 2025 22:45:22 +0100 Subject: [PATCH 7/9] Update Project.toml --- Project.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/Project.toml b/Project.toml index 2813c961..ab782442 100644 --- a/Project.toml +++ b/Project.toml @@ -26,6 +26,7 @@ Random = "<0.0.1, 1" Random123 = "1.3" Requires = "1.0" StatsFuns = "0.9, 1" +SSMProblems = "0.5" julia = "1.7" [extras] From c258ee32e3a1afaf763426aa5c83254e5dadf61a Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Thu, 1 May 2025 10:31:32 -0400 Subject: [PATCH 8/9] remove unused constructor --- src/model.jl | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/model.jl b/src/model.jl index bb212ed9..c836c74f 100644 --- a/src/model.jl +++ b/src/model.jl @@ -22,12 +22,6 @@ mutable struct TracedSSM{SSM,XT,YT} <: SSMProblems.AbstractStateSpaceModel @assert eltype(OP) == YT return new{SSMProblems.StateSpaceModel{T,LD,OP},XT,YT}(model, Vector{XT}(), Y) end - - function TracedSSM(model::SSMProblems.StateSpaceModel{T,LD,OP}) where {T,LD,OP} - XT = eltype(LD) - YT = eltype(OP) - return new{SSMProblems.StateSpaceModel{T,LD,OP},XT,YT}(model, Vector{YT}()) - end end (model::SSMProblems.StateSpaceModel)(Y::AbstractVector) = TracedSSM(model, Y) From 563a27a1954e90a2fd4e6b144adf48c35138ee44 Mon Sep 17 00:00:00 2001 From: Charles Knipp Date: Thu, 1 May 2025 10:32:25 -0400 Subject: [PATCH 9/9] update Levy SSM --- examples/levy-ssm/script.jl | 79 +++++++++++-------------------------- 1 file changed, 23 insertions(+), 56 deletions(-) diff --git a/examples/levy-ssm/script.jl b/examples/levy-ssm/script.jl index 7c2465a3..e42a2eec 100644 --- a/examples/levy-ssm/script.jl +++ b/examples/levy-ssm/script.jl @@ -85,59 +85,6 @@ function meancov( end end -# Gamma Process -C = 1.0 -β = 1.0 -process = GammaProcess(C, β) - -# Normal Mean-Variance representation -μw = 0.0 -σw = 1.0 -nvm = NormalMeanVariance(μw, σw) - -# Levy SSM with Langevin dynamics -# ```math -# dx_{t} = A x_{t} dt + L dW_{t} -# ``` -# ```math -# y_{t} = H x_{t} + ϵ{t} -# ``` -θ = -0.5 -A = [ - 0.0 1.0 - 0.0 θ -] -L = [0.0; 1.0] -σe = 1.0 -H = [1.0, 0] -dyn = LangevinDynamics(A, L, θ, H, σe) - -# Simulation parameters -N = 200 -ts = range(0, 100; length=N) - -rng = Random.MersenneTwister(seed) -X = zeros(Float64, (N, 2)) -Y = zeros(Float64, N) -for (i, t) in enumerate(ts) - if i > 1 - s = ts[i - 1] - dt = t - s - path = simulate(rng, process, dt, s, t) - μ, Σ = meancov(t, dyn, path, nvm) - X[i, :] .= rand(rng, MultivariateNormal(exp(dyn, dt) * X[i - 1, :] + μ, Σ)) - end - - let H = dyn.H, σe = dyn.σe - Y[i] = transpose(H) * X[i, :] + rand(rng, Normal(0, σe)) - end -end - -# NOTE: doesn't match 1:1, but I think that's okay -rng = Random.MersenneTwister(seed) -_, x, y = sample(rng, levyssm, N) - -# TODO: this can surely be optimized struct LevyLangevin{T} <: LatentDynamics{T,Vector{T}} dt::T dyn::LangevinDynamics{T} @@ -165,7 +112,11 @@ function SSMProblems.distribution(proc::LinearGaussianObservation, step::Int, st return Normal(transpose(proc.H) * state, proc.R) end -function LevyModel(dt, A, L, θ, H, σe, C, β, ϵ, μw, σw) +function LevyModel(dt, θ, σe, C, β, μw, σw; ϵ=1e-10) + A = [0.0 1.0; 0.0 θ] + L = [0.0; 1.0] + H = [1.0, 0] + dyn = LevyLangevin( dt, LangevinDynamics(A, L, θ, H, σe), @@ -177,8 +128,24 @@ function LevyModel(dt, A, L, θ, H, σe, C, β, ϵ, μw, σw) return StateSpaceModel(dyn, obs) end -levyssm = LevyModel(0.5025125628140756, A, L, θ, H, σe, C, β, ϵ, μw, σw); +# Levy SSM with Langevin dynamics +# ```math +# dx_{t} = A x_{t} dt + L dW_{t} +# ``` +# ```math +# y_{t} = H x_{t} + ϵ{t} +# ``` + +# Simulation parameters +N = 200 +ts = range(0, 100; length=N) +levyssm = LevyModel(step(ts), θ, 1.0, 1.0, 1.0, 0.0, 1.0); + +# Simulate data +rng = Random.MersenneTwister(1234); +_, X, Y = sample(rng, levyssm, N); +# Run sampler pg = AdvancedPS.PGAS(50); chains = sample(rng, levyssm(Y), pg, 100); @@ -221,7 +188,7 @@ for d in 1:2 fillalpha=0.2, title="Marginal State Trajectories (X$d)", ) - plot!(p, ts, X[:, d]; color=:dodgerblue, label="True Trajectory") + plot!(p, ts, getindex.(X, d); color=:dodgerblue, label="True Trajectory") push!(ps, p) end plot(ps...; layout=(2, 1), size=(600, 600))