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/Project.toml b/Project.toml index ab5ba56..d0b19fc 100644 --- a/Project.toml +++ b/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" 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..879bc85 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,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", diff --git a/docs/src/ddm.md b/docs/src/ddm.md new file mode 100644 index 0000000..cf512a6 --- /dev/null +++ b/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. 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/ext/EmissionModelsSequentialSamplingModelsExt.jl b/ext/EmissionModelsSequentialSamplingModelsExt.jl new file mode 100644 index 0000000..e6c487b --- /dev/null +++ b/ext/EmissionModelsSequentialSamplingModelsExt.jl @@ -0,0 +1,48 @@ +#= Backs the DDM emission hooks (density, sampler, CDF) with + SequentialSamplingModels' Wiener implementation. Everything else lives in + src/ssm/ddm.jl. =# +module EmissionModelsSequentialSamplingModelsExt + +using EmissionModels: EmissionModels +using SequentialSamplingModels: SequentialSamplingModels, DDM + +function EmissionModels._ddm_logpdf( + ν::Real, α::Real, z::Real, τ::Real, choice::Real, rt::Real +) + T = float(promote_type(typeof(ν), typeof(α), typeof(z), typeof(τ), typeof(rt))) + # Cases SequentialSamplingModels.pdf rejects; here they have zero density. + (choice == 1 || choice == 2) || return T(-Inf) + rt > τ || return T(-Inf) + #= Promote parameters and rt together before handing them to SSM: its + series internals require one homogeneous element type, and mixed + precisions (e.g. Float32 emission, Float64 rt) would MethodError. + The series can also overflow at extreme parameters; treat both as + zero density. =# + p = try + SequentialSamplingModels.pdf(DDM(T(ν), T(α), T(z), T(τ)), Int(choice), T(rt)) + catch + return T(-Inf) + end + return (isfinite(p) && p > 0) ? T(log(p)) : T(-Inf) +end + +function EmissionModels._ddm_rand(rng, ν::Real, α::Real, z::Real, τ::Real) + return rand(rng, DDM(ν, α, z, τ)) +end + +function EmissionModels._ddm_cdf(ν::Real, α::Real, z::Real, τ::Real, choice::Real, rt::Real) + T = float(promote_type(typeof(ν), typeof(α), typeof(z), typeof(τ), typeof(rt))) + # Cases SequentialSamplingModels.cdf rejects; here they carry zero mass. + (choice == 1 || choice == 2) || return zero(T) + rt > τ || return zero(T) + # Same homogeneous-type promotion and overflow guard as `_ddm_logpdf`. + F = try + SequentialSamplingModels.cdf(DDM(T(ν), T(α), T(z), T(τ)), Int(choice), T(rt)) + catch + return zero(T) + end + # Series truncation can stray slightly outside [0, 1]. + return isfinite(F) ? clamp(T(F), zero(T), one(T)) : zero(T) +end + +end diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index d6684f4..7e1abea 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -5,13 +5,15 @@ using Distributions: cdf, quantile using Distributions: ContinuousUnivariateDistribution, DiscreteUnivariateDistribution, AbstractMvNormal using DensityInterface -using HiddenMarkovModels: ControlledEmission +using DiffResults: DiffResults +using ForwardDiff: ForwardDiff +using HiddenMarkovModels: ControlledEmission, ControlBoundEmission using HiddenMarkovModels: AbstractHMM, obs_distributions, forward_backward using LinearAlgebra -using LogExpFunctions: logaddexp, logsumexp, log1pexp, logistic +using LogExpFunctions: logaddexp, logsumexp, log1pexp, logistic, logit using NearestNeighbors: KDTree, knn -using Optim: Optim, optimize, TwiceDifferentiable, Newton -using Optim.NLSolversBase: only_fgh! +using Optim: Optim, optimize, OnceDifferentiable, TwiceDifferentiable, Newton, LBFGS +using Optim.NLSolversBase: only_fgh!, only_fg! using Random using SpecialFunctions: loggamma, digamma, trigamma using Statistics: mean, var, cov @@ -21,6 +23,7 @@ using StatsAPI: fit! include("zeroinflated/poisson.jl") include("multivariate/t.jl") include("glms/glm.jl") +include("ssm/ddm.jl") include("acdc/interface.jl") include("acdc/drivers.jl") include("acdc/hmm.jl") @@ -32,6 +35,7 @@ export MultivariateT, MultivariateTDiag export GaussianGLM, BernoulliGLM, PoissonGLM, MultinomialGLM export MvGaussianGLM, MvBernoulliGLM, MvPoissonGLM export AbstractPrior, NoPrior, RidgePrior +export StimulusCodedDDM, CoherenceDDM export neglogprior, neglogprior_grad!, neglogprior_hess! # ACDC model selection diff --git a/src/acdc/drivers.jl b/src/acdc/drivers.jl index 10d098b..e3b87e7 100644 --- a/src/acdc/drivers.jl +++ b/src/acdc/drivers.jl @@ -201,6 +201,35 @@ function _emission_to_driver(rng::AbstractRNG, g::MultinomialGLM, obs::Real, x) return _emission_to_driver(rng, g, _OneHot(Int(obs), g.out_dim), x) end +#= DDM emission (control -> drift): Rosenblatt transform of the (choice, rt) pair + into a randomized PIT on the boundary choice and the conditional RT PIT, both + uniform on (0,1) under the true model. Uses the extension's defective CDF. =# +function _emission_to_driver(rng::AbstractRNG, d::AbstractDDMEmission, obs, control) + ν = _drift(d, control) + choice = obs[1] + rt = obs[2] + T = float(promote_type(typeof(ν), typeof(d.α), typeof(rt))) + # P(choice = 1), the t → ∞ limit of the defective CDF, in closed form. + p1 = _clamp01(T(_ddm_prob_upper(ν, d.α, d.z))) + Fc = T(_ddm_cdf(ν, d.α, d.z, d.τ, choice, rt)) # P(choice, RT ≤ rt) + + if choice == 1 + lower, upper, p_choice = zero(T), p1, p1 + else + lower, upper, p_choice = p1, one(T), one(T) - p1 + end + ε_choice = lower + rand(rng, T) * (upper - lower) + # A never-hit boundary has an uninformative conditional RT: draw uniformly. + ε_rt = p_choice > 0 ? _clamp01(Fc / p_choice) : rand(rng, T) + return [_clamp01(ε_choice), ε_rt] +end + +#= `obs_distributions` on a ControlledEmissionHMM yields ControlBoundEmissions; + unwrap to the emission and its control for the conditional PIT methods. =# +function _emission_to_driver(rng::AbstractRNG, ce::ControlBoundEmission, obs, control) + return _emission_to_driver(rng, ce.dist, obs, control) +end + # 3-arg form on a GLM: no covariate to condition on, so point at the 4-arg hook. function _emission_to_driver(::AbstractRNG, g::AbstractGLM, obs) throw( diff --git a/src/ssm/ddm.jl b/src/ssm/ddm.jl new file mode 100644 index 0000000..ce9914e --- /dev/null +++ b/src/ssm/ddm.jl @@ -0,0 +1,375 @@ +""" + AbstractDDMEmission <: HiddenMarkovModels.ControlledEmission + +Abstract type for drift-diffusion-model (DDM) emission distributions whose +drift rate depends on a scalar, trial-specific control (a stimulus code or a +signed coherence). + +Subtyping `ControlledEmission` lets a `Vector` of DDM emissions serve as the +`dists` of a `HiddenMarkovModels.ControlledEmissionHMM` (the DDM-HMM). Each +concrete type implements the `ControlledEmission` positional interface: + +- `DensityInterface.logdensityof(d, obs, control)`: log density of one trial +- `Random.rand(rng, d, control)`: sample one trial as `(; choice, rt)` +- `StatsAPI.fit!(d, obs_seq, control_seq, weights)`: weighted in-place update + +Observations are `(choice, rt)` pairs — any positionally indexable pair such +as a `Tuple` or the `(; choice, rt)` `NamedTuple` returned by `rand` — with +`choice ∈ {1, 2}` (1 = upper boundary, 2 = lower boundary, matching +SequentialSamplingModels.jl) and `rt` the reaction time in seconds. +`DensityKind` is inherited from `ControlledEmission`. + +The Wiener first-passage-time density and sampler are provided by +SequentialSamplingModels.jl through a package extension: run +`using SequentialSamplingModels` to enable `logdensityof`, `rand`, and `fit!`. +""" +abstract type AbstractDDMEmission <: ControlledEmission end + +# These hooks are overloaded in the SequentialSamplingModels extension. +@noinline function _require_ssm() + throw( + ArgumentError( + "DDM emission models require SequentialSamplingModels.jl: run " * + "`using SequentialSamplingModels` to enable `logdensityof`, " * + "`rand`, and `fit!`.", + ), + ) +end +_ddm_logpdf(ν, α, z, τ, choice, rt) = _require_ssm() +_ddm_rand(rng, ν, α, z, τ) = _require_ssm() +# Defective CDF P(choice, RT ≤ rt), used by ACDC driver recovery. +_ddm_cdf(ν, α, z, τ, choice, rt) = _require_ssm() + +#= Probability of absorption at the upper boundary (choice = 1), the t → ∞ + limit of the defective CDF. Standard Wiener first-passage result for unit + diffusion, boundaries at 0 and α, start z·α: + + P(upper) = (1 - e^{-2ναz}) / (1 - e^{-2να}) + + Written with expm1 for accuracy near ν = 0 and mirrored through + P(upper; ν, z) = 1 - P(upper; -ν, 1-z) for ν < 0 so the exponentials never + overflow. Pure closed form, so it needs no extension hook. =# +function _ddm_prob_upper(ν::Real, α::Real, z::Real) + T = float(promote_type(typeof(ν), typeof(α), typeof(z))) + iszero(ν) && return T(z) + if ν > 0 + return T(expm1(-2 * ν * α * z) / expm1(-2 * ν * α)) + end + return T(1 - expm1(2 * ν * α * (1 - z)) / expm1(2 * ν * α)) +end + +""" + StimulusCodedDDM{T<:Real} <: AbstractDDMEmission + +Stimulus-coded drift diffusion model emission for two-alternative forced +choice data. + +Under stimulus coding the two boundaries correspond to stimulus/choice +identity (e.g. right vs. left) rather than correct vs. error, and the drift +on each trial is the drift magnitude multiplied by a trial-specific stimulus +code: + +`v_trial = s_trial · ν, s_trial ∈ {-1, +1}` + +The control passed to `logdensityof`/`rand`/`fit!` is the scalar stimulus +code `s` (`+1` when the upper-boundary stimulus is presented, `-1` for the +lower; `0` is valid for no-signal trials). Correctness is determined after +the fact from the trial condition. + +# Fields +- `ν::T`: drift magnitude (ν > 0); the stimulus code carries the sign +- `α::T`: boundary separation (α > 0) +- `z::T`: relative starting point (0 < z < 1; 0.5 is unbiased). Stimulus coding + makes this a side bias, not an accuracy bias. +- `τ::T`: non-decision time in seconds (τ ≥ 0) + +# Example +```julia +using EmissionModels, SequentialSamplingModels + +d = StimulusCodedDDM(; ν=2.0, α=1.0, z=0.5, τ=0.3) +obs = rand(rng, d, -1) # (; choice, rt) for a lower stimulus +logdensityof(d, obs, -1) +fit!(d, obs_seq, weights; control_seq=stimulus_codes) +``` +""" +mutable struct StimulusCodedDDM{T<:Real} <: AbstractDDMEmission + ν::T + α::T + z::T + τ::T + + function StimulusCodedDDM{T}(ν::T, α::T, z::T, τ::T) where {T<:Real} + ν > 0 || throw(ArgumentError("ν must be positive, got $ν")) + α > 0 || throw(ArgumentError("α must be positive, got $α")) + 0 < z < 1 || throw(ArgumentError("z must be in (0,1), got $z")) + τ >= 0 || throw(ArgumentError("τ must be non-negative, got $τ")) + return new{T}(ν, α, z, τ) + end +end + +function StimulusCodedDDM(ν::Real, α::Real, z::Real, τ::Real) + T = float(promote_type(typeof(ν), typeof(α), typeof(z), typeof(τ))) + return StimulusCodedDDM{T}(T(ν), T(α), T(z), T(τ)) +end +StimulusCodedDDM(; ν=1.0, α=1.0, z=0.5, τ=0.3) = StimulusCodedDDM(ν, α, z, τ) + +""" + CoherenceDDM{T<:Real} <: AbstractDDMEmission + +Drift diffusion model emission whose drift is a (possibly nonlinear) function +of signed stimulus coherence: + +`v_trial = k · sign(c) · |c|^γ` + +where `c` is the trial's signed coherence (sign = stimulus side, magnitude = +stimulus strength), `k` is the drift gain, and `γ` controls the nonlinearity +(`γ = 1` recovers the classic linear drift-coherence relationship). The +control passed to `logdensityof`/`rand`/`fit!` is the scalar signed coherence +`c`; a no-signal trial (`c = 0`) has zero drift. As with +[`StimulusCodedDDM`](@ref), the boundaries are stimulus-coded. + +# Fields +- `k::T`: drift gain (k > 0); the signed coherence carries the sign +- `γ::T`: coherence exponent (γ > 0) +- `α::T`: boundary separation (α > 0) +- `z::T`: relative starting point / side bias (0 < z < 1) +- `τ::T`: non-decision time in seconds (τ ≥ 0) + +# Example +```julia +using EmissionModels, SequentialSamplingModels + +d = CoherenceDDM(; k=8.0, γ=0.7, α=1.2, z=0.5, τ=0.25) +obs = rand(rng, d, 0.256) # (; choice, rt) at coherence 0.256 +logdensityof(d, obs, 0.256) +fit!(d, obs_seq, weights; control_seq=coherences) +``` +""" +mutable struct CoherenceDDM{T<:Real} <: AbstractDDMEmission + k::T + γ::T + α::T + z::T + τ::T + + function CoherenceDDM{T}(k::T, γ::T, α::T, z::T, τ::T) where {T<:Real} + k > 0 || throw(ArgumentError("k must be positive, got $k")) + γ > 0 || throw(ArgumentError("γ must be positive, got $γ")) + α > 0 || throw(ArgumentError("α must be positive, got $α")) + 0 < z < 1 || throw(ArgumentError("z must be in (0,1), got $z")) + τ >= 0 || throw(ArgumentError("τ must be non-negative, got $τ")) + return new{T}(k, γ, α, z, τ) + end +end + +function CoherenceDDM(k::Real, γ::Real, α::Real, z::Real, τ::Real) + T = float(promote_type(typeof(k), typeof(γ), typeof(α), typeof(z), typeof(τ))) + return CoherenceDDM{T}(T(k), T(γ), T(α), T(z), T(τ)) +end +CoherenceDDM(; k=1.0, γ=1.0, α=1.0, z=0.5, τ=0.3) = CoherenceDDM(k, γ, α, z, τ) + +_signedpow(c::Real, γ::Real) = sign(c) * abs(c)^γ + +_drift(d::StimulusCodedDDM, s::Real) = d.ν * s +_drift(d::CoherenceDDM, c::Real) = d.k * _signedpow(c, d.γ) + +function DensityInterface.logdensityof(d::AbstractDDMEmission, obs, control::Real) + return _ddm_logpdf(_drift(d, control), d.α, d.z, d.τ, obs[1], obs[2]) +end + +function Random.rand(rng::AbstractRNG, d::AbstractDDMEmission, control::Real) + return _ddm_rand(rng, _drift(d, control), d.α, d.z, d.τ) +end + +# Disambiguates from Base's `rand(rng, S, dims::Integer...)` for integer codes. +function Random.rand(rng::AbstractRNG, d::AbstractDDMEmission, control::Integer) + return _ddm_rand(rng, _drift(d, control), d.α, d.z, d.τ) +end + +#= `_pack`/`_unpack` map a model's fields to and from the unconstrained vector θ + the optimizer sees: ν, k, γ, α through exp; z through logistic; τ through a + logistic scaled into (0, rt_min). =# +function _pack(d::StimulusCodedDDM{T}, rt_min::T) where {T<:Real} + ϵ = sqrt(eps(T)) + return T[log(d.ν), log(d.α), logit(d.z), logit(clamp(d.τ / rt_min, ϵ, 1 - ϵ))] +end + +function _unpack(::StimulusCodedDDM, θ::AbstractVector{<:Real}, rt_min::Real) + # logistic saturates to exactly 0/1 for large |θ|; keep z interior + ϵ = eps(float(typeof(rt_min))) + return ( + ν=exp(θ[1]), + α=exp(θ[2]), + z=clamp(logistic(θ[3]), ϵ, 1 - ϵ), + τ=rt_min * logistic(θ[4]), + ) +end + +_drift_at(::StimulusCodedDDM, pars, s::Real) = pars.ν * s + +function _setparams!(d::StimulusCodedDDM, pars) + d.ν = pars.ν + d.α = pars.α + d.z = pars.z + d.τ = pars.τ + return d +end + +function _pack(d::CoherenceDDM{T}, rt_min::T) where {T<:Real} + ϵ = sqrt(eps(T)) + return T[log(d.k), log(d.γ), log(d.α), logit(d.z), logit(clamp(d.τ / rt_min, ϵ, 1 - ϵ))] +end + +function _unpack(::CoherenceDDM, θ::AbstractVector{<:Real}, rt_min::Real) + ϵ = eps(float(typeof(rt_min))) + return ( + k=exp(θ[1]), + γ=exp(θ[2]), + α=exp(θ[3]), + z=clamp(logistic(θ[4]), ϵ, 1 - ϵ), + τ=rt_min * logistic(θ[5]), + ) +end + +_drift_at(::CoherenceDDM, pars, c::Real) = pars.k * _signedpow(c, pars.γ) + +function _setparams!(d::CoherenceDDM, pars) + d.k = pars.k + d.γ = pars.γ + d.α = pars.α + d.z = pars.z + d.τ = pars.τ + return d +end + +# Weighted negative log-likelihood in θ; a zero-density trial returns +Inf. +struct _DDMNLL{ + D<:AbstractDDMEmission, + O<:AbstractVector, + W<:AbstractVector{<:Real}, + C<:AbstractVector{<:Real}, + T<:Real, +} + d::D + obs_seq::O + weight_seq::W + control_seq::C + rt_min::T +end + +function (o::_DDMNLL)(θ::AbstractVector{T}) where {T<:Real} + pars = _unpack(o.d, θ, o.rt_min) + nll = zero(T) + for i in eachindex(o.obs_seq, o.weight_seq, o.control_seq) + w = o.weight_seq[i] + w > 0 || continue + obs = o.obs_seq[i] + ν = _drift_at(o.d, pars, o.control_seq[i]) + lp = _ddm_logpdf(ν, pars.α, pars.z, pars.τ, obs[1], obs[2]) + isfinite(lp) || return T(Inf) + nll -= T(w) * lp + end + return nll +end + +# Fused value+gradient in the `fg!(F, G, θ)` form `Optim.only_fg!` expects. +struct _DDMFG{N<:_DDMNLL,C<:ForwardDiff.GradientConfig} + nll::N + cfg::C +end + +function (o::_DDMFG)(F, G, θ::AbstractVector{T}) where {T<:Real} + if G !== nothing + dr = DiffResults.MutableDiffResult(zero(T), (G,)) + ForwardDiff.gradient!(dr, o.nll, θ, o.cfg) + return F === nothing ? nothing : DiffResults.value(dr) + end + return F === nothing ? nothing : o.nll(θ) +end + +""" + fit!(d::AbstractDDMEmission, obs_seq, weight_seq; + control_seq, max_iter=100, gtol=1e-8) + +Fit the DDM emission parameters to weighted `(choice, rt)` observations by +maximizing `Σᵢ wᵢ · log p(choiceᵢ, rtᵢ | controlᵢ)` in place, warm-started +from the current parameters (so repeated EM calls refine, not restart). + +The optimizer is Optim's LBFGS over unconstrained transformed parameters, +with the gradient supplied by ForwardDiff through the first-passage density. +The non-decision time is constrained to `(0, rt_min)` where `rt_min` is the +smallest reaction time with positive weight. If no positive weight is present +the parameters are left unchanged. + +# Arguments +- `d`: [`StimulusCodedDDM`](@ref) or [`CoherenceDDM`](@ref), updated in place +- `obs_seq`: sequence of `(choice, rt)` pairs with `choice ∈ {1, 2}` +- `weight_seq`: per-observation weights (e.g. HMM posterior state + probabilities) +- `control_seq`: per-observation scalar controls (stimulus codes or signed + coherences) +- `max_iter`: LBFGS iteration cap per call +- `gtol`: gradient-norm convergence tolerance +""" +function StatsAPI.fit!( + d::AbstractDDMEmission, + obs_seq::AbstractVector, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractVector{<:Real}, + max_iter::Int=100, + gtol::Real=1e-8, +) + n = length(obs_seq) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ obs_seq length $n") + ) + length(control_seq) == n || throw( + DimensionMismatch("control_seq length $(length(control_seq)) ≠ obs_seq length $n"), + ) + + T = typeof(d.α) + + # Only positive-weight trials constrain the fit and bound τ from above. + total_weight = zero(T) + rt_min = T(Inf) + for i in 1:n + w = weight_seq[i] + w > 0 || continue + total_weight += T(w) + obs = obs_seq[i] + choice = obs[1] + rt = obs[2] + (choice == 1 || choice == 2) || throw( + ArgumentError( + "observations must be (choice, rt) pairs with choice ∈ {1, 2}, " * + "got choice = $choice", + ), + ) + rt > 0 || throw(ArgumentError("reaction times must be positive, got $rt")) + rt_min = min(rt_min, T(rt)) + end + total_weight > 0 || return d + + θ = _pack(d, rt_min) + nll = _DDMNLL(d, obs_seq, weight_seq, control_seq, rt_min) + fg = _DDMFG(nll, ForwardDiff.GradientConfig(nll, θ)) + od = OnceDifferentiable(only_fg!(fg), θ) + result = optimize(od, θ, LBFGS(), Optim.Options(; iterations=max_iter, g_abstol=gtol)) + # Keep the current parameters if no finite-likelihood optimum was found. + if isfinite(Optim.minimum(result)) + _setparams!(d, _unpack(d, Optim.minimizer(result), rt_min)) + end + return d +end + +# ControlledEmission positional fit signature; delegates to the keyword method. +function StatsAPI.fit!( + d::AbstractDDMEmission, + obs_seq::AbstractVector, + control_seq::AbstractVector{<:Real}, + weights::AbstractVector{<:Real}; + kwargs..., +) + return fit!(d, obs_seq, weights; control_seq=control_seq, kwargs...) +end diff --git a/test/Project.toml b/test/Project.toml index 21a95b6..dccf8f5 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -8,6 +8,7 @@ JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SequentialSamplingModels = "0e71a2a6-2b30-4447-8742-d083a85e82d1" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" @@ -17,3 +18,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" HiddenMarkovModels = "0.7.1" JET = "0.9, 0.10, 0.11" JuliaFormatter = "1.0.62" +SequentialSamplingModels = "0.13" diff --git a/test/acdc/test_acdc.jl b/test/acdc/test_acdc.jl index b7c3aa9..e60e4a9 100644 --- a/test/acdc/test_acdc.jl +++ b/test/acdc/test_acdc.jl @@ -210,6 +210,30 @@ end @test_throws ArgumentError EM._emission_to_driver(rng, GaussianGLM([0.5], 1.0), 1.0) end +@testset "stochastic_drivers on a ControlledEmissionHMM" begin + #= End-to-end regression: driver recovery on a controlled HMM must unwrap the + ControlBoundEmission that obs_distributions binds to each control. Under + the generating model the recovered drivers are near-uniform. =# + rng = Random.MersenneTwister(5) + p, T = 3, 3000 + dists = [PoissonGLM([0.2, 0.5, -0.3]), PoissonGLM([1.2, -0.4, 0.6])] + hmm = ControlledEmissionHMM([0.6, 0.4], [0.9 0.1; 0.1 0.9], dists) + control_seq = [vcat(1.0, randn(rng, p - 1)) for _ in 1:T] + obs_seq = rand(rng, hmm, control_seq).obs_seq + + acdc = component_discrepancies( + hmm, + obs_seq, + KSDiscrepancy(); + control_seq=control_seq, + seq_ends=(T,), + rng=Random.MersenneTwister(1), + ) + @test length(acdc.component_discrepancies) == 2 + @test all(isfinite, acdc.component_discrepancies) + @test maximum(acdc.component_discrepancies) < 0.1 +end + @testset "Discrepancy edge cases" begin rng = Random.MersenneTwister(123) diff --git a/test/allocations.jl b/test/allocations.jl index 3416f23..8099bb2 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -4,6 +4,7 @@ using Random using LinearAlgebra using DensityInterface using StatsAPI +using SequentialSamplingModels: SequentialSamplingModels #= Allocation regression tests. Warm up first, then call the operation many @@ -72,6 +73,16 @@ for _ in 1:n s += out[1] end; s) +bench_logd_ddm(d, y, c, n) = (s = 0.0; +for _ in 1:n + s += logdensityof(d, y, c) +end; +s) +bench_rand_ddm(rng, d, c, n) = (s = 0.0; +for _ in 1:n + s += rand(rng, d, c).rt +end; +s) @testset "Allocations (steady state)" begin rng = Random.MersenneTwister(0) @@ -264,6 +275,36 @@ s) @test (@allocated fit!(mvt2, obs, w; max_iter=5)) ≤ 1_000_000 end + @testset "DDM emissions" begin + d = StimulusCodedDDM(; ν=2.0, α=1.0, z=0.5, τ=0.2) + + #= Each density evaluation constructs SSM's mutable `DDM` (twice for + choice 1, which flips the boundary) that is SSM's public API, so + the bound is a small per-call constant, not zero. =# + bench_logd_ddm(d, (1, 0.6), 1.0, 1) + bench_logd_ddm(d, (2, 0.6), -1.0, 1) + @test (@allocated bench_logd_ddm(d, (1, 0.6), 1.0, REPS)) ≤ 256 * REPS + @test (@allocated bench_logd_ddm(d, (2, 0.6), -1.0, REPS)) ≤ 256 * REPS + + # rand: one DDM construction; SSM's rejection sampler is scalar math. + bench_rand_ddm(rng, d, 1.0, 1) + @test (@allocated bench_rand_ddm(rng, d, 1.0, REPS)) ≤ 256 * REPS + + #= fit!: every NLL/gradient evaluation rebuilds a Dual-typed DDM per + observation, so allocations scale with n × line-search evaluations + for as long as the density goes through SSM's mutable struct. + Regression guard against something worse, not a zero-alloc claim. =# + rngf = Random.MersenneTwister(1) + n = 200 + controls = [rand(rngf, (-1.0, 1.0)) for _ in 1:n] + obs = [rand(rngf, d, controls[i]) for i in 1:n] + w = ones(n) + d2 = StimulusCodedDDM(; ν=1.5, α=0.9, z=0.5, τ=0.15) + fit!(d2, obs, w; control_seq=controls, max_iter=5) + d2 = StimulusCodedDDM(; ν=1.5, α=0.9, z=0.5, τ=0.15) + @test (@allocated fit!(d2, obs, w; control_seq=controls, max_iter=5)) ≤ 10_000_000 + end + @testset "MultivariateTDiag" begin d = 2 mvtd = MultivariateTDiag([0.0, 0.0], [1.0, 1.0], 5.0) diff --git a/test/runtests.jl b/test/runtests.jl index f56c6fc..3bdfe87 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -31,10 +31,28 @@ using JuliaFormatter include("multivariate/test_t.jl") end + @testset "DDM emissions (SequentialSamplingModels extension)" begin + include("ssm/test_ddm.jl") + end + @testset "ACDC model selection" begin 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 diff --git a/test/ssm/test_ddm.jl b/test/ssm/test_ddm.jl new file mode 100644 index 0000000..2c88653 --- /dev/null +++ b/test/ssm/test_ddm.jl @@ -0,0 +1,355 @@ +using EmissionModels +using EmissionModels: stochastic_drivers, component_discrepancies, KSDiscrepancy +using SequentialSamplingModels +using SequentialSamplingModels: DDM +using HiddenMarkovModels: ControlledEmission, ControlledEmissionHMM, baum_welch, forward +using DensityInterface +using StatsAPI +using Statistics: mean, var +using Random +using Test + +const EM = EmissionModels + +@testset "DDM emissions (SequentialSamplingModels extension)" begin + @testset "construction and validation" begin + d = StimulusCodedDDM(; ν=2.0, α=1.0, z=0.4, τ=0.3) + @test d isa StimulusCodedDDM{Float64} + @test (d.ν, d.α, d.z, d.τ) == (2.0, 1.0, 0.4, 0.3) + + # promoting positional constructor + @test StimulusCodedDDM(2, 1, 0.5, 0.3) isa StimulusCodedDDM{Float64} + + c = CoherenceDDM(; k=8.0, γ=0.7, α=1.2, z=0.5, τ=0.25) + @test c isa CoherenceDDM{Float64} + @test (c.k, c.γ, c.α, c.z, c.τ) == (8.0, 0.7, 1.2, 0.5, 0.25) + + # gains are magnitudes: the control carries the sign + @test_throws ArgumentError StimulusCodedDDM(; ν=-1.0) + @test_throws ArgumentError StimulusCodedDDM(; ν=0.0) + @test_throws ArgumentError CoherenceDDM(; k=-1.0) + @test_throws ArgumentError CoherenceDDM(; k=0.0) + @test_throws ArgumentError StimulusCodedDDM(; α=-1.0) + @test_throws ArgumentError StimulusCodedDDM(; z=0.0) + @test_throws ArgumentError StimulusCodedDDM(; z=1.0) + @test_throws ArgumentError StimulusCodedDDM(; τ=-0.1) + @test_throws ArgumentError CoherenceDDM(; γ=0.0) + + for D in (StimulusCodedDDM, CoherenceDDM) + @test D <: ControlledEmission + end + end + + @testset "logdensityof matches SequentialSamplingModels" begin + d = StimulusCodedDDM(; ν=1.5, α=1.0, z=0.4, τ=0.2) + # stimulus code multiplies the drift gain + @test logdensityof(d, (1, 0.6), 1.0) ≈ logpdf(DDM(1.5, 1.0, 0.4, 0.2), 1, 0.6) + @test logdensityof(d, (2, 0.6), -1.0) ≈ logpdf(DDM(-1.5, 1.0, 0.4, 0.2), 2, 0.6) + # zero-drift no-signal trial + @test logdensityof(d, (1, 0.6), 0) ≈ logpdf(DDM(0.0, 1.0, 0.4, 0.2), 1, 0.6) + + # NamedTuple observations (as returned by rand) index positionally + @test logdensityof(d, (; choice=1, rt=0.6), 1.0) == logdensityof(d, (1, 0.6), 1.0) + + # coherence model: v = k * sign(c) * |c|^γ + c = CoherenceDDM(; k=8.0, γ=0.7, α=1.2, z=0.5, τ=0.25) + v = 8.0 * 0.256^0.7 + @test logdensityof(c, (1, 0.7), 0.256) ≈ logpdf(DDM(v, 1.2, 0.5, 0.25), 1, 0.7) + @test logdensityof(c, (2, 0.7), -0.256) ≈ logpdf(DDM(-v, 1.2, 0.5, 0.25), 2, 0.7) + + #= stimulus-coding mirror symmetry: with no side bias (z = 0.5), + hitting the upper boundary under s = +1 is as likely as hitting + the lower boundary under s = -1 =# + du = StimulusCodedDDM(; ν=2.0, α=1.0, z=0.5, τ=0.2) + for rt in (0.3, 0.5, 1.0) + @test logdensityof(du, (1, rt), 1.0) ≈ logdensityof(du, (2, rt), -1.0) + end + + # zero-density observations return -Inf instead of throwing + @test logdensityof(d, (1, 0.1), 1.0) == -Inf # rt < τ + @test logdensityof(d, (1, 0.2), 1.0) == -Inf # rt == τ + @test logdensityof(d, (3, 0.6), 1.0) == -Inf # invalid choice + end + + @testset "_ddm_prob_upper closed form" begin + # matches the t → ∞ limit of the defective CDF computed by SSM + for (ν, α, z) in + ((1.5, 1.0, 0.4), (-2.0, 0.8, 0.6), (0.3, 2.0, 0.5), (-0.1, 1.2, 0.3)) + @test EM._ddm_prob_upper(ν, α, z) ≈ + SequentialSamplingModels.cdf(DDM(ν, α, z, 0.0), 1, 1e5) atol = 1e-8 + end + # zero drift: P(upper) is the relative start point + @test EM._ddm_prob_upper(0.0, 1.0, 0.3) == 0.3 + # extreme drifts saturate without overflow (a naive expm1 ratio NaNs) + @test EM._ddm_prob_upper(500.0, 2.0, 0.5) ≈ 1.0 + @test EM._ddm_prob_upper(-500.0, 2.0, 0.5) ≈ 0.0 atol = 1e-12 + @test EM._ddm_prob_upper(1.5f0, 1.0f0, 0.4f0) isa Float32 + end + + @testset "_ddm_cdf guards" begin + ν, α, z, τ = 1.5, 1.0, 0.4, 0.2 + # valid calls defer to SSM, whatever Real type carries the choice + ref = SequentialSamplingModels.cdf(DDM(ν, α, z, τ), 1, 0.6) + @test EM._ddm_cdf(ν, α, z, τ, 1, 0.6) == clamp(ref, 0.0, 1.0) + @test EM._ddm_cdf(ν, α, z, τ, 1.0, 0.6) == EM._ddm_cdf(ν, α, z, τ, 1, 0.6) + # cases SSM's cdf throws on carry zero mass instead + @test EM._ddm_cdf(ν, α, z, τ, 3, 0.6) == 0.0 # invalid choice + @test EM._ddm_cdf(ν, α, z, τ, 1.5, 0.6) == 0.0 # non-integral choice + @test EM._ddm_cdf(ν, α, z, τ, 1, 0.1) == 0.0 # rt < τ + @test EM._ddm_cdf(ν, α, z, τ, 1, τ) == 0.0 # rt == τ + # defective CDF: monotone in rt, saturating at the boundary-hit mass + Fs = [EM._ddm_cdf(ν, α, z, τ, 1, rt) for rt in (0.25, 0.4, 0.8, 2.0, 10.0)] + @test issorted(Fs) + @test all(0 .≤ Fs .≤ 1) + @test last(Fs) ≈ EM._ddm_prob_upper(ν, α, z) atol = 1e-6 + end + + @testset "type stability" begin + rng = MersenneTwister(0) + d = StimulusCodedDDM(; ν=1.5, α=1.0, z=0.4, τ=0.2) + c = CoherenceDDM(; k=8.0, γ=0.7, α=1.2, z=0.5, τ=0.25) + @test @inferred(logdensityof(d, (1, 0.6), 1.0)) isa Float64 + @test @inferred(logdensityof(c, (2, 0.7), -0.256)) isa Float64 + @test @inferred(EM._ddm_cdf(1.5, 1.0, 0.4, 0.2, 1, 0.6)) isa Float64 + @test @inferred(EM._ddm_prob_upper(1.5, 1.0, 0.4)) isa Float64 + @test @inferred(EM._emission_to_driver(rng, d, (1, 0.6), 1.0)) isa Vector{Float64} + end + + @testset "rand" begin + rng = MersenneTwister(0) + d = StimulusCodedDDM(; ν=2.0, α=1.0, z=0.5, τ=0.3) + for control in (1.0, -1.0, 1, -1) # Real and Integer control paths + obs = rand(rng, d, control) + @test obs.choice in (1, 2) + @test obs.rt > d.τ + @test isfinite(logdensityof(d, obs, control)) + end + + # a strong positive drift should mostly hit the upper boundary + n_upper = count(_ -> rand(rng, d, 1.0).choice == 1, 1:200) + @test n_upper > 150 + end + + @testset "fit! recovers StimulusCodedDDM parameters" begin + rng = MersenneTwister(1) + ν, α, z, τ = 2.0, 1.2, 0.55, 0.30 + truth = StimulusCodedDDM(; ν, α, z, τ) + n = 1500 + controls = [rand(rng, (-1.0, 1.0)) for _ in 1:n] + obs = [rand(rng, truth, controls[i]) for i in 1:n] + w = ones(n) + + d = StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.15) + fit!(d, obs, w; control_seq=controls) + + @test isapprox(d.ν, ν; rtol=0.15) + @test isapprox(d.α, α; rtol=0.15) + @test isapprox(d.z, z; atol=0.05) + @test isapprox(d.τ, τ; atol=0.05) + + # the fit must improve the weighted log-likelihood over the start + ll(m) = sum(logdensityof(m, obs[i], controls[i]) for i in 1:n) + @test ll(d) > ll(StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.15)) + end + + @testset "fit! recovers CoherenceDDM drift function" begin + rng = MersenneTwister(2) + truth = CoherenceDDM(; k=8.0, γ=0.7, α=1.2, z=0.5, τ=0.25) + levels = (0.032, 0.064, 0.128, 0.256, 0.512) + n = 2500 + controls = [rand(rng, (-1, 1)) * rand(rng, levels) for _ in 1:n] + obs = [rand(rng, truth, controls[i]) for i in 1:n] + w = ones(n) + + d = CoherenceDDM(; k=4.0, γ=1.0, α=1.0, z=0.5, τ=0.15) + fit!(d, obs, w; control_seq=controls) + + #= k and γ trade off against each other, so compare the drift + function they parameterize at the sampled coherence levels + rather than the raw parameters =# + for c in levels + v_true = 8.0 * c^0.7 + v_fit = d.k * c^d.γ + @test isapprox(v_fit, v_true; rtol=0.2) + end + @test isapprox(d.α, 1.2; rtol=0.15) + @test isapprox(d.τ, 0.25; atol=0.05) + end + + @testset "fit! weights and edge cases" begin + rng = MersenneTwister(3) + truth = StimulusCodedDDM(; ν=2.0, α=1.0, z=0.5, τ=0.25) + n = 400 + controls = [rand(rng, (-1.0, 1.0)) for _ in 1:n] + obs = [rand(rng, truth, controls[i]) for i in 1:n] + + # zero-weight observations must not affect the fit (even invalid ones) + obs_dirty = vcat(obs, [(3, -1.0)]) + controls_dirty = vcat(controls, [1.0]) + w_dirty = vcat(ones(n), [0.0]) + d_clean = StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.1) + d_dirty = StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.1) + fit!(d_clean, obs, ones(n); control_seq=controls) + fit!(d_dirty, obs_dirty, w_dirty; control_seq=controls_dirty) + @test (d_dirty.ν, d_dirty.α, d_dirty.z, d_dirty.τ) == + (d_clean.ν, d_clean.α, d_clean.z, d_clean.τ) + + # all-zero weights leave the parameters unchanged + d0 = StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.1) + fit!(d0, obs, zeros(n); control_seq=controls) + @test (d0.ν, d0.α, d0.z, d0.τ) == (1.0, 0.8, 0.5, 0.1) + + # positional ControlledEmission signature delegates to the keyword one + d_pos = StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.1) + d_kw = StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.1) + fit!(d_pos, obs, controls, ones(n)) + fit!(d_kw, obs, ones(n); control_seq=controls) + @test (d_pos.ν, d_pos.α, d_pos.z, d_pos.τ) == (d_kw.ν, d_kw.α, d_kw.z, d_kw.τ) + + # invalid weighted observations are rejected up front + @test_throws ArgumentError fit!( + StimulusCodedDDM(), [(3, 0.5)], [1.0]; control_seq=[1.0] + ) + @test_throws ArgumentError fit!( + StimulusCodedDDM(), [(1, -0.5)], [1.0]; control_seq=[1.0] + ) + @test_throws DimensionMismatch fit!( + StimulusCodedDDM(), obs, ones(n - 1); control_seq=controls + ) + @test_throws DimensionMismatch fit!( + StimulusCodedDDM(), obs, ones(n); control_seq=controls[1:(n - 1)] + ) + + # solver kwargs are accepted and a capped run stays well-defined + d_cap = StimulusCodedDDM(; ν=1.0, α=0.8, z=0.5, τ=0.1) + fit!(d_cap, obs, ones(n); control_seq=controls, max_iter=2, gtol=1e-3) + @test all(isfinite, (d_cap.ν, d_cap.α, d_cap.z, d_cap.τ)) + end + + @testset "Float32 parameters" begin + rng = MersenneTwister(5) + truth = StimulusCodedDDM(; ν=2.0f0, α=1.0f0, z=0.5f0, τ=0.2f0) + @test truth isa StimulusCodedDDM{Float32} + + n = 300 + controls = [rand(rng, (-1.0f0, 1.0f0)) for _ in 1:n] + obs = [rand(rng, truth, controls[i]) for i in 1:n] + + d = StimulusCodedDDM(; ν=1.0f0, α=0.8f0, z=0.5f0, τ=0.1f0) + ll(m) = sum(logdensityof(m, obs[i], controls[i]) for i in 1:n) + ll0 = ll(d) + #= SSM's sampler returns Float64 rts, so this exercises the mixed + Float32-emission / Float64-observation path, which must promote + rather than silently hit -Inf through SSM's series internals. =# + @test isfinite(ll0) + fit!(d, obs, ones(Float32, n); control_seq=controls, gtol=1e-4) + @test typeof((d.ν, d.α, d.z, d.τ)) == NTuple{4,Float32} + @test all(isfinite, (d.ν, d.α, d.z, d.τ)) + @test ll(d) > ll0 + end + + @testset "DDM-HMM: sample, forward, baum_welch" begin + rng = MersenneTwister(42) + T = 400 + init = [0.6, 0.4] + trans = [0.95 0.05; 0.1 0.9] + # an "engaged" state (strong drift, wide bounds) and a "lapse" state + dists = [ + StimulusCodedDDM(; ν=2.5, α=1.2, z=0.5, τ=0.25), + StimulusCodedDDM(; ν=0.3, α=0.7, z=0.5, τ=0.2), + ] + hmm = ControlledEmissionHMM(init, trans, dists) + + control_seq = [rand(rng, (-1.0, 1.0)) for _ in 1:T] + obs_seq = rand(rng, hmm, control_seq).obs_seq + @test length(obs_seq) == T + @test all(o -> o.choice in (1, 2) && o.rt > 0, obs_seq) + + logL = last(forward(hmm, obs_seq, control_seq; seq_ends=[T])) + @test all(isfinite, logL) + + # fit from a perturbed start; baum_welch must be (weakly) monotone + dists0 = [ + StimulusCodedDDM(; ν=1.5, α=1.0, z=0.5, τ=0.15), + StimulusCodedDDM(; ν=0.8, α=0.9, z=0.5, τ=0.15), + ] + hmm0 = ControlledEmissionHMM([0.5, 0.5], copy(trans), dists0) + _, lls = baum_welch(hmm0, obs_seq, control_seq; seq_ends=[T], max_iterations=8) + @test all(diff(lls) .>= -1e-6) + @test last(lls) >= first(lls) + end + + @testset "ACDC driver recovery" begin + #= Under the true model the Rosenblatt drivers (choice PIT, RT PIT) are + independent uniforms on (0,1): mean 1/2, variance 1/12. =# + rng = MersenneTwister(7) + for d in ( + StimulusCodedDDM(; ν=1.8, α=1.0, z=0.45, τ=0.2), + CoherenceDDM(; k=6.0, γ=0.8, α=1.1, z=0.55, τ=0.25), + ) + controls = d isa StimulusCodedDDM ? (-1.0, 1.0) : (-0.5, -0.2, 0.2, 0.5) + N = 40_000 + E = Matrix{Float64}(undef, 2, N) + for i in 1:N + c = rand(rng, controls) + obs = rand(rng, d, c) + E[:, i] = EM._emission_to_driver(rng, d, obs, c) + end + @test all(0 .< E .< 1) + for row in 1:2 + @test isapprox(mean(view(E, row, :)), 0.5; atol=0.02) + @test isapprox(var(view(E, row, :)), 1 / 12; atol=0.01) + end + end + end + + @testset "ACDC on a DDM-HMM (end-to-end)" begin + #= Exercises the full stochastic_drivers path on a ControlledEmissionHMM, + including the ControlBoundEmission unwrap. Drivers recovered under the + generating model should look uniform, i.e. low discrepancy. =# + rng = MersenneTwister(11) + T = 3000 + trans = [0.95 0.05; 0.08 0.92] + dists = [ + StimulusCodedDDM(; ν=2.5, α=1.2, z=0.5, τ=0.25), + StimulusCodedDDM(; ν=0.8, α=0.8, z=0.5, τ=0.2), + ] + hmm = ControlledEmissionHMM([0.6, 0.4], trans, dists) + control_seq = [rand(rng, (-1.0, 1.0)) for _ in 1:T] + obs_seq = rand(rng, hmm, control_seq).obs_seq + + sd = stochastic_drivers( + hmm, obs_seq; control_seq=control_seq, seq_ends=(T,), rng=MersenneTwister(3) + ) + @test length(sd.ε_pools) == 2 + @test all(p -> size(p, 1) == 2, sd.ε_pools) # (choice, rt) drivers + @test all(p -> all(0 .< p .< 1), sd.ε_pools) + @test isapprox(sum(sd.usage), 1; atol=1e-8) + + acdc = component_discrepancies( + hmm, + obs_seq, + KSDiscrepancy(); + control_seq=control_seq, + seq_ends=(T,), + rng=MersenneTwister(3), + ) + # well-specified emissions ⇒ near-uniform drivers ⇒ small KS discrepancy + @test all(isfinite, acdc.component_discrepancies) + @test maximum(acdc.component_discrepancies) < 0.1 + end + + @testset "extension fallback hooks" begin + #= The src fallbacks hat a user without SequentialSamplingModels + loaded hits for any argument types stay reachable =# + @test_throws "SequentialSamplingModels" EM._ddm_logpdf( + nothing, nothing, nothing, nothing, nothing, nothing + ) + @test_throws "SequentialSamplingModels" EM._ddm_rand( + nothing, nothing, nothing, nothing, nothing + ) + @test_throws "SequentialSamplingModels" EM._ddm_cdf( + nothing, nothing, nothing, nothing, nothing, nothing + ) + end +end