Kalman filter

[r2cp]

Hi,

I would like to know if there is a plan to implement Kalman filtering/smoothing routines? I see there is some work in the dev_dfm branch, but I am not sure whether it's complete or how to test it. I could help with tests and documentation.

Thank you for your great work!

[bbejanov]

Hi @r2cp, thank you for your interest.

Indeed, branch dev_dfm has an implementation of Kalman filter and smoother. It's not yet published because its interface is not user friendly. Also, it is a part of a bigger effort to add DFM models, which is still under development. We may be able to split the Kalman Filter from the DFM and publish it on its own sooner, if that would be of interest.

The implementation for linear stationary models has been tested - it works well and is okay for speed. The general version has not been tested, so there may be some bugs left in it and it's probably slower than it could be.

If you want to give it a try, please let me know. I can prepare a simple example and some instructions to help you get started.

Cheers!

[r2cp]

Dear @bbejanov,

Yes, please! If it's possible for you to setup a simple example, I would like to give it a try.

And in case of splitting the development of the Kalman routines and the DFM models, please let me know if I could collaborate with the documentation or tests.

Thank you again!

[bbejanov]

Thank you, @r2cp, for your patience!

I have extracted the Kalman Filter and Smoother code into its own branch called kalman_filter. You can try it by adding StateSpaceEcon#kalman_filter to your project environment.

The code is contained into its own submodule called StateSpaceEcon.Kalman, and for now it is unrelated to the rest of the package. Like I said in my earlier comment, the linear stationary case should work fine. This is in functions Kalman.dk_filter! and Kalman.dk_smoother!. The general case (Kalman.filter and Kalman.smoother) are probably not going to work, but feel free to give them a try and let me know how it goes.

Following is a script that shows an example how to setup a model and call Kalman.dk_filter!() and Kalman.dk_smoother!().
To run the script you will need have these packages installed:

] add LinearAlgebra Plots Random TimeSeriesEcon 
] add StateSpaceEcon#kalman_filter

Then the following script should be able to run.

using LinearAlgebra
using Random

using Plots

using TimeSeriesEcon
using StateSpaceEcon
Kalman = StateSpaceEcon.Kalman

##

# simple linear stationary state space model
#
#    yₜ = mu + H * xₜ + vₜ      vₜ ~ N(0, Q)
#    xₜ = F * xₜ₋₁ + wₜ         wₜ ~ N(0, R)
#

struct MyModel
    mu
    H
    F
    Q
    R
    function MyModel(mu, H, F, Q, R)
        ny, nx = size(H)
        if size(mu) != (ny,) || size(F) != (nx, nx) ||
           size(Q) != (ny, ny) || size(R) != (nx, nx)
            throw(ArgumentError("Incompatible sizes"))
        end
        new(mu, H, F, Q, R)
    end
end

# Simple DFM example to try:

#       y₁ₜ =  6 + x₁ₜ + 0.8fₜ + v₁ₜ
#       y₂ₜ = -2 + x₂ₜ + -0.4fₜ + v₂ₜ
#       x₁ₜ =  0.3x₁ₜ₋₁ + w₁ₜ           # idiosyncratic component for y₁
#       x₂ₜ =  0.4x₂ₜ₋₁ + w₂ₜ           # idiosyncratic component for y₂
#       fₜ  =  0.7fₜ₋₁ + w₃ₜ            # common factor

# Create MyModel instance for this model
model = MyModel(
    [6.0, -2.0], # mu
    [1.0 0.0 0.7; 0.0 1.0 -0.4],   # H
    Matrix(Diagonal([0.3, 0.4, 0.7])),   # F
    Diagonal([0.9, 1.2]),                # Q
    Diagonal([0.2, 0.1, 1.6])
)

##   Simulate data

Random.seed!(0x007)

# initial states and state covariance
sim_x0 = zeros(3)
sim_Px0 = model.R

# simulation range
rng = 2024M1 .+ (0:59)   # 5 years of monthly data

# allocate containers for observed and states
sim_y = MVTSeries(rng, (:y1, :y2))
sim_x = MVTSeries(rng, (:x1, :x2, :f))

let
    # simulate states (xₜ)
    local curr_state = copy(sim_x0)
    for t in rng
        curr_state = model.F * curr_state + sqrt.(model.R) * randn(3)
        sim_x[t, :] = curr_state
    end
    # simulate observations (yₜ)
    sim_y .= sim_x * transpose(model.H) + randn(length(rng), 2) * sqrt.(model.Q)
    sim_y .+= transpose(model.mu)
