Add Literate tutorials with embedded tests

.gitignore

@@ -5,6 +5,7 @@
 
 /docs/Manifest*.toml
 /docs/build/
+/docs/src/examples/
 
 /test/Manifest*.toml
 

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"

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",

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

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

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