Feat/ssm ext

.gitignore

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

Project.toml

@@ -5,7 +5,9 @@ authors = ["Ryan Senne"]
 
 [deps]
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
+DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HiddenMarkovModels = "84ca31d5-effc-45e0-bfda-5a68cd981f47"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
@@ -16,15 +18,24 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
 
+[weakdeps]
+SequentialSamplingModels = "0e71a2a6-2b30-4447-8742-d083a85e82d1"
+
+[extensions]
+EmissionModelsSequentialSamplingModelsExt = "SequentialSamplingModels"
+
 [compat]
 DensityInterface = "0.4.0"
+DiffResults = "1.1.0"
 Distributions = "0.25.122"
+ForwardDiff = "1.4.1"
 HiddenMarkovModels = "0.7.1"
 LinearAlgebra = "1"
 LogExpFunctions = "0.3.29, 1"
 NearestNeighbors = "0.4.27"
 Optim = "2"
 Random = "1"
+SequentialSamplingModels = "0.13"
 SpecialFunctions = "2.6.1"
 Statistics = "1"
 StatsAPI = "1.8.0"

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,8 +34,14 @@ 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",
+        "DDM Emissions" => "ddm.md",
         "Priors" => "priors.md",
         "ACDC Model Selection" => "acdc.md",
         "Custom Emission Models" => "custom.md",

docs/src/ddm.md

@@ -0,0 +1,94 @@
+# DDM Emissions
+
+Drift diffusion model (DDM) emissions for two-alternative forced choice data. Pairing them with a hidden Markov model gives the DDM-HMM of Senne et al. (2026): the hidden states are decision-making regimes, and each state emits a `(choice, rt)` trial from its own DDM.
+
+The Wiener first-passage-time density and sampler come from [SequentialSamplingModels.jl](https://github.com/itsdfish/SequentialSamplingModels.jl), a weak dependency. The emission types always construct, but `logdensityof`, `rand`, and `fit!` need the extension loaded:
+
+```julia
+using EmissionModels
+using SequentialSamplingModels   # activates the extension
+```
+
+## Observations and controls
+
+An observation is a `(choice, rt)` pair (a plain tuple or the `(; choice, rt)` `NamedTuple` returned by `rand`), with `choice ∈ {1, 2}` for the upper/lower boundary (matching SequentialSamplingModels.jl) and `rt` in seconds. Each trial also carries a scalar control that sets its drift.
+
+Both models are stimulus-coded: the boundaries mark stimulus/choice identity (e.g. right vs. left), not correct vs. error, so correctness is read off afterward from the trial condition. The starting point ``z`` is therefore a side bias rather than an accuracy bias. The [HSSM stimulus-coding tutorial](https://lnccbrown.github.io/HSSM/tutorials/tutorial_stim_coding/) walks through how this changes parameter interpretation relative to accuracy coding.
+
+## Provided DDM types
+
+### `StimulusCodedDDM(; ν, α, z, τ)`
+
+The drift is the gain ``ν`` multiplied by a trial-specific stimulus code:
+
+``v_{\mathrm{trial}} = s_{\mathrm{trial}} \, ν, \qquad s_{\mathrm{trial}} \in \{-1, +1\}``
+
+The control is the stimulus code ``s`` (`+1` when the upper-boundary stimulus is shown, `-1` for the lower, `0` for no-signal trials). ``ν > 0`` is a magnitude, so the stimulus code carries the sign.
+
+```julia
+d = StimulusCodedDDM(; ν=2.0, α=1.0, z=0.5, τ=0.3)
+obs = rand(rng, d, -1)          # (; choice, rt) on a lower-stimulus trial
+logdensityof(d, obs, -1)
+```
+
+### `CoherenceDDM(; k, γ, α, z, τ)`
+
+The drift is a (possibly nonlinear) function of signed stimulus coherence:
+
+``v_{\mathrm{trial}} = k \, \operatorname{sign}(c) \, |c|^{γ}``
+
+The control is the signed coherence ``c`` (sign = stimulus side, magnitude = stimulus strength); ``γ = 1`` recovers the classic linear drift/coherence relationship. Both ``k > 0`` and ``γ > 0``, so the coherence carries the sign.
+
+```julia
+d = CoherenceDDM(; k=8.0, γ=0.7, α=1.2, z=0.5, τ=0.25)
+obs = rand(rng, d, 0.256)
+logdensityof(d, obs, 0.256)
+```
+
+## Fitting DDM emissions
+
+`fit!` maximizes the weighted log-likelihood ``\sum_i w_i \log p(\mathrm{choice}_i, \mathrm{rt}_i \mid \mathrm{control}_i)`` in place. The density is the Navarro & Fuss (2009) series for the Wiener first passage time; positive parameters are optimized in ``\log`` space and bounded ones (``z``, ``τ``) through a logistic map, so LBFGS runs unconstrained with gradients from ForwardDiff. ``τ`` is capped at the smallest reaction time carrying positive weight. Each call warm-starts from the current parameters, so it refines an EM iterate rather than restarting.
+
+```julia
+fit!(d, obs_seq, weights; control_seq=stimulus_codes, max_iter=100, gtol=1e-8)
+```
+
+## The DDM-HMM
+
+Both types subtype `HiddenMarkovModels.ControlledEmission`, so a vector of them is a valid `dists` for a `ControlledEmissionHMM`:
+
+```julia
+using HiddenMarkovModels
+
+dists = [
+    StimulusCodedDDM(; ν=2.5, α=1.2, z=0.5, τ=0.25),  # engaged
+    StimulusCodedDDM(; ν=0.3, α=0.7, z=0.5, τ=0.20),  # lapsed
+]
+hmm = ControlledEmissionHMM([0.5, 0.5], [0.95 0.05; 0.1 0.9], dists)
+
+sim = rand(rng, hmm, stimulus_codes)                  # simulate trials
+hmm_fit, lls = baum_welch(hmm, obs_seq, stimulus_codes; seq_ends=[T])
+```
+
+## Model selection with ACDC
+
+DDM emissions support [ACDC model selection](acdc.md). Driver recovery inverts each `(choice, rt)` trial through the Rosenblatt transform of its DDM: a randomized PIT on the boundary choice followed by the conditional reaction-time PIT, giving two drivers per trial that are uniform on ``[0,1]`` under a well-specified model.
+
+```julia
+acdc = component_discrepancies(hmm, obs_seq, KSDiscrepancy();
+                               control_seq=stimulus_codes, seq_ends=[T])
+```
+
+## API Reference
+
+```@docs
+EmissionModels.AbstractDDMEmission
+StimulusCodedDDM
+CoherenceDDM
+StatsAPI.fit!(::EmissionModels.AbstractDDMEmission, ::AbstractVector, ::AbstractVector{<:Real})
+```
+
+## References
+
+- Senne, R. A., Xia, H., Duebel, H. F., Do, Q., Kane, G., Fourie, J., Ramirez, S., Scott, B., & DePasquale, B. (2026). Diurnal rhythms of choice: a novel state-dependent drift diffusion model uncovers time-dependent changes in rat decision making. *bioRxiv* [2026.05.25.727672](https://doi.org/10.64898/2026.05.25.727672).
+- Navarro, D. J., & Fuss, I. G. (2009). Fast and accurate calculations for first-passage times in Wiener diffusion models. *Journal of Mathematical Psychology*, 53(4), 222-230.

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