end


# visualize simulated data
common_plot_args = (;
    vars=(colnames(sim_y)..., colnames(sim_x)...),
    lw=2.0,
    layout=(3, 2),
    size=(800, 600),
)

fig = plot([sim_y sim_x];
    label=["Simulated"],
    common_plot_args...)
# display(fig)
# savefig(fig, "simulated.pdf")


##  Prep for using Kalman 

# Tell Kalman the number of observed variables in MyModel ...
Kalman.kf_length_y(model::MyModel) = size(model.H, 1)
# ... and the number of states in MyModel
Kalman.kf_length_x(model::MyModel) = size(model.H, 2)

# allocate Kalman data for our range and model
kfd = Kalman.KFDataSmoother(rng, model)
# allocate filter instance for our `kfd`
kf = Kalman.KFilter(kfd)

##  Filter

# run filter (updates `kfd` in place)
Kalman.dk_filter!(kf, sim_y,
    model.mu, model.H, model.F, model.Q, model.R, I,
    sim_x0, sim_Px0, false)

# extract filtered (a.k.a predicted) data from `kfd`
pred_y = copyto!(similar(sim_y), transpose(kfd.y_pred))
pred_x = copyto!(similar(sim_x), transpose(kfd.x_pred))

# visualize
fig = plot([sim_y sim_x], [pred_y pred_x];
    label=["Simulated" "Filtered"],
    common_plot_args...)
# display(fig)
# savefig(fig, "filtered.pdf")

##   Smoother

# run smoother (updates `kfd`)
Kalman.dk_smoother!(kf, sim_y,
    model.mu, model.H, model.F, model.Q, model.R, I, false)

# extract smoothed data from kfd
smooth_y = copyto!(similar(sim_y), transpose(kfd.y_smooth))
smooth_x = copyto!(similar(sim_x), transpose(kfd.x_smooth))

# viz
fig = plot([sim_y sim_x], [smooth_y smooth_x];
    label=["Simulated" "Smoothed"],
    common_plot_args...)
# display(fig)
# savefig(fig, "smoothed.pdf")

[r2cp]

Dear @bbejanov, thank you so much for setting up this example!

I would like to know if there is any way to apply the filter/smoother to a Model object from ModelBaseEcon.jl? I am working with a very simple Local-Level model to showcase the use of this package for simulations in a similar fashion to what I would do in MATLAB/IRIS.

My Local-Level model is as follows:

# LocalLevel.jl
module LocalLevel
using ModelBaseEcon 
const model = Model()

model.flags.ssZeroSlope = true 
model.flags.linear = true
setoption!(model) do o
    o.tol = 1e-12
    o.maxiter = 100
    o.substitutions = false
    o.factorization = :default
    o.verbose = true
end
@parameters model begin 
    c = 0.01
    ρ = 0.9
end
@variables model begin
    y
    alpha
end
@shocks model begin 
    s_y
    s_alpha
end
@autoexogenize model begin 
    y = s_y
    alpha = s_alpha
end
@equations model begin
    y[t] = alpha[t] + s_y[t]
    alpha[t] = c + ρ*alpha[t-1] + s_alpha[t] 
end
@initialize(model)
end

What I would like to do is to give data for y and get smoothed alpha. Something like this:

# LocalLevel_test.jl
include("LocalLevel.jl")
using .LocalLevel
m = LocalLevel.model

using TimeSeriesEcon, StateSpaceEcon, ModelBaseEcon
Kalman = StateSpaceEcon.Kalman

# Check model status
m.sstate
clear_sstate!(m)
sssolve!(m; method=:auto)
check_sstate(m)
printsstate(m)

# Create data
simrng = 2001Q1:2010Q4
T = length(simrng)
y_true = zeros(T)
y_true[1] = rand()
for t in 2:T 
    y_true[t] = 0.9*y_true[t-1] + 0.05*randn()
end

p  = Plan(m, simrng)
zz = zerodata(m, p)
zz[begin+1:end, :y] .= y_true
zz[begin, :alpha] = rand() # Initial state for alpha

## Apply Kalman filter to recover alpha from `p`, `m`, `zz` 
# ... 

Is there a way to extract the State-Space representation matrices from the model object m to apply what you showed me on your example?

Thank you again!