Add filer and smoothing for regime switching model

src/QuantEcon.jl

@@ -147,7 +147,12 @@ export
 
 # filter
     hp_filter,
-    hamilton_filter
+    hamilton_filter,
+    filter,
+    filter!,
+    smooth,
+    smooth!,
+    RegimeSwitchingModel
 
 
 include("sampler.jl")

src/filter.jl

@@ -85,3 +85,206 @@ function hamilton_filter(y::AbstractVector, h::Integer)
     y_trend = y - y_cycle
     return y_cycle, y_trend
 end
+
+"""
+- `g::Function`: Parameterized `Function` that recieves three arguments 
+                 (data at each period, state at each period, parameter) and return the density.
+                 For example, if the data `x <: Real` is generated by `N(mu_s, sigma_s)` where mean `mu` and std `sigma` is
+                 state dependent, `g` should be defined as 
+                 ```
+                 g(x, s, parameter) = exp(-((x-parameter.mu[s])^2)/(2*parameter.sigma[s]^2))/sqrt(2*pi*parameter.sigma[s]^2)
+                 ```
+- `y::AbstractArray`: Collection of data whose row corresponds to periods
+- `parameter`: Collection of parameters in the model. Each parameter must be a vector that stores the values for each state.
+               For example, if you want to set `mu` at state 1 is 0 and `mu` at state 2 is 0.5 and 
+               `sigma` at state 1 is 3 and `sigma` at state 2 is 4, `parameter` should be
+               ```
+               parameter = (mu = [0, 0.5], sigma = [3, 4])
+               ```
+               Alternatively, you can define your own type and use it:
+               ```
+               struct ParaType
+                   mu::Vector
+                   sigma::Vector
+               end
+               parameter = ParaType([0, 0.5], [3, 4])
+               ```
+- `P::AbstractMatrix`: Transition matrix of state. It must be square.
+- `M::Integer`: Number of states. It must be same as the number of rows and columns of `P`. If `M` is not passed, 
+                number of rows of `P` is set as `M`.
+"""
+mutable struct RegimeSwitchingModel{TD <: AbstractArray, TP <: AbstractMatrix, TPS <: AbstractMatrix}
+    g::Function
+    y::TD
+    parameter
+    P::TP
+    M::Int
+    prob_smoothed::TPS
+    logL::Float64
+    function RegimeSwitchingModel(g, y, parameter, P, M, prob_smoothed, logL)
+        if size(P, 1) != size(P, 2)
+            error("P must be square")
+        end
+        if size(P, 1) != M
+            error("the number of rows and columns of P must be equal to M")
+        end
+        return new{typeof(y), typeof(P), typeof(prob_smoothed)}(g, y, parameter, P, M, prob_smoothed, logL)
+    end
+end
+RegimeSwitchingModel(g::Function, y::AbstractArray, parameter, P::AbstractMatrix) =
+    RegimeSwitchingModel(g, y, parameter, P, size(P, 1), fill(NaN, size(y, 1), size(P, 1)), NaN)
+
+"""
+Apply the filter developed by Hamilton (1989, Econometrica) to a regime switching model.
+
+##### Arguments
+- `rsm::RegimeSwitchingModel`: `RegimeSwitchingModel`
+- `prob_update_pre::AbstractArray`: Probability distribution of state at period 0 conditional on the information up to
+                                   period 0.
+
+##### Returns
+- `p`: likelihood at each period
+- `p_s_update`: Probability of period `t` data conditional on the information up to period `t`.
+- `p_s_forecast`: Probability of period `t` data conditional on the information up to period `t-1`.
+"""
+function filter!(rsm::RegimeSwitchingModel;
+                 prob_update_pre::AbstractArray = stationary_distributions(MarkovChain(rsm.P))[1])
+    T = size(rsm.y, 1)
+    p_s_forecast = Matrix{Float64}(undef, T, rsm.M)
+    p_s_update = Matrix{Float64}(undef, T, rsm.M)
+    p, _, _ = filter!(rsm, p_s_update, p_s_forecast, prob_update_pre = prob_update_pre)
+    return p, p_s_update, p_s_forecast
+end
+
+"""
+Apply the filter developed by Hamilton (1989, Econometrica) to a regime switching model.
+
+##### Arguments
+- `rsm::RegimeSwitchingModel`: `RegimeSwitchingModel` specifying the model
+- `p_s_update::AbstractMatrix`: Probability of period `t` data conditional on the information up to period `t`.
+- `p_s_forecast::AbstractMatrix`: Probability of period `t` data conditional on the information up to period `t-1`.
+- `prob_update_pre::AbstractArray`: Probability distribution of state at period 0 conditional on the information up to
+                                   period 0.
+
+##### Returns
+- `p`: Likelihood at each period.
+- `p_s_update`: Probability of period `t` data conditional on the information up to period `t`.
+- `p_s_forecast`: Probability of period `t` data conditional on the information up to period `t-1`.
+"""
+function filter!(rsm::RegimeSwitchingModel, 
+                 p_s_update::AbstractMatrix, p_s_forecast::AbstractMatrix;
+                 prob_update_pre::AbstractArray = stationary_distributions(MarkovChain(rsm.P))[1])
+    g, y = rsm.g, rsm.y
+    T = size(y, 1)
+    p = Vector{Float64}(undef, T)
+    logL = 0
+    for t in 1:T
+        forecast!(view(p_s_forecast, t, :), prob_update_pre, rsm.P)
+        p[t] = update!(view(p_s_update, t, :), rsm, p_s_forecast[t, :], t)
+        logL += log(p[t])
+        prob_update_pre .= p_s_update[t, :]
+    end
+    rsm.logL = logL
+    return p, p_s_update, p_s_forecast
+end
+
+"""
+Given the state distribution of period `t-1` conditional on the information up to period `t-1`, 
+forecast the distribution of period `t`.
+
+##### Arguments
+- `p_s_forecast::AbstractVector`: Probability distribution of state at period `t` conditional on the 
+                                  information up to period `t-1`.
+- `p_s_update_pre::AbstractArray`: Probability distribution of state at period `t-1` conditional on the 
+                                   information up to period `t-1`.
+- `P::AbstractMatrix`: Transition matrix of state.
+
+#### Return
+- `nothing`
+"""
+function forecast!(p_s_forecast::AbstractVector, p_s_update_pre::AbstractArray, P::AbstractMatrix)
+    p_s_forecast .= P' * p_s_update_pre
+    return nothing
+end
+
+"""
+Given the state distribution of period `t` conditional on the information up to period `t-1`, 
+update the distribution using the infomration at period `t`.
+
+##### Arguments
+- `p_s_update::AbstractVector`: Probability distribution of state at period `t` conditional on the 
+                         information up to period `t`.
+- `rsm::RegimeSwitchingModel`: `RegimeSwitchingModel` that specifies the model.
+- `p_s_forecast::Array`: Probability distribution of state at period `t` conditional on the 
+                information up to period `t-1`.
+- `t::Integer`: Period.
+
+#### Returns
+- `p`: Likelihood of data at period `t`.
+"""
+function update!(p_s_update::AbstractVector, rsm::RegimeSwitchingModel, p_s_forecast::AbstractArray, t::Integer)
+    y_at_period_t = get_y_at_period_t(rsm.y, t)
+    eta = [max(0, rsm.g(y_at_period_t, s, rsm.parameter)) for s in 1:rsm.M]
+    tmp = eta .* p_s_forecast
+    p = sum(tmp)
+    p_s_update .= tmp./p
+    return p
+end
+get_y_at_period_t(y::AbstractMatrix, t::Union{Integer, AbstractVector}) = y[t, :]
+get_y_at_period_t(y::AbstractVector, t::Union{Integer, AbstractVector}) = y[t]
+
+"""
+Apply the backward smoother developed by Kim to a regime switching model.
+
+##### Arguments
+- `rsm::RegimeSwitchingModel`: `RegimeSwitchingModel` specifying the model.
+- `prob_update_pre::AbstractArray`: Probability distribution of state at period 0 conditional on the information up to
+                                   period 0.
+- `verbose::Bool`: If `true`, the funciton returns joint distribution. If `false`, `nothing` is returned.
+
+##### Returns (only when `verbose` is `true`)
+- `p_s_joint_smoothed`: Joint distribution of state at period `t` and `t+1` data conditional on the all information.
+"""
+function smooth!(rsm::RegimeSwitchingModel;
+                 prob_update_pre::AbstractArray = stationary_distributions(MarkovChain(rsm.P))[1],
+                 verbose::Bool=false)
+    p_s_joint_smoothed = Array{Float64, 3}(undef, size(rsm.y, 1), rsm.M, rsm.M)
+    smooth!(rsm, p_s_joint_smoothed, prob_update_pre = prob_update_pre)
+    verbose ? (return p_s_joint_smoothed) : (return nothing) 
+end
+
+"""
+Apply the backward smoother developed by Kim to a regime switching model.
+
+Same as `smooth` except that the result is stored in the perallocated first argument.
+
+##### Arguments
+- `rsm::RegimeSwitchingModel`: `RegimeSwitchingModel` specifying the model.
+- `p_s_joint_smoothed`: `T x n_state x n_state`
+- `prob_update_pre::AbstractArray`: Probability distribution of state at period 0 conditional on the information up to
+                                   period 0.
+
+##### Returns
+- `nothing`
+"""
+function smooth!(rsm::RegimeSwitchingModel,
+                 p_s_joint_smoothed::AbstractArray;
+                 prob_update_pre::AbstractArray = stationary_distributions(MarkovChain(rsm.P))[1],
+                 verbose::Bool=false)
+    y = rsm.y
+    M = rsm.M
+    T = size(y, 1)
+    rsm_tmp = RegimeSwitchingModel(rsm.g, y, rsm.parameter, rsm.P)
+    ps, p_s_update, p_s_forecast = filter!(rsm_tmp, prob_update_pre = prob_update_pre)
+    p_s_init = Vector{Float64}(undef, M)
+    rsm.prob_smoothed[T, :] = p_s_update[T, :]
+    for t = T-1:-1:1
+        for s in 1:M
+            for sp in 1:M
+                p_s_joint_smoothed[t, s, sp] = rsm.prob_smoothed[t+1, sp] * p_s_update[t, s] * rsm.P[s, sp]/p_s_forecast[t+1, sp]
+            end
+            rsm.prob_smoothed[t, s] = sum(p_s_joint_smoothed[t, s, :])
+        end
+    end
+    verbose ? (return p_s_joint_smoothed) : (return nothing) 
+end

