diff --git a/.gitignore b/.gitignore index 039cda9..2e89a79 100644 --- a/.gitignore +++ b/.gitignore @@ -5,6 +5,7 @@ /docs/Manifest*.toml /docs/build/ +/docs/src/examples/ /test/Manifest*.toml diff --git a/docs/Project.toml b/docs/Project.toml index be32084..c2af5cd 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,5 +1,16 @@ [deps] DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" EmissionModels = "1e2dd27c-41a5-43b2-863c-3eddd0c72c67" +HiddenMarkovModels = "84ca31d5-effc-45e0-bfda-5a68cd981f47" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" + +[compat] +Documenter = "1" +HiddenMarkovModels = "0.7.1" +Literate = "2" diff --git a/docs/make.jl b/docs/make.jl index 240631e..f3e976c 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,11 +1,28 @@ using EmissionModels using Documenter using DensityInterface +using Literate using Random using StatsAPI DocMeta.setdocmeta!(EmissionModels, :DocTestSetup, :(using EmissionModels); recursive=true) +# Convert the Literate tutorials into Documenter pages. The `#src` lines +# (embedded tests) are stripped from the generated markdown. +examples_path = joinpath(dirname(@__DIR__), "examples") +examples_md_path = joinpath(@__DIR__, "src", "examples") +mkpath(examples_md_path) +for file in readdir(examples_md_path) + if endswith(file, ".md") + rm(joinpath(examples_md_path, file)) + end +end +for file in readdir(examples_path) + if endswith(file, ".jl") + Literate.markdown(joinpath(examples_path, file), examples_md_path) + end +end + makedocs(; modules=[EmissionModels], authors="Ryan Senne", @@ -17,6 +34,11 @@ makedocs(; ), pages=[ "Home" => "index.md", + "Tutorials" => [ + joinpath("examples", "basics.md"), + joinpath("examples", "glm.md"), + joinpath("examples", "acdc.md"), + ], "Distributions" => "distributions.md", "GLM Emissions" => "glm.md", "Priors" => "priors.md", diff --git a/docs/src/index.md b/docs/src/index.md index 6150a35..1f46361 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -36,6 +36,8 @@ hmm_est, loglik_trace = baum_welch(hmm, obs_seq) ## Where to go from here +- [Tutorials](examples/basics.md): worked examples, from built-in emissions + to GLMs and model selection. - [Distributions](distributions.md): all available count and multivariate emission models. - [GLM Emissions](glm.md): regression-based emissions that take a control diff --git a/examples/acdc.jl b/examples/acdc.jl new file mode 100644 index 0000000..8fb772e --- /dev/null +++ b/examples/acdc.jl @@ -0,0 +1,150 @@ +# # Model selection with ACDC + +#= +Here we use the Accumulated Cutoff Discrepancy Criterion (ACDC, Li et al. +2026) to choose the number of hidden states, and to detect a misspecified +emission model. +=# + +using Distributions +using EmissionModels +using HiddenMarkovModels +using Random +using Statistics +using Test #src + +#- + +rng = MersenneTwister(63); + +# ## Stochastic drivers + +#= +If a fitted model is correctly specified, the probability integral transform +(PIT) of each observation under its generating emission is uniform on +$[0, 1]$. ACDC recovers these *stochastic drivers* state by state, scores each +state by the discrepancy of its drivers from uniformity, and selects the +smallest number of states whose scores all fall below a cutoff $\rho$. Unlike +AIC or BIC, no parameter count is involved. + +Let us generate data from a three-state Gaussian HMM. +=# + +hmm_true = HMM( + [0.4, 0.3, 0.3], + [0.90 0.05 0.05; 0.05 0.90 0.05; 0.05 0.05 0.90], + [Normal(-5.0, 1.0), Normal(0.0, 1.0), Normal(5.0, 1.0)], +) +_, obs_seq = rand(rng, hmm_true, 3000); + +#= +[`stochastic_drivers`](@ref) inverts each emission through the PIT and pools +the drivers by state, assigning each observation according to its posterior +state membership. There is one pool of size $D \times N_k$ per state (the +observations here are scalar, so $D = 1$). +=# + +sd = stochastic_drivers(hmm_true, obs_seq; rng=rng) +size.(sd.ε_pools) + +#= +Under the true model, every pool should look uniform, with mean $1/2$. +=# + +[mean(pool) for pool in sd.ε_pools] + +# ## Choosing the number of states + +#= +We fit candidate models with $K = 1, \dots, 4$ states by Baum-Welch, starting +each from evenly spread quantiles of the data, and score every fitted model +with [`component_discrepancies`](@ref). Here we measure uniformity with the +Kolmogorov-Smirnov statistic ([`KSDiscrepancy`](@ref)); see the +[ACDC page](../acdc.md) for the other available measures. +=# + +function candidate_hmm(K, obs_seq) + stay = K == 1 ? 1.0 : 0.9 + init = fill(1 / K, K) + trans = [i == j ? stay : (1 - stay) / (K - 1) for i in 1:K, j in 1:K] + dists = [Normal(quantile(obs_seq, (2k - 1) / (2K)), 1.0) for k in 1:K] + return HMM(init, trans, dists) +end + +results = map(1:4) do K + hmm_est, _ = baum_welch(candidate_hmm(K, obs_seq), obs_seq) + return component_discrepancies(hmm_est, obs_seq, KSDiscrepancy(); rng=rng) +end +[maximum(r.component_discrepancies) for r in results] + +#= +Underfitted models leave clearly non-uniform drivers, while the three-state +model passes. [`acdc_select`](@ref) turns these scores into a choice: it +minimizes the accumulated loss $\sum_k \max(0, \hat{D}_k - \rho)$, breaking +ties in favor of fewer states (which is how an overfitted but well-calibrated +four-state model loses to the three-state one). +=# + +K_selected = acdc_select(results, 0.1) + +#= +[`get_critical_rho_values`](@ref) lists the cutoffs at which the selection +changes, which is useful to check how sensitive the choice is to $\rho$. +=# + +get_critical_rho_values(results) + +# ## Detecting a misspecified emission + +#= +Because ACDC scores each state separately, it also reveals *which* emission is +wrong, even when the number of states is correct. Let us generate zero-inflated +counts and compare a plain Poisson fit against a zero-inflated one. +=# + +hmm_zi = HMM( + [0.5, 0.5], + [0.95 0.05; 0.05 0.95], + [PoissonZeroInflated(3.0, 0.4), PoissonZeroInflated(12.0, 0.4)], +) +_, obs_zi = rand(rng, hmm_zi, 2000); + +#= +First, a two-state HMM with plain Poisson emissions, which cannot account for +the excess zeros. +=# + +guess_pois = HMM([0.5, 0.5], [0.9 0.1; 0.1 0.9], [Poisson(2.0), Poisson(8.0)]) +hmm_pois, _ = baum_welch(guess_pois, obs_zi) +res_pois = component_discrepancies(hmm_pois, obs_zi, KSDiscrepancy(); rng=rng) +res_pois.component_discrepancies + +#= +Then the correctly specified zero-inflated model. +=# + +guess_zi = HMM( + [0.5, 0.5], + [0.9 0.1; 0.1 0.9], + [PoissonZeroInflated(2.0, 0.3), PoissonZeroInflated(8.0, 0.3)], +) +hmm_zi_est, _ = baum_welch(guess_zi, obs_zi) +res_zi = component_discrepancies(hmm_zi_est, obs_zi, KSDiscrepancy(); rng=rng) +res_zi.component_discrepancies + +#= +The misspecified Poisson states are flagged with large discrepancies, while +the zero-inflated fit passes. +=# + +# ## Tests #src + +@test sum(size.(sd.ε_pools, 2)) == 3000 #src +@test all(pool -> all(0 .<= pool .<= 1), sd.ε_pools) #src +@test all(pool -> isapprox(mean(pool), 0.5; atol=0.05), sd.ε_pools) #src +@test K_selected == 3 #src +@test maximum(results[1].component_discrepancies) > 0.1 #src +@test maximum(results[2].component_discrepancies) > 0.1 #src +@test maximum(results[3].component_discrepancies) < 0.1 #src +@test maximum(res_pois.component_discrepancies) > 0.1 #src +@test maximum(res_zi.component_discrepancies) < 0.1 #src diff --git a/examples/basics.jl b/examples/basics.jl new file mode 100644 index 0000000..90d6ed5 --- /dev/null +++ b/examples/basics.jl @@ -0,0 +1,144 @@ +# # Basics + +#= +Here we show how to use the emission models of this package inside a hidden +Markov model from +[HiddenMarkovModels.jl](https://github.com/JuliaStats/HiddenMarkovModels.jl). +=# + +using Distributions +using EmissionModels +using HiddenMarkovModels +using Random +using Statistics +using Test #src + +#- + +rng = MersenneTwister(63); + +# ## Emission models + +#= +An *emission model* describes how observations are generated conditioned on the +latent state of an HMM. Every distribution in this package implements the +interface that HiddenMarkovModels.jl expects from an observation distribution: + +- `rand(rng, dist)` samples an observation +- `logdensityof(dist, obs)` evaluates the loglikelihood of an observation +- `fit!(dist, obs_seq, weight_seq)` updates the parameters from weighted observations + +Take the zero-inflated Poisson ([`PoissonZeroInflated`](@ref)) as an example. +It models count data with excess zeros: with probability $\pi$ the observation +is a structural zero, otherwise it is drawn from a Poisson with rate $\lambda$. +=# + +dist = PoissonZeroInflated(5.0, 0.3) + +#- + +x = rand(rng, dist) + +#- + +logdensityof(dist, x) + +#= +The `fit!` method performs a weighted maximum-likelihood update, which is +exactly the operation Baum-Welch needs during learning. With unit weights it +reduces to ordinary maximum likelihood. +=# + +obs = [rand(rng, dist) for _ in 1:2000] +dist_est = PoissonZeroInflated(1.0, 0.5) +fit!(dist_est, obs, ones(2000)) +(dist_est.λ, dist_est.π) + +# ## Model + +#= +Because the interface is the standard one, a vector of emission models is a +valid `dists` argument for the `HMM` type of HiddenMarkovModels.jl. Here is a +two-state model for count data: a quiet state with many structural zeros and +an active state with a high rate. +=# + +init = [0.6, 0.4] +trans = [0.9 0.1; 0.2 0.8] +dists = [PoissonZeroInflated(2.0, 0.5), PoissonZeroInflated(15.0, 0.05)] +hmm = HMM(init, trans, dists) + +# ## Simulation + +#= +You can simulate a pair of state and observation sequences with `rand` by +specifying how long you want them to be. +=# + +state_seq, obs_seq = rand(rng, hmm, 1000); + +#- + +obs_seq[1:5] + +# ## Inference + +#= +All the inference algorithms of HiddenMarkovModels.jl work out of the box. +Since the two states are well separated, the Viterbi algorithm recovers most +of the true state sequence. +=# + +best_state_seq, _ = viterbi(hmm, obs_seq); +mean(best_state_seq .== state_seq) + +#= +The loglikelihood of the observation sequence is computed with the forward +algorithm, which `logdensityof` wraps. +=# + +logdensityof(hmm, obs_seq) + +# ## Learning + +#= +The Baum-Welch algorithm is a local optimization procedure, so it requires a +starting point that is close enough to the true model. +=# + +dists_guess = [PoissonZeroInflated(1.0, 0.3), PoissonZeroInflated(10.0, 0.2)] +hmm_guess = HMM([0.5, 0.5], [0.8 0.2; 0.3 0.7], dists_guess); + +#- + +hmm_est, loglikelihood_evolution = baum_welch(hmm_guess, obs_seq); +first(loglikelihood_evolution), last(loglikelihood_evolution) + +#= +The estimated rates and zero-inflation probabilities are close to the truth. +Keep in mind that HMMs are only identifiable up to a permutation of the +states, so in general the estimated states must be matched to the true ones +before comparing parameters. +=# + +hcat( + [(d.λ, d.π) for d in obs_distributions(hmm_est)], + [(d.λ, d.π) for d in obs_distributions(hmm)], +) + +#= +Multivariate emissions like [`MultivariateT`](@ref) work in exactly the same +way, and the [GLM emissions](glm.md) additionally condition on a control +vector. +=# + +# ## Tests #src + +@test dist_est.λ ≈ 5.0 atol = 0.3 #src +@test dist_est.π ≈ 0.3 atol = 0.03 #src +@test mean(best_state_seq .== state_seq) > 0.9 #src +@test all(diff(loglikelihood_evolution) .>= -1e-6) #src +ds_est = obs_distributions(hmm_est) #src +order = sortperm([d.λ for d in ds_est]) #src +@test [d.λ for d in ds_est[order]] ≈ [2.0, 15.0] atol = 1.0 #src +@test [d.π for d in ds_est[order]] ≈ [0.5, 0.05] atol = 0.1 #src diff --git a/examples/glm.jl b/examples/glm.jl new file mode 100644 index 0000000..640b163 --- /dev/null +++ b/examples/glm.jl @@ -0,0 +1,143 @@ +# # GLM emissions + +#= +Here we build a Markov switching regression: an HMM whose emissions are +generalized linear models, conditioned on a control (design) vector at each +time step. +=# + +using Distributions +using EmissionModels +using HiddenMarkovModels +using HiddenMarkovModels: ControlledEmissionHMM +using LinearAlgebra +using Random +using Statistics +using Test #src + +#- + +rng = MersenneTwister(63); + +# ## Model + +#= +A GLM emission maps a control vector $x$ to the parameters of an observation +distribution. For instance [`PoissonGLM`](@ref) models +$y \mid x \sim \mathrm{Poisson}(\exp(\beta^\top x))$. No intercept is added +automatically: if you want one, include a constant entry in $x$. +=# + +glm = PoissonGLM([0.5, -0.3]) + +#= +Outside of any HMM, a GLM can be sampled and evaluated by passing the control +vector as a keyword argument. +=# + +x = [1.0, 2.0] +y = rand(rng, glm; control_seq=x) + +#- + +logdensityof(glm, y; control_seq=x) + +#= +All GLMs subtype `ControlledEmission` from HiddenMarkovModels.jl, so a vector +of them is a valid `dists` argument for a `ControlledEmissionHMM`. Here is a +two-state Markov switching Poisson regression. +=# + +p = 3 +init = [0.6, 0.4] +trans = [0.92 0.08; 0.15 0.85] +dists = [PoissonGLM([0.2, 0.5, -0.3]), PoissonGLM([1.2, -0.4, 0.6])] +hmm = ControlledEmissionHMM(init, trans, dists) + +# ## Simulation + +#= +A controlled HMM cannot be simulated for a bare length `T`: there is no +sensible default control. Provide one control vector per time step instead, +here with a constant first entry acting as the intercept. +=# + +T = 1000 +control_seq = [vcat(1.0, randn(rng, p - 1)) for _ in 1:T]; +state_seq, obs_seq = rand(rng, hmm, control_seq); + +#- + +obs_seq[1:5] + +# ## Inference + +#= +Inference algorithms take the control sequence as an additional positional +argument. +=# + +best_state_seq, _ = viterbi(hmm, obs_seq, control_seq); +mean(best_state_seq .== state_seq) + +#- + +logdensityof(hmm, obs_seq, control_seq) + +# ## Learning + +#= +During Baum-Welch, each state's `fit!` solves a weighted GLM problem: a +closed-form weighted least squares for [`GaussianGLM`](@ref) and +[`MvGaussianGLM`](@ref), a Newton solve for the others. As always with a +local optimization procedure, start from a guess close enough to the truth. +=# + +dists_guess = [PoissonGLM([0.0, 0.3, -0.1]), PoissonGLM([1.0, -0.2, 0.4])] +hmm_guess = ControlledEmissionHMM([0.5, 0.5], [0.8 0.2; 0.3 0.7], dists_guess); + +#- + +hmm_est, loglikelihood_evolution = baum_welch(hmm_guess, obs_seq, control_seq); +first(loglikelihood_evolution), last(loglikelihood_evolution) + +#= +How did we perform? +=# + +hcat(hmm_est.dists[1].β, hmm.dists[1].β) + +#- + +hcat(hmm_est.dists[2].β, hmm.dists[2].β) + +# ## Priors + +#= +Every GLM accepts an optional prior on the coefficients, which turns the +M-step into a penalized (MAP) update. [`RidgePrior`](@ref) adds an $\ell_2$ +penalty $\frac{\lambda}{2} \lVert \beta \rVert^2$, which is useful when +controls are correlated or data per state is scarce. Fitting the same data +with and without a ridge prior shrinks the coefficients toward zero. +=# + +control_mat = permutedims(reduce(hcat, control_seq)) +glm_ridge = PoissonGLM(zeros(p), RidgePrior(10.0)) +glm_flat = PoissonGLM(zeros(p)) +fit!(glm_ridge, obs_seq, ones(T); control_seq=control_mat) +fit!(glm_flat, obs_seq, ones(T); control_seq=control_mat) +norm(glm_ridge.β), norm(glm_flat.β) + +#= +The other GLM families ([`GaussianGLM`](@ref), [`BernoulliGLM`](@ref), +[`MultinomialGLM`](@ref) and the multivariate variants) follow the same +interface. +=# + +# ## Tests #src + +@test mean(best_state_seq .== state_seq) > 0.7 #src +@test all(diff(loglikelihood_evolution) .>= -1e-6) #src +@test hmm_est.dists[1].β ≈ hmm.dists[1].β atol = 0.2 #src +@test hmm_est.dists[2].β ≈ hmm.dists[2].β atol = 0.2 #src +@test norm(glm_ridge.β) < norm(glm_flat.β) #src diff --git a/test/runtests.jl b/test/runtests.jl index f56c6fc..eb94249 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -35,6 +35,20 @@ using JuliaFormatter include("acdc/test_acdc.jl") end + #= The Literate tutorials double as integration tests: their `@test` lines + are hidden from the rendered docs with `#src` but run here as ordinary + code. =# + @testset "Examples" begin + examples_path = joinpath(dirname(@__DIR__), "examples") + for file in readdir(examples_path) + if endswith(file, ".jl") + @testset "Example - $file" begin + include(joinpath(examples_path, file)) + end + end + end + end + @testset "Allocations" begin include("allocations.jl") end