test/runtests.jl

@@ -3,7 +3,7 @@ activate(joinpath(@__DIR__, ".."))
 using QuantEcon
 using Base.Iterators: take, cycle
 using DataStructures: counter
-using Distributions: LogNormal, pdf
+using Distributions: Normal, LogNormal, pdf
 using LinearAlgebra
 using Random
 using FFTW

test/test_filter.jl

@@ -20,4 +20,81 @@
         @test isapprox(data["ham_c"], data["ham_c_mat"], nans=true, rtol=1e-7, atol=1e-7)
         @test isapprox(data["ham_rw_c"], data["ham_rw_c_mat"], nans=true)
     end
+
+    @testset "test hamilton regime switching filter" begin
+        # The model here is
+        # $y_t=a_{S_t}+\rho_{S_t}(y_{t−1}−a_{S_{t−1}})+\epsilon_t$
+        # where $\epsilon_t \sim i.i.d. N(0, \sigma_{S_t})$.
+        # Although this model has a history dependent observation equation,
+        # it can be written in the form of i.i.d. observation equation with Markov state transition
+        # by defining state and parameters appropriately.
+        data = [1.5, 2.4, 3.1, 2.9, 3.0]
+        a1 = 1.0
+        a2 = 3.0
+        sigma1 = 0.5
+        sigma2 = 0.6
+        rho1 = 0.1
+        rho2 = 0.2
+        p1 = 0.6
+        p2 = 0.3
+        parameter = (mu = [a1-rho1*a1, a1-rho1*a2, a2-rho2*a1, a2-rho2*a2], 
+                     sigma=[sigma1, sigma1, sigma2, sigma2],
+                     rho=[rho1, rho1, rho2, rho2])
+        P = [p1  0   (1-p1)   0;
+             p1  0   (1-p1)   0;
+             0   p2  0        (1-p2);
+             0   p2  0        (1-p2)]
+        ss = P^100000
+        g(x, s, parameter) = pdf(Normal(parameter.rho[s]*x[2], parameter.sigma[s]), x[1]-parameter.mu[s])
+        y = hcat(data[2:end], data[1:end-1])
+        rsm = RegimeSwitchingModel(g, y, parameter, P)
+        ps, p_s_update, p_s_forecast = QuantEcon.filter!(rsm, prob_update_pre = ss[1, :])
+
+        # compare results with the following python code
+        # data = pd.DataFrame(data=[1.5, 2.4, 3.1, 2.9, 3.0], index=[1,2,3,4,5])
+        # model = sm.tsa.MarkovAutoregression(data, k_regimes=2, order=1, 
+        #                                     switching_ar=True, switching_variance=True)
+        # para = pd.Series()
+        # para['p[0->0]'] = 0.6
+        # para['p[1->0]'] = 0.3
+        # para['const[0]'] = 1
+        # para['const[1]'] = 3
+        # para['sigma2[0]'] = 0.25
+        # para['sigma2[1]'] = 0.36
+        # para['ar.L1[0]'] = 0.1
+        # para['ar.L1[1]'] = 0.2
+        # fil_res_hamilton = model.filter(para)
+        #
+        # values used in the test are obtained as follows
+        # log_likelihood: fil_res_hamilton.llf
+        # filtered_probability: fil_res_hamilton.filtered_marginal_probabilities
+        # forecasted_probability: fil_res_hamilton.predicted_marginal_probabilities
+
+        log_likelihood = -3.5984510703365626
+        @test isapprox(rsm.logL, log_likelihood)
+        filtered_probability = [0.0216771960143715 0.9783228039856280;
+                                0.0000591976220513 0.9999408023779480;
+                                0.0004142177820710 0.9995857822179280;
+                                0.0001598680201190 0.9998401319798810]
+        @test all(isapprox.(p_s_update[:, 1]+p_s_update[:, 2], filtered_probability[:, 1]))
+        @test all(isapprox.(p_s_update[:, 3]+p_s_update[:, 4], filtered_probability[:, 2]))
+        forecasted_probability = [0.4285714285714280 0.5714285714285710;
+                                  0.3065031588043110 0.6934968411956880;
+                                  0.3000177592866150 0.6999822407133840;
+                                  0.3001242653346210 0.6998757346653780]
+        @test all(isapprox.(p_s_forecast[:, 1]+p_s_forecast[:, 2], forecasted_probability[:, 1]))
+        @test all(isapprox.(p_s_forecast[:, 3]+p_s_forecast[:, 4], forecasted_probability[:, 2]))
+
+        # test smoothing
+        # smooth_res_hamilton = mod_hamilton.smooth(para)
+        # `smoothed_probability` is obtained by smooth_res_hamilton.smoothed_marginal_probabilities
+        smoothed_probability = [0.0127846430957809	0.9872153569042180;
+                                0.0000237981350077	0.9999762018649920;
+                                0.0001944102709486	0.9998055897290510;
+                                0.0001598680201190	0.9998401319798810]
+        QuantEcon.smooth!(rsm, prob_update_pre = ss[1, :])
+        @test all(isapprox.(rsm.prob_smoothed[:, 1]+rsm.prob_smoothed[:, 2], smoothed_probability[:, 1]))
+        @test all(isapprox.(rsm.prob_smoothed[:, 3]+rsm.prob_smoothed[:, 4], smoothed_probability[:, 2]))
+
+    end
 end