From 164bce2725868d8971b016d6116f8dc5d5ed628f Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Fri, 5 Jun 2026 22:04:07 -0400 Subject: [PATCH 1/6] Add ACDC model selection for HiddenMarkovModels.jl HMMs Port the Accumulated Cutoff Discrepancy Criterion (Li et al., 2026) from StateSpaceDynamics.jl PR #87, re-scoped to work with any AbstractHMM from HiddenMarkovModels.jl rather than SSD-internal model types. - Model-agnostic core (src/acdc/interface.jl): result types, the KL/KS/ Wasserstein/SquaredError/MMD discrepancies, and acdc_loss/acdc_select/ get_critical_rho_values. - Driver recovery (src/acdc/drivers.jl) runs off the standard Distributions PIT: cdf for continuous, randomized PIT for discrete, Cholesky-Rosenblatt for MvNormal. This subsumes the original Gaussian/Poisson/Bernoulli/ regression adapters and covers any standard emission; custom emissions add one _emission_to_driver method. - HiddenMarkovModels is a weak dependency: the stochastic_drivers(::AbstractHMM) method lives in a package extension, keeping EmissionModels usable standalone. - New hard deps: NearestNeighbors (KL kNN), Statistics. GMM/PMM/PPCA adapters and WrappedCauchy support were intentionally dropped (no counterpart types in HiddenMarkovModels.jl / EmissionModels.jl). Co-Authored-By: nguyenston --- Project.toml | 11 + ext/EmissionModelsHiddenMarkovModelsExt.jl | 88 ++++ src/EmissionModels.jl | 15 + src/acdc/drivers.jl | 82 ++++ src/acdc/interface.jl | 472 +++++++++++++++++++++ test/acdc/test_acdc.jl | 103 +++++ test/runtests.jl | 4 + 7 files changed, 775 insertions(+) create mode 100644 ext/EmissionModelsHiddenMarkovModelsExt.jl create mode 100644 src/acdc/drivers.jl create mode 100644 src/acdc/interface.jl create mode 100644 test/acdc/test_acdc.jl diff --git a/Project.toml b/Project.toml index b372185..e2ec111 100644 --- a/Project.toml +++ b/Project.toml @@ -8,18 +8,29 @@ DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688" +NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" +[weakdeps] +HiddenMarkovModels = "84ca31d5-effc-45e0-bfda-5a68cd981f47" + +[extensions] +EmissionModelsHiddenMarkovModelsExt = "HiddenMarkovModels" + [compat] DensityInterface = "0.4.0" Distributions = "0.25.122" +HiddenMarkovModels = "0.7" LinearAlgebra = "1.12.0" LogExpFunctions = "0.3.29" +NearestNeighbors = "0.4.27" Optim = "1.13.3, 2" Random = "1.11.0" SpecialFunctions = "2.6.1" +Statistics = "1.11.1" StatsAPI = "1.8.0" julia = "1.9" diff --git a/ext/EmissionModelsHiddenMarkovModelsExt.jl b/ext/EmissionModelsHiddenMarkovModelsExt.jl new file mode 100644 index 0000000..ce83903 --- /dev/null +++ b/ext/EmissionModelsHiddenMarkovModelsExt.jl @@ -0,0 +1,88 @@ +""" +EmissionModelsHiddenMarkovModelsExt + +Loads when `HiddenMarkovModels` is available alongside `EmissionModels`, adding +the ACDC `stochastic_drivers` method for any `AbstractHMM`. Driver recovery uses +forward-backward state posteriors plus the per-emission PIT defined in +`EmissionModels._emission_to_driver`, so it works with any HMM whose emissions +are standard `Distributions` (or a type with a custom driver method). +""" +module EmissionModelsHiddenMarkovModelsExt + +using EmissionModels: EmissionModels, StochasticDriverResult +using HiddenMarkovModels: AbstractHMM, obs_distributions, forward_backward +using Random: rand + +""" + stochastic_drivers(hmm::AbstractHMM, obs_seq; control_seq, seq_ends, n_samples=1) + +Recover the ACDC stochastic drivers for a fitted HiddenMarkovModels.jl `hmm`. + +For each time step the hidden state is sampled from its forward-backward +posterior, and that state's emission is inverted to a driver via the probability +integral transform (see [`EmissionModels._emission_to_driver`](@ref)). Component +"usage" is the posterior expected time spent in each state. + +# Arguments +- `hmm::AbstractHMM`: a fitted HMM satisfying the HiddenMarkovModels.jl interface. +- `obs_seq::AbstractVector`: observation sequence (scalars or vectors), as passed + to `forward_backward`. + +# Keyword Arguments +- `control_seq::AbstractVector`: controls for control-dependent HMMs; defaults to + `fill(nothing, length(obs_seq))`. +- `seq_ends`: end indices when `obs_seq` concatenates multiple sequences; defaults + to `(length(obs_seq),)`. +- `n_samples::Int=1`: number of posterior sampling passes over the data. + +# Returns +- [`StochasticDriverResult`](@ref) with per-state driver pools and usage. +""" +function EmissionModels.stochastic_drivers( + hmm::AbstractHMM, + obs_seq::AbstractVector; + control_seq::AbstractVector=fill(nothing, length(obs_seq)), + seq_ends=(length(obs_seq),), + n_samples::Int=1, +) + n_samples > 0 || throw(ArgumentError("n_samples must be positive")) + T_len = length(obs_seq) + K = length(hmm) + + # State posteriors γ (K × T) from forward-backward. + γ, _ = forward_backward(hmm, obs_seq, control_seq; seq_ends=seq_ends) + + # Usage: expected fraction of time in each state. + usage = vec(sum(γ; dims=2)) ./ T_len + + # Driver dimension from a probe on the first observation's sampled-state-able + # emission (all states share the emission dimension in an HMM). + probe = EmissionModels._emission_to_driver( + obs_distributions(hmm, control_seq[1])[1], obs_seq[1] + ) + D = length(probe) + Tε = eltype(probe) + + ε_lists = [Vector{Vector{Tε}}() for _ in 1:K] + for _ in 1:n_samples + for t in 1:T_len + z = EmissionModels._sample_categorical(view(γ, :, t)) + dist = obs_distributions(hmm, control_seq[t])[z] + push!(ε_lists[z], EmissionModels._emission_to_driver(dist, obs_seq[t])) + end + end + + # Pack each state's drivers into a D × n_k matrix. + ε_pools = Vector{Matrix{Tε}}(undef, K) + for k in 1:K + if isempty(ε_lists[k]) + ε_pools[k] = Matrix{Tε}(undef, D, 0) + else + ε_pools[k] = reduce(hcat, ε_lists[k]) + end + end + + return StochasticDriverResult(ε_pools, collect(Tε, usage)) +end + +end # module diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index e059dcc..1328dbf 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -1,19 +1,26 @@ module EmissionModels using Distributions: Normal, Bernoulli, Poisson, Chisq +using Distributions: cdf, quantile +using Distributions: + ContinuousUnivariateDistribution, DiscreteUnivariateDistribution, AbstractMvNormal using DensityInterface using LinearAlgebra using LogExpFunctions: logaddexp, logsumexp, log1pexp, logistic +using NearestNeighbors: KDTree, knn using Optim: optimize, TwiceDifferentiable, Newton, LBFGS, LineSearches using Optim using Random using SpecialFunctions: logfactorial, loggamma, digamma, trigamma, polygamma +using Statistics: mean, var, cov using StatsAPI using StatsAPI: fit! include("zeroinflated/poisson.jl") include("multivariate/t.jl") include("glms/glm.jl") +include("acdc/interface.jl") +include("acdc/drivers.jl") # exports export rand, logdensityof, fit! @@ -24,4 +31,12 @@ export MvGaussianGLM, MvBernoulliGLM, MvPoissonGLM export AbstractPrior, NoPrior, RidgePrior export neglogprior, neglogprior_grad!, neglogprior_hess! +# ACDC model selection +export ACDCResult, StochasticDriverResult, ComponentDiscrepancy +export stochastic_drivers, component_discrepancies +export acdc_loss, acdc_select, get_critical_rho_values +export compute_discrepancy +export KLDiscrepancy, KSDiscrepancy, WassersteinDiscrepancy +export SquaredErrorDiscrepancy, MMDDiscrepancy + end diff --git a/src/acdc/drivers.jl b/src/acdc/drivers.jl new file mode 100644 index 0000000..934ce9d --- /dev/null +++ b/src/acdc/drivers.jl @@ -0,0 +1,82 @@ +""" +Stochastic-driver recovery for individual emission distributions. + +Given an emission distribution and an observation, `_emission_to_driver` returns +the driver vector ``\\varepsilon \\in [0,1]^D`` via the probability integral +transform (PIT): + +- continuous univariate: ``\\varepsilon = F(x)``; +- discrete univariate: randomized PIT, ``\\varepsilon \\sim U(F(x^-), F(x))``, which + is exactly uniform under the true model; +- `MvNormal`: the Rosenblatt transform, which for a Gaussian reduces to whitening + the residual with the Cholesky factor and pushing each coordinate through + ``\\Phi``. + +These methods dispatch on `Distributions` types, so ACDC works with any +HiddenMarkovModels.jl HMM whose emissions are standard distributions. To support +a custom emission type, add a `_emission_to_driver(dist, obs)` method returning a +`Vector` of drivers in ``[0,1]``. +""" + +# Clamp a PIT value strictly inside (0,1); the discrepancy measures map drivers +# through the probit transform, which is ±Inf at the boundary. +_clamp01(u::T) where {T<:Real} = clamp(u, eps(T), one(T) - eps(T)) + +""" + _sample_categorical(p::AbstractVector) -> Int + +Draw a category index from probability vector `p` (assumed to sum to ≈1). +""" +function _sample_categorical(p::AbstractVector{T}) where {T<:Real} + u = rand(T) + cumsum_p = zero(T) + for i in eachindex(p) + cumsum_p += p[i] + u <= cumsum_p && return i + end + return lastindex(p) +end + +""" + _emission_to_driver(dist, obs) -> Vector{Float64} + +Recover the stochastic drivers for a single observation under emission `dist`. +Returns a length-`D` vector of values in ``(0,1)``. +""" +function _emission_to_driver end + +# Continuous univariate: standard PIT through the CDF. +function _emission_to_driver(d::ContinuousUnivariateDistribution, obs::Real) + return [_clamp01(float(cdf(d, obs)))] +end + +# Discrete univariate: randomized PIT, ε ~ U(F(x⁻), F(x)). +function _emission_to_driver(d::DiscreteUnivariateDistribution, obs::Real) + upper = float(cdf(d, obs)) + lower = float(cdf(d, obs - 1)) + return [_clamp01(lower + rand() * (upper - lower))] +end + +# Multivariate normal: Rosenblatt transform = whiten residual, then push each +# whitened coordinate through Φ. Whitening with the lower-Cholesky factor yields +# independent N(0,1) coordinates, so the result is uniform under the true model. +function _emission_to_driver(d::AbstractMvNormal, obs::AbstractVector) + μ = mean(d) + L = cholesky(Symmetric(Matrix(cov(d)))).L + z = L \ (collect(float.(obs)) .- μ) + return [_clamp01(float(cdf(Normal(), zi))) for zi in z] +end + +# Clear error for emission types without a PIT recipe, pointing the user at the +# extension hook. +function _emission_to_driver(d, obs) + throw( + ArgumentError( + "ACDC has no `_emission_to_driver` method for emission of type " * + "$(typeof(d)). Supported out of the box: continuous/discrete " * + "univariate `Distributions` and `MvNormal`. Define " * + "`EmissionModels._emission_to_driver(dist, obs)` returning a vector " * + "of drivers in [0,1] to add support.", + ), + ) +end diff --git a/src/acdc/interface.jl b/src/acdc/interface.jl new file mode 100644 index 0000000..79e0b26 --- /dev/null +++ b/src/acdc/interface.jl @@ -0,0 +1,472 @@ +#= +ACDC — Accumulated Cutoff Discrepancy Criterion. + +Robust model selection that measures component-level discrepancy via the +"stochastic drivers" framework (Li et al., 2026). The generative process is + +```math +\\begin{aligned} + x_n &= \\sum_k y_{n,k} \\\\ + y_{n,k} &= f(z_{n,k}, \\phi_k, \\varepsilon_{n,k}) \\quad \\text{where } \\varepsilon_{n,k} \\sim U(0,1) +\\end{aligned} +``` + +If the model is correctly specified, the recovered drivers ``\\varepsilon_{n,k}`` +are uniform on ``[0,1]``. ACDC scores each component by how far its drivers +deviate from uniformity, then selects the smallest component count whose +per-component discrepancies all fall below a cutoff ``\\rho``. + +This file holds the model-agnostic core: result types, discrepancy measures, +and the loss/selection machinery. Model-specific recovery of the drivers is +done by `stochastic_drivers`, whose methods live alongside the model types they +support (for HiddenMarkovModels.jl HMMs, in the package extension). + +This is a diagnostic / model-selection tool run once on a fitted model, not a +hot path, so the code favors clarity over the zero-allocation discipline the +emission `fit!`/`logdensityof` paths follow. +=# + +""" + ComponentDiscrepancy + +Abstract supertype for discrepancy measures between a sample of stochastic +drivers and the reference ``U([0,1]^D)`` distribution. +""" +abstract type ComponentDiscrepancy end + +""" + StochasticDriverResult{T<:Real} + +Stochastic drivers recovered from a fitted model. + +# Fields +- `ε_pools::Vector{Matrix{T}}`: per-component driver pools, each a `D × n_k` + matrix where `n_k` is the number of samples assigned to component `k`. +- `usage::Vector{T}`: per-component contribution magnitudes (length `K`). +""" +struct StochasticDriverResult{T<:Real} + ε_pools::Vector{Matrix{T}} + usage::Vector{T} + + function StochasticDriverResult( + ε_pools::Vector{Matrix{T}}, usage::Vector{T} + ) where {T<:Real} + K = length(ε_pools) + length(usage) == K || throw(ArgumentError("usage must have length K=$(K)")) + return new{T}(ε_pools, usage) + end +end + +""" + ACDCResult{T<:Real} + +Per-component ACDC discrepancies for a model with `K` components. + +# Fields +- `K::Int`: number of components. +- `component_discrepancies::Vector{T}`: per-component discrepancy values ``\\hat{D}_k``. +- `component_usage::Vector{T}`: per-component contribution magnitudes. +""" +struct ACDCResult{T<:Real} + K::Int + component_discrepancies::Vector{T} + component_usage::Vector{T} + + function ACDCResult(K::Int, discs::Vector{T}, usage::Vector{T}) where {T<:Real} + length(discs) == K || throw(ArgumentError("must have K discrepancy values")) + length(usage) == K || throw(ArgumentError("must have K usage values")) + return new{T}(K, discs, usage) + end +end + +""" + stochastic_drivers(model, data; n_samples=1, kwargs...) -> StochasticDriverResult + +Recover the stochastic drivers ``\\varepsilon_{n,k}`` for a fitted `model` by +inverting its generative process. Returns a [`StochasticDriverResult`](@ref). + +This is a generic function; methods are defined per model type. The method for +HiddenMarkovModels.jl `AbstractHMM`s is provided by the package extension that +loads with `HiddenMarkovModels`. The fallback below errors for unsupported models. +""" +function stochastic_drivers(model, data; kwargs...) + throw( + ArgumentError( + "ACDC has no `stochastic_drivers` method for a model of type " * + "$(typeof(model)). Load `HiddenMarkovModels` to enable support for " * + "`AbstractHMM` models, or define a `stochastic_drivers` method.", + ), + ) +end + +""" + component_discrepancies(model, data, discrepancy; n_samples=1, kwargs...) -> ACDCResult + +Compute per-component discrepancies from ``U(0,1)`` for a fitted `model`. + +Recovers the stochastic drivers via [`stochastic_drivers`](@ref) and scores each +component's driver pool with `discrepancy`. Extra keyword arguments are forwarded +to `stochastic_drivers`. +""" +function component_discrepancies( + model, data, discrepancy::ComponentDiscrepancy; n_samples::Int=1, kwargs... +) + result = stochastic_drivers(model, data; n_samples=n_samples, kwargs...) + usage = result.usage + K = length(usage) + T = eltype(usage) + + discs = Vector{T}(undef, K) + for k in 1:K + discs[k] = T(compute_discrepancy(discrepancy, result.ε_pools[k])) + end + + return ACDCResult(K, discs, usage) +end + +""" + KLDiscrepancy{T<:Real} <: ComponentDiscrepancy + +KL divergence ``D_{\\text{KL}}(P \\| U([0,1]^D))`` estimated via k-nearest-neighbor +density estimation. Returns 0 for perfectly uniform drivers, positive otherwise. + +To avoid boundary bias on the bounded hypercube, samples are mapped to +``\\mathbb{R}^D`` with the probit transform and compared against ``\\mathcal{N}(0, I)``. + +# Fields +- `k_neighbors::Int`: number of neighbors for kNN density estimation (default 5). +""" +struct KLDiscrepancy{T<:Real} <: ComponentDiscrepancy + k_neighbors::Int + + function KLDiscrepancy{T}(; k_neighbors::Int=5) where {T<:Real} + k_neighbors > 0 || throw(ArgumentError("k_neighbors must be positive")) + return new{T}(k_neighbors) + end +end + +KLDiscrepancy(; kwargs...) = KLDiscrepancy{Float64}(; kwargs...) + +""" + KSDiscrepancy{T<:Real} <: ComponentDiscrepancy + +Kolmogorov-Smirnov statistic for uniformity on ``[0,1]^D``. For `D=1` it is the +standard KS statistic against ``F(x)=x``; for `D>1` it is the maximum KS +statistic across marginals. + +!!! note + The multivariate version tests marginal uniformity but not independence; use + [`KLDiscrepancy`](@ref) or [`MMDDiscrepancy`](@ref) for a joint test. +""" +struct KSDiscrepancy{T<:Real} <: ComponentDiscrepancy + KSDiscrepancy{T}() where {T<:Real} = new{T}() +end + +KSDiscrepancy() = KSDiscrepancy{Float64}() + +""" + WassersteinDiscrepancy{T<:Real} <: ComponentDiscrepancy + +Wasserstein-`p` distance between the empirical distribution and ``U([0,1]^D)``. +Closed-form (sorted samples vs uniform quantiles) for `D=1`; sliced Wasserstein +(average over random 1D projections) for `D>1`. + +# Fields +- `p::Int`: order of the Wasserstein distance (default 2). +- `regularization::T`: unused; kept for API compatibility (default 0.1). +""" +struct WassersteinDiscrepancy{T<:Real} <: ComponentDiscrepancy + p::Int + regularization::T + + function WassersteinDiscrepancy{T}(; p::Int=2, regularization::T=T(0.1)) where {T<:Real} + p > 0 || throw(ArgumentError("p must be positive")) + return new{T}(p, regularization) + end +end + +WassersteinDiscrepancy(; kwargs...) = WassersteinDiscrepancy{Float64}(; kwargs...) + +""" + SquaredErrorDiscrepancy{T<:Real} <: ComponentDiscrepancy + +Moment-based discrepancy from ``U([0,1]^D)``: squared deviation of marginal means +from `0.5`, marginal variances from `1/12`, and cross-covariances from `0`. Fast +but less sensitive than KL or Wasserstein. +""" +struct SquaredErrorDiscrepancy{T<:Real} <: ComponentDiscrepancy + SquaredErrorDiscrepancy{T}() where {T<:Real} = new{T}() +end + +SquaredErrorDiscrepancy() = SquaredErrorDiscrepancy{Float64}() + +""" + MMDDiscrepancy{T<:Real} <: ComponentDiscrepancy + +Unbiased Maximum Mean Discrepancy with a Gaussian (RBF) kernel between the +empirical distribution and ``U([0,1]^D)``. Detects cross-dimensional dependence. +For `N > block_size`, computes the average MMD over blocks to stay ``O(N)``. + +# Fields +- `sigma::T`: kernel bandwidth (default 0.5). +- `block_size::Int`: max samples per block (default 5000). +""" +struct MMDDiscrepancy{T<:Real} <: ComponentDiscrepancy + sigma::T + block_size::Int + + function MMDDiscrepancy{T}(; sigma::T=T(0.5), block_size::Int=5000) where {T<:Real} + sigma > 0 || throw(ArgumentError("sigma must be positive")) + block_size > 1 || throw(ArgumentError("block_size must be > 1")) + return new{T}(sigma, block_size) + end +end + +MMDDiscrepancy(; kwargs...) = MMDDiscrepancy{Float64}(; kwargs...) + +""" + compute_discrepancy(d::ComponentDiscrepancy, samples::AbstractMatrix) -> Real + +Divergence between the empirical distribution of `samples` (a `D × N` matrix of +drivers in ``[0,1]``) and ``U([0,1]^D)``. Non-negative. +""" +function compute_discrepancy end + +function compute_discrepancy( + d::KLDiscrepancy{T}, samples::AbstractMatrix{T} +) where {T<:Real} + D, N = size(samples) + k = d.k_neighbors + + if N < k + 1 + @warn "Too few samples ($N) for k-NN KL estimation (k=$k)" + return T(Inf) + end + + # Probit transform to N(0,1) space; clamp to avoid ±Inf at the boundaries. + ϵ = eps(T) + samples_clamped = clamp.(samples, ϵ, one(T) - ϵ) + samples_normal = quantile.(Normal(zero(T), one(T)), samples_clamped) + + # E[log p_data] via k-NN in R^D. + mean_log_p = _mean_log_pdf_knn_multivariate(samples_normal, k) + + # E[log q_ref] for q_ref = N(0, I): log N(x;0,I) = -D/2 log(2π) - ½‖x‖². + log_2pi = log(T(2) * T(π)) + mean_sq_norm = mean(sum(samples_normal .^ 2; dims=1)) + mean_log_q = -T(0.5) * (D * log_2pi + mean_sq_norm) + + return max(zero(T), mean_log_p - mean_log_q) +end + +function compute_discrepancy( + d::KSDiscrepancy{T}, samples::AbstractMatrix{T} +) where {T<:Real} + D, N = size(samples) + + max_ks = zero(T) + for dim in 1:D + x = sort(view(samples, dim, :)) + ecdf_vals = collect(T, 1:N) ./ N + ref_cdf_vals = clamp.(x, zero(T), one(T)) + ks_stat = maximum( + max.( + abs.(ecdf_vals .- ref_cdf_vals), + abs.((ecdf_vals .- one(T) / N) .- ref_cdf_vals), + ), + ) + max_ks = max(max_ks, ks_stat) + end + return max_ks +end + +function compute_discrepancy( + d::SquaredErrorDiscrepancy{T}, samples::AbstractMatrix{T} +) where {T<:Real} + D, N = size(samples) + total_err = zero(T) + + # Marginal moments: mean → 0.5, variance → 1/12. + for dim in 1:D + x = view(samples, dim, :) + total_err += (mean(x) - T(0.5))^2 + total_err += (var(x) - T(1 / 12))^2 + end + + # Cross-covariances should be 0. + for i in 1:D + for j in (i + 1):D + total_err += cov(view(samples, i, :), view(samples, j, :))^2 + end + end + + return total_err / D +end + +function compute_discrepancy( + d::WassersteinDiscrepancy{T}, samples::AbstractMatrix{T} +) where {T<:Real} + D, N = size(samples) + + if D == 1 + x_sorted = sort(vec(samples)) + quantiles = [(i - T(0.5)) / N for i in 1:N] + w_dist = zero(T) + for i in 1:N + w_dist += abs(x_sorted[i] - quantiles[i])^d.p + end + return (w_dist / N)^(one(T) / d.p) + else + # Sliced Wasserstein: average over random 1D projections. + n_projections = 50 + total_w = zero(T) + for _ in 1:n_projections + direction = randn(T, D) + direction ./= norm(direction) + + x_sorted = sort(vec(direction' * samples)) + ref_samples = sort(vec(direction' * rand(T, D, N))) + + w_dist = zero(T) + for i in 1:N + w_dist += abs(x_sorted[i] - ref_samples[i])^d.p + end + total_w += (w_dist / N)^(one(T) / d.p) + end + return total_w / n_projections + end +end + +function compute_discrepancy( + d::MMDDiscrepancy{T}, samples::AbstractMatrix{T} +) where {T<:Real} + D, N = size(samples) + + if N <= d.block_size + reference_uniform = rand(T, D, N) + return _compute_mmd_quadratic_unbiased(samples, reference_uniform, d.sigma) + end + + # Block strategy for large N: average MMD over disjoint blocks. + n_blocks = div(N, d.block_size) + total_mmd = zero(T) + for b in 1:n_blocks + start_idx = (b - 1) * d.block_size + 1 + end_idx = b * d.block_size + block_samples = view(samples, :, start_idx:end_idx) + reference = rand(T, D, d.block_size) + total_mmd += _compute_mmd_quadratic_unbiased(block_samples, reference, d.sigma) + end + return max(zero(T), total_mmd / n_blocks) +end + +""" + acdc_loss(result::ACDCResult, ρ::Real) -> Real + +ACDC loss for cutoff ``\\rho``: ``R^\\rho(K) = \\sum_k \\max(0, \\hat{D}_k - \\rho)``. +""" +function acdc_loss(result::ACDCResult{T}, ρ::Real) where {T<:Real} + return sum(max(zero(T), d - T(ρ)) for d in result.component_discrepancies) +end + +""" + acdc_select(results::Vector{ACDCResult}, ρ::Real) -> Int + +Select the component count `K` with minimum ACDC loss at cutoff ``\\rho``, +breaking ties toward smaller `K`. `results` must be ordered by `K`. +""" +function acdc_select(results::Vector{ACDCResult{T}}, ρ::Real) where {T<:Real} + losses = [acdc_loss(r, ρ) for r in results] + min_loss = minimum(losses) + for (i, loss) in enumerate(losses) + if loss ≈ min_loss + return results[i].K + end + end + return results[argmin(losses)].K +end + +""" + get_critical_rho_values(results::Vector{ACDCResult}) -> Vector + +Sorted unique ``\\rho`` values at which the ACDC loss changes slope — exactly the +component discrepancy values across all `K`. +""" +function get_critical_rho_values(results::Vector{ACDCResult{T}}) where {T<:Real} + all_discs = T[] + for r in results + append!(all_discs, r.component_discrepancies) + end + return unique(sort(all_discs)) +end + +""" + _mean_log_pdf_knn_multivariate(samples::AbstractMatrix, k::Int) -> Real + +Estimate ``\\mathbb{E}_P[\\log p(x)]`` via the Kozachenko-Leonenko k-NN estimator +with KDTree acceleration. +""" +function _mean_log_pdf_knn_multivariate(samples::AbstractMatrix{T}, k::Int) where {T<:Real} + D, N = size(samples) + + tree = KDTree(samples) + # k+1 neighbors because the nearest is the point itself. + _, dists = knn(tree, samples, k + 1, true) + + log_c_D = (D / 2) * log(T(π)) - loggamma(T(D) / 2 + 1) + bias_correction = digamma(T(k)) - digamma(T(N)) + + sum_log_p = zero(T) + for i in 1:N + rho_k = max(dists[i][k + 1], eps(T)) + sum_log_p += bias_correction - log_c_D - D * log(rho_k) + end + return sum_log_p / N +end + +""" + _compute_mmd_quadratic_unbiased(X, Y, sigma) -> Real + +Unbiased U-statistic ``\\text{MMD}^2`` between `D × N` matrices `X` and `Y` with a +Gaussian RBF kernel ``k(x,y) = \\exp(-\\|x-y\\|^2 / (2\\sigma^2))``. ``O(N^2)``. +""" +function _compute_mmd_quadratic_unbiased( + X::AbstractMatrix{T}, Y::AbstractMatrix{T}, sigma::T +) where {T<:Real} + D, N = size(X) + gamma = one(T) / (T(2) * sigma^2) + + sum_xx = zero(T) + sum_yy = zero(T) + sum_xy = zero(T) + + for i in 1:N + # Intra-group sums, excluding the diagonal (unbiased estimator). + for j in (i + 1):N + dist_sq_xx = zero(T) + dist_sq_yy = zero(T) + for d in 1:D + diff_x = X[d, i] - X[d, j] + diff_y = Y[d, i] - Y[d, j] + dist_sq_xx += diff_x^2 + dist_sq_yy += diff_y^2 + end + sum_xx += exp(-gamma * dist_sq_xx) + sum_yy += exp(-gamma * dist_sq_yy) + end + # Cross-group sum over all pairs. + for j in 1:N + dist_sq_xy = zero(T) + for d in 1:D + diff_xy = X[d, i] - Y[d, j] + dist_sq_xy += diff_xy^2 + end + sum_xy += exp(-gamma * dist_sq_xy) + end + end + + # 2/(N(N-1)) restores the symmetric lower triangle skipped above. + norm_intra = T(2) / (N * (N - 1)) + norm_inter = T(2) / (N * N) + return (norm_intra * sum_xx) + (norm_intra * sum_yy) - (norm_inter * sum_xy) +end diff --git a/test/acdc/test_acdc.jl b/test/acdc/test_acdc.jl new file mode 100644 index 0000000..0b4f531 --- /dev/null +++ b/test/acdc/test_acdc.jl @@ -0,0 +1,103 @@ +using EmissionModels +using HiddenMarkovModels +using Distributions +using LinearAlgebra +using Random +using Test + +@testset "Discrepancy measures on (non)uniform samples" begin + rng = Random.MersenneTwister(0) + + # Uniform drivers ⇒ every discrepancy is small. + U = rand(rng, 1, 4000) + @test compute_discrepancy(KSDiscrepancy(), U) < 0.05 + @test compute_discrepancy(SquaredErrorDiscrepancy(), U) < 1e-2 + @test compute_discrepancy(WassersteinDiscrepancy(), U) < 0.05 + @test abs(compute_discrepancy(KLDiscrepancy(), U)) < 0.1 + @test abs(compute_discrepancy(MMDDiscrepancy(; block_size=4000), U)) < 1e-2 + + # Strongly non-uniform drivers (concentrated near 0.5) ⇒ larger scores. + B = clamp.(0.5 .+ 0.05 .* randn(rng, 1, 4000), 1e-6, 1 - 1e-6) + @test compute_discrepancy(KSDiscrepancy(), B) > compute_discrepancy(KSDiscrepancy(), U) + @test compute_discrepancy(KLDiscrepancy(), B) > compute_discrepancy(KLDiscrepancy(), U) + @test compute_discrepancy(SquaredErrorDiscrepancy(), B) > + compute_discrepancy(SquaredErrorDiscrepancy(), U) + + # Multivariate uniform ⇒ small KS / MMD. + U2 = rand(rng, 2, 3000) + @test compute_discrepancy(KSDiscrepancy(), U2) < 0.05 + @test compute_discrepancy(MMDDiscrepancy(; block_size=3000), U2) < 1e-2 +end + +@testset "ACDC loss and selection" begin + # K=2 fits well (both components below cutoff), K=3 has one bad component. + r2 = ACDCResult(2, [0.01, 0.02], [0.5, 0.5]) + r3 = ACDCResult(3, [0.01, 0.02, 0.40], [0.4, 0.4, 0.2]) + + @test acdc_loss(r2, 0.1) == 0.0 + @test acdc_loss(r3, 0.1) ≈ 0.30 + @test acdc_select([r2, r3], 0.1) == 2 # prefer smaller K at min loss + @test acdc_select([r2, r3], 0.5) == 2 # both zero ⇒ smaller K + + crit = get_critical_rho_values([r2, r3]) + @test crit == sort(unique([0.01, 0.02, 0.01, 0.02, 0.40])) +end + +@testset "HMM stochastic drivers (Normal emissions)" begin + rng = Random.MersenneTwister(42) + init = [0.34, 0.33, 0.33] + trans = [0.9 0.05 0.05; 0.05 0.9 0.05; 0.05 0.05 0.9] + dists = [Normal(-6.0, 1.0), Normal(0.0, 1.0), Normal(6.0, 1.0)] + hmm = HMM(init, trans, dists) + _, obs_seq = rand(rng, hmm, 3000) + + sd = stochastic_drivers(hmm, obs_seq) + @test length(sd.ε_pools) == 3 + @test all(p -> size(p, 1) == 1, sd.ε_pools) # univariate ⇒ D=1 + @test sum(size.(sd.ε_pools, 2)) == 3000 # every step assigned once + @test isapprox(sum(sd.usage), 1.0; atol=1e-8) + @test all(p -> all(0 .<= p .<= 1), sd.ε_pools) # drivers in [0,1] + + # Well-specified model ⇒ drivers ≈ uniform ⇒ small per-state discrepancy. + res = component_discrepancies(hmm, obs_seq, KSDiscrepancy()) + @test res.K == 3 + @test all(<(0.1), res.component_discrepancies) +end + +@testset "HMM well-specified vs misspecified" begin + rng = Random.MersenneTwister(7) + init = [0.5, 0.5] + trans = [0.95 0.05; 0.05 0.95] + true_hmm = HMM(init, trans, [Normal(-4.0, 1.0), Normal(4.0, 1.0)]) + _, obs_seq = rand(rng, true_hmm, 3000) + + good = component_discrepancies(true_hmm, obs_seq, KLDiscrepancy()) + # Misspecified emission variances ⇒ drivers deviate from uniform. + bad_hmm = HMM(init, trans, [Normal(-4.0, 3.0), Normal(4.0, 3.0)]) + bad = component_discrepancies(bad_hmm, obs_seq, KLDiscrepancy()) + + @test maximum(bad.component_discrepancies) > maximum(good.component_discrepancies) +end + +@testset "HMM discrete (Poisson) and multivariate (MvNormal) emissions" begin + rng = Random.MersenneTwister(11) + init = [0.5, 0.5] + trans = [0.9 0.1; 0.1 0.9] + + pois = HMM(init, trans, [Poisson(2.0), Poisson(20.0)]) + _, pobs = rand(rng, pois, 2500) + pres = component_discrepancies(pois, pobs, KSDiscrepancy()) + @test pres.K == 2 + @test all(<(0.1), pres.component_discrepancies) + + mv = HMM( + init, + trans, + [MvNormal([-3.0, -3.0], I(2)), MvNormal([3.0, 3.0], [1.0 0.5; 0.5 1.0])], + ) + _, mobs = rand(rng, mv, 2500) + sd = stochastic_drivers(mv, mobs) + @test all(p -> size(p, 1) == 2, sd.ε_pools) # bivariate ⇒ D=2 + mres = component_discrepancies(mv, mobs, KSDiscrepancy()) + @test all(<(0.1), mres.component_discrepancies) +end diff --git a/test/runtests.jl b/test/runtests.jl index 278fce6..3e2dffb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -26,6 +26,10 @@ using StatsAPI include("multivariate/test_t.jl") end + @testset "ACDC model selection" begin + include("acdc/test_acdc.jl") + end + @testset "Allocations" begin include("allocations.jl") end From b2551ec00d79c349a53f386eb92f3557d47c09f9 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Fri, 5 Jun 2026 22:38:31 -0400 Subject: [PATCH 2/6] Fix CI for ACDC: formatting, docs, and coverage - Format: relocate an inline block comment in glm.jl that JuliaFormatter v2 (the version CI pins) strips; repo now passes `format(".", overwrite=false)`. - Docs: add docs/src/acdc.md documenting the ACDC public API and register it in make.jl; convert ACDC internal-helper docstrings to comments (matching the project convention) so Documenter's checkdocs=:all passes. - Coverage: add tests for the KL too-few-samples path, multivariate SquaredError/Wasserstein, the MMD block-averaging path, the unsupported- emission error, and HMM adapter argument validation. --- docs/make.jl | 1 + docs/src/acdc.md | 83 ++++++++++++++++++++++++++++++++++++++++++ src/acdc/drivers.jl | 45 +++++++++-------------- src/acdc/interface.jl | 16 ++------ src/glms/glm.jl | 3 +- test/acdc/test_acdc.jl | 34 +++++++++++++++++ 6 files changed, 141 insertions(+), 41 deletions(-) create mode 100644 docs/src/acdc.md diff --git a/docs/make.jl b/docs/make.jl index f5bcbd8..240631e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -20,6 +20,7 @@ makedocs(; "Distributions" => "distributions.md", "GLM Emissions" => "glm.md", "Priors" => "priors.md", + "ACDC Model Selection" => "acdc.md", "Custom Emission Models" => "custom.md", ], ) diff --git a/docs/src/acdc.md b/docs/src/acdc.md new file mode 100644 index 0000000..e3dc6bd --- /dev/null +++ b/docs/src/acdc.md @@ -0,0 +1,83 @@ +# ACDC Model Selection + +The **Accumulated Cutoff Discrepancy Criterion** (ACDC, Li et al. 2026) is a +robust model-selection method for hidden Markov models. Instead of penalizing +likelihood by parameter count (AIC/BIC), ACDC scores each state by how far its +*stochastic drivers* deviate from uniformity. + +## Idea + +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 drivers per state, measures their discrepancy from +``U([0,1]^D)``, and selects the smallest number of states whose per-state +discrepancies all fall below a cutoff ``\rho``. + +This works with **any `AbstractHMM` from +[HiddenMarkovModels.jl](https://github.com/gdalle/HiddenMarkovModels.jl)** whose +emissions are standard `Distributions` (continuous via the CDF, discrete via a +randomized PIT, `MvNormal` via the Cholesky–Rosenblatt transform). The HMM +method loads automatically once `HiddenMarkovModels` is imported. + +## Example + +```julia +using EmissionModels, HiddenMarkovModels, Distributions + +hmm = HMM([0.5, 0.5], [0.95 0.05; 0.05 0.95], + [Normal(-4.0, 1.0), Normal(4.0, 1.0)]) +_, obs_seq = rand(hmm, 3000) + +# Per-state discrepancy from uniform; small ⇒ well-specified. +result = component_discrepancies(hmm, obs_seq, KSDiscrepancy()) + +# Pick the number of states with smallest ACDC loss at cutoff ρ. +K = acdc_select([result], 0.05) +``` + +To score a candidate set of fitted models, build one [`ACDCResult`](@ref) per +model (ordered by state count) and pass the vector to [`acdc_select`](@ref). +[`get_critical_rho_values`](@ref) returns the cutoffs at which the selection +changes. + +## Custom emissions + +To use ACDC with an emission type that is not a standard `Distributions` object, +add a method `EmissionModels._emission_to_driver(dist, obs)` returning the driver +vector in ``[0,1]``. + +## API Reference + +### Driver recovery and discrepancies + +```@docs +stochastic_drivers +component_discrepancies +compute_discrepancy +``` + +### Discrepancy measures + +```@docs +ComponentDiscrepancy +KLDiscrepancy +KSDiscrepancy +WassersteinDiscrepancy +SquaredErrorDiscrepancy +MMDDiscrepancy +``` + +### Loss and selection + +```@docs +acdc_loss +acdc_select +get_critical_rho_values +``` + +### Result types + +```@docs +StochasticDriverResult +ACDCResult +``` diff --git a/src/acdc/drivers.jl b/src/acdc/drivers.jl index 934ce9d..d77f949 100644 --- a/src/acdc/drivers.jl +++ b/src/acdc/drivers.jl @@ -1,32 +1,25 @@ -""" -Stochastic-driver recovery for individual emission distributions. +#= Stochastic-driver recovery for individual emission distributions. -Given an emission distribution and an observation, `_emission_to_driver` returns -the driver vector ``\\varepsilon \\in [0,1]^D`` via the probability integral -transform (PIT): + Given an emission distribution and an observation, `_emission_to_driver` + returns the driver vector ε ∈ [0,1]^D via the probability integral transform + (PIT): + - continuous univariate: ε = F(x); + - discrete univariate: randomized PIT, ε ~ U(F(x⁻), F(x)), exactly uniform + under the true model; + - MvNormal: the Rosenblatt transform, which for a Gaussian reduces to + whitening the residual with the Cholesky factor and pushing each coordinate + through Φ. -- continuous univariate: ``\\varepsilon = F(x)``; -- discrete univariate: randomized PIT, ``\\varepsilon \\sim U(F(x^-), F(x))``, which - is exactly uniform under the true model; -- `MvNormal`: the Rosenblatt transform, which for a Gaussian reduces to whitening - the residual with the Cholesky factor and pushing each coordinate through - ``\\Phi``. - -These methods dispatch on `Distributions` types, so ACDC works with any -HiddenMarkovModels.jl HMM whose emissions are standard distributions. To support -a custom emission type, add a `_emission_to_driver(dist, obs)` method returning a -`Vector` of drivers in ``[0,1]``. -""" + These methods dispatch on `Distributions` types, so ACDC works with any + HiddenMarkovModels.jl HMM whose emissions are standard distributions. To + support a custom emission type, add a `_emission_to_driver(dist, obs)` method + returning a `Vector` of drivers in [0,1]. =# # Clamp a PIT value strictly inside (0,1); the discrepancy measures map drivers # through the probit transform, which is ±Inf at the boundary. _clamp01(u::T) where {T<:Real} = clamp(u, eps(T), one(T) - eps(T)) -""" - _sample_categorical(p::AbstractVector) -> Int - -Draw a category index from probability vector `p` (assumed to sum to ≈1). -""" +# Draw a category index from probability vector `p` (assumed to sum to ≈1). function _sample_categorical(p::AbstractVector{T}) where {T<:Real} u = rand(T) cumsum_p = zero(T) @@ -37,12 +30,8 @@ function _sample_categorical(p::AbstractVector{T}) where {T<:Real} return lastindex(p) end -""" - _emission_to_driver(dist, obs) -> Vector{Float64} - -Recover the stochastic drivers for a single observation under emission `dist`. -Returns a length-`D` vector of values in ``(0,1)``. -""" +# Recover the stochastic drivers for a single observation under emission `dist`. +# Returns a length-`D` vector of values in (0,1). function _emission_to_driver end # Continuous univariate: standard PIT through the CDF. diff --git a/src/acdc/interface.jl b/src/acdc/interface.jl index 79e0b26..3087c17 100644 --- a/src/acdc/interface.jl +++ b/src/acdc/interface.jl @@ -400,12 +400,8 @@ function get_critical_rho_values(results::Vector{ACDCResult{T}}) where {T<:Real} return unique(sort(all_discs)) end -""" - _mean_log_pdf_knn_multivariate(samples::AbstractMatrix, k::Int) -> Real - -Estimate ``\\mathbb{E}_P[\\log p(x)]`` via the Kozachenko-Leonenko k-NN estimator -with KDTree acceleration. -""" +# Estimate E_P[log p(x)] via the Kozachenko-Leonenko k-NN estimator with KDTree +# acceleration. function _mean_log_pdf_knn_multivariate(samples::AbstractMatrix{T}, k::Int) where {T<:Real} D, N = size(samples) @@ -424,12 +420,8 @@ function _mean_log_pdf_knn_multivariate(samples::AbstractMatrix{T}, k::Int) wher return sum_log_p / N end -""" - _compute_mmd_quadratic_unbiased(X, Y, sigma) -> Real - -Unbiased U-statistic ``\\text{MMD}^2`` between `D × N` matrices `X` and `Y` with a -Gaussian RBF kernel ``k(x,y) = \\exp(-\\|x-y\\|^2 / (2\\sigma^2))``. ``O(N^2)``. -""" +# Unbiased U-statistic MMD² between D × N matrices X and Y with a Gaussian RBF +# kernel k(x,y) = exp(-‖x-y‖² / (2σ²)). O(N²). function _compute_mmd_quadratic_unbiased( X::AbstractMatrix{T}, Y::AbstractMatrix{T}, sigma::T ) where {T<:Real} diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 3a1d237..b3bcac7 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -828,8 +828,9 @@ function StatsAPI.fit!( XWX = zeros(T, p, p) XWY = zeros(T, p, k) - wsum = zero(T)#= Build XᵀWX and XᵀWY in a single pass — no Y matrix, no temporaries. =# + wsum = zero(T) + # Build XᵀWX and XᵀWY in a single pass — no Y matrix, no temporaries. for i in 1:n obs_i = obs_seq[i] length(obs_i) == k || diff --git a/test/acdc/test_acdc.jl b/test/acdc/test_acdc.jl index 0b4f531..f200ed8 100644 --- a/test/acdc/test_acdc.jl +++ b/test/acdc/test_acdc.jl @@ -101,3 +101,37 @@ end mres = component_discrepancies(mv, mobs, KSDiscrepancy()) @test all(<(0.1), mres.component_discrepancies) end + +@testset "Discrepancy edge cases" begin + rng = Random.MersenneTwister(123) + + # KL with too few samples (N < k+1) ⇒ Inf. + @test compute_discrepancy(KLDiscrepancy(; k_neighbors=5), rand(rng, 2, 3)) == Inf + + # Multivariate paths: SquaredError cross-covariance loop, sliced Wasserstein. + U3 = rand(rng, 3, 1500) + @test compute_discrepancy(SquaredErrorDiscrepancy(), U3) < 1e-2 + @test compute_discrepancy(WassersteinDiscrepancy(), U3) < 0.1 + + # MMD block-averaging path (N > block_size). + @test compute_discrepancy(MMDDiscrepancy(; block_size=200), rand(rng, 2, 600)) < 0.05 + + # Type stability through a parametric discrepancy. + @test compute_discrepancy(KSDiscrepancy{Float32}(), rand(rng, Float32, 1, 200)) isa + Float32 + + # Unsupported emission ⇒ clear error. + @test_throws ArgumentError EmissionModels._emission_to_driver(missing, 1.0) +end + +@testset "HMM adapter argument validation" begin + rng = Random.MersenneTwister(5) + hmm = HMM([0.5, 0.5], [0.9 0.1; 0.1 0.9], [Normal(0.0, 1.0), Normal(5.0, 1.0)]) + _, obs_seq = rand(rng, hmm, 100) + + @test_throws ArgumentError stochastic_drivers(hmm, obs_seq; n_samples=0) + + # n_samples > 1 multiplies the pool size (one assignment per step per pass). + sd = stochastic_drivers(hmm, obs_seq; n_samples=2) + @test sum(size.(sd.ε_pools, 2)) == 200 +end From 5fde721831ff06b2e2abe84d9f7099d685cd289f Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sat, 6 Jun 2026 09:14:08 -0400 Subject: [PATCH 3/6] Test stochastic_drivers fallback error for unsupported models --- test/acdc/test_acdc.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test/acdc/test_acdc.jl b/test/acdc/test_acdc.jl index f200ed8..5c69fc9 100644 --- a/test/acdc/test_acdc.jl +++ b/test/acdc/test_acdc.jl @@ -122,6 +122,11 @@ end # Unsupported emission ⇒ clear error. @test_throws ArgumentError EmissionModels._emission_to_driver(missing, 1.0) + + # Unsupported model ⇒ stochastic_drivers fallback errors (also via the + # component_discrepancies path, which forwards to stochastic_drivers). + @test_throws ArgumentError stochastic_drivers(missing, [1.0, 2.0]) + @test_throws ArgumentError component_discrepancies(missing, [1.0, 2.0], KSDiscrepancy()) end @testset "HMM adapter argument validation" begin From 3fcece3c8598bb671a5e90127ccb2ce52e2fe0ef Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sun, 7 Jun 2026 21:43:53 -0400 Subject: [PATCH 4/6] add more drivers --- ext/EmissionModelsHiddenMarkovModelsExt.jl | 7 +- src/EmissionModels.jl | 2 +- src/acdc/drivers.jl | 111 ++++++++++++++++++++- test/acdc/test_acdc.jl | 81 +++++++++++++++ 4 files changed, 193 insertions(+), 8 deletions(-) diff --git a/ext/EmissionModelsHiddenMarkovModelsExt.jl b/ext/EmissionModelsHiddenMarkovModelsExt.jl index ce83903..a085e27 100644 --- a/ext/EmissionModelsHiddenMarkovModelsExt.jl +++ b/ext/EmissionModelsHiddenMarkovModelsExt.jl @@ -58,7 +58,7 @@ function EmissionModels.stochastic_drivers( # Driver dimension from a probe on the first observation's sampled-state-able # emission (all states share the emission dimension in an HMM). probe = EmissionModels._emission_to_driver( - obs_distributions(hmm, control_seq[1])[1], obs_seq[1] + obs_distributions(hmm, control_seq[1])[1], obs_seq[1], control_seq[1] ) D = length(probe) Tε = eltype(probe) @@ -68,7 +68,10 @@ function EmissionModels.stochastic_drivers( for t in 1:T_len z = EmissionModels._sample_categorical(view(γ, :, t)) dist = obs_distributions(hmm, control_seq[t])[z] - push!(ε_lists[z], EmissionModels._emission_to_driver(dist, obs_seq[t])) + push!( + ε_lists[z], + EmissionModels._emission_to_driver(dist, obs_seq[t], control_seq[t]), + ) end end diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index 1328dbf..1cc13aa 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -1,6 +1,6 @@ module EmissionModels -using Distributions: Normal, Bernoulli, Poisson, Chisq +using Distributions: Normal, Bernoulli, Poisson, Chisq, TDist, MvNormal using Distributions: cdf, quantile using Distributions: ContinuousUnivariateDistribution, DiscreteUnivariateDistribution, AbstractMvNormal diff --git a/src/acdc/drivers.jl b/src/acdc/drivers.jl index d77f949..c94d5c4 100644 --- a/src/acdc/drivers.jl +++ b/src/acdc/drivers.jl @@ -9,11 +9,22 @@ - MvNormal: the Rosenblatt transform, which for a Gaussian reduces to whitening the residual with the Cholesky factor and pushing each coordinate through Φ. + - PoissonZeroInflated: randomized PIT against the ZIP CDF + F(k) = π + (1-π)·F_Poisson(k). + - MultivariateT / MultivariateTDiag: the conditional-t Rosenblatt — whiten + the residual to a standardized spherical t, then push each coordinate + through its conditional Student-t CDF. (A Gaussian-style whiten-then-Φ is + wrong here: whitened t-coordinates are uncorrelated but not independent.) + - GLMs (`AbstractGLM`): conditional emissions f(y | x). Recovered through a + 3-arg `_emission_to_driver(dist, obs, x)` that reduces the GLM at covariate + `x` to the standard emission it is (Normal / Bernoulli / Poisson / MvNormal) + and reuses the recipes above. - These methods dispatch on `Distributions` types, so ACDC works with any - HiddenMarkovModels.jl HMM whose emissions are standard distributions. To - support a custom emission type, add a `_emission_to_driver(dist, obs)` method - returning a `Vector` of drivers in [0,1]. =# + These methods dispatch on `Distributions` types and the package's own emission + types, so ACDC works with any HiddenMarkovModels.jl HMM whose emissions are + standard distributions or those types. To support a further custom emission + type, add a `_emission_to_driver(dist, obs)` method (or the 3-arg form for a + covariate-dependent emission) returning a `Vector` of drivers in [0,1]. =# # Clamp a PIT value strictly inside (0,1); the discrepancy measures map drivers # through the probit transform, which is ±Inf at the boundary. @@ -56,6 +67,95 @@ function _emission_to_driver(d::AbstractMvNormal, obs::AbstractVector) return [_clamp01(float(cdf(Normal(), zi))) for zi in z] end +# Zero-inflated Poisson: randomized PIT against the ZIP CDF. The mixture CDF is +# F(k) = π + (1-π)·F_Poisson(k) for k ≥ 0; the structural-zero mass collapses +# into F(0), so a true zero still maps uniformly into [0, F(0)]. +function _emission_to_driver(d::PoissonZeroInflated, obs::Real) + pois = Poisson(d.λ) + upper = d.π + (1 - d.π) * cdf(pois, obs) + lower = obs > 0 ? d.π + (1 - d.π) * cdf(pois, obs - 1) : zero(upper) + return [_clamp01(float(lower + rand() * (upper - lower)))] +end + +# Multivariate Student-t (full and diagonal scale): the conditional-t Rosenblatt +# transform. Whitening the residual yields a standardized spherical t whose +# coordinates are uncorrelated but share the common χ² scale, so they are NOT +# independent — pushing each through Φ (as for a Gaussian) would be wrong. The +# correct map sends each whitened coordinate through its conditional Student-t +# CDF, which has ν+(j-1) degrees of freedom and scale √((ν+Σ_{i Date: Sun, 28 Jun 2026 09:59:11 -0400 Subject: [PATCH 5/6] fix merge issues --- Project.toml | 7 +--- src/EmissionModels.jl | 2 ++ .../acdc/hmm.jl | 35 ++++++------------- src/acdc/interface.jl | 15 +++----- 4 files changed, 19 insertions(+), 40 deletions(-) rename ext/EmissionModelsHiddenMarkovModelsExt.jl => src/acdc/hmm.jl (67%) diff --git a/Project.toml b/Project.toml index c53550d..757eb06 100644 --- a/Project.toml +++ b/Project.toml @@ -16,18 +16,13 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" -[weakdeps] -HiddenMarkovModels = "84ca31d5-effc-45e0-bfda-5a68cd981f47" - -[extensions] -EmissionModelsHiddenMarkovModelsExt = "HiddenMarkovModels" - [compat] DensityInterface = "0.4.0" Distributions = "0.25.122" HiddenMarkovModels = "0.7.1" LinearAlgebra = "1.12.0" LogExpFunctions = "0.3.29, 1" +NearestNeighbors = "0.4.27" Optim = "1.13.3, 2" Random = "1.11.0" SpecialFunctions = "2.6.1" diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index b054073..922b0c8 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -6,6 +6,7 @@ using Distributions: ContinuousUnivariateDistribution, DiscreteUnivariateDistribution, AbstractMvNormal using DensityInterface using HiddenMarkovModels: ControlledEmission +using HiddenMarkovModels: AbstractHMM, obs_distributions, forward_backward using LinearAlgebra using LogExpFunctions: logaddexp, logsumexp, log1pexp, logistic using NearestNeighbors: KDTree, knn @@ -22,6 +23,7 @@ include("multivariate/t.jl") include("glms/glm.jl") include("acdc/interface.jl") include("acdc/drivers.jl") +include("acdc/hmm.jl") # exports export rand, logdensityof, fit! diff --git a/ext/EmissionModelsHiddenMarkovModelsExt.jl b/src/acdc/hmm.jl similarity index 67% rename from ext/EmissionModelsHiddenMarkovModelsExt.jl rename to src/acdc/hmm.jl index a085e27..752cd50 100644 --- a/ext/EmissionModelsHiddenMarkovModelsExt.jl +++ b/src/acdc/hmm.jl @@ -1,17 +1,9 @@ -""" -EmissionModelsHiddenMarkovModelsExt - -Loads when `HiddenMarkovModels` is available alongside `EmissionModels`, adding -the ACDC `stochastic_drivers` method for any `AbstractHMM`. Driver recovery uses -forward-backward state posteriors plus the per-emission PIT defined in -`EmissionModels._emission_to_driver`, so it works with any HMM whose emissions -are standard `Distributions` (or a type with a custom driver method). -""" -module EmissionModelsHiddenMarkovModelsExt +#= ACDC stochastic-driver recovery for HiddenMarkovModels.jl HMMs. -using EmissionModels: EmissionModels, StochasticDriverResult -using HiddenMarkovModels: AbstractHMM, obs_distributions, forward_backward -using Random: rand + Driver recovery uses the forward-backward state posteriors plus the + per-emission PIT defined in `_emission_to_driver` (see `drivers.jl`), so it + works with any HMM whose emissions are standard `Distributions` (or a type + with a custom driver method). =# """ stochastic_drivers(hmm::AbstractHMM, obs_seq; control_seq, seq_ends, n_samples=1) @@ -20,8 +12,8 @@ Recover the ACDC stochastic drivers for a fitted HiddenMarkovModels.jl `hmm`. For each time step the hidden state is sampled from its forward-backward posterior, and that state's emission is inverted to a driver via the probability -integral transform (see [`EmissionModels._emission_to_driver`](@ref)). Component -"usage" is the posterior expected time spent in each state. +integral transform (see [`_emission_to_driver`](@ref)). Component "usage" is the +posterior expected time spent in each state. # Arguments - `hmm::AbstractHMM`: a fitted HMM satisfying the HiddenMarkovModels.jl interface. @@ -38,7 +30,7 @@ integral transform (see [`EmissionModels._emission_to_driver`](@ref)). Component # Returns - [`StochasticDriverResult`](@ref) with per-state driver pools and usage. """ -function EmissionModels.stochastic_drivers( +function stochastic_drivers( hmm::AbstractHMM, obs_seq::AbstractVector; control_seq::AbstractVector=fill(nothing, length(obs_seq)), @@ -57,7 +49,7 @@ function EmissionModels.stochastic_drivers( # Driver dimension from a probe on the first observation's sampled-state-able # emission (all states share the emission dimension in an HMM). - probe = EmissionModels._emission_to_driver( + probe = _emission_to_driver( obs_distributions(hmm, control_seq[1])[1], obs_seq[1], control_seq[1] ) D = length(probe) @@ -66,12 +58,9 @@ function EmissionModels.stochastic_drivers( ε_lists = [Vector{Vector{Tε}}() for _ in 1:K] for _ in 1:n_samples for t in 1:T_len - z = EmissionModels._sample_categorical(view(γ, :, t)) + z = _sample_categorical(view(γ, :, t)) dist = obs_distributions(hmm, control_seq[t])[z] - push!( - ε_lists[z], - EmissionModels._emission_to_driver(dist, obs_seq[t], control_seq[t]), - ) + push!(ε_lists[z], _emission_to_driver(dist, obs_seq[t], control_seq[t])) end end @@ -87,5 +76,3 @@ function EmissionModels.stochastic_drivers( return StochasticDriverResult(ε_pools, collect(Tε, usage)) end - -end # module diff --git a/src/acdc/interface.jl b/src/acdc/interface.jl index 3087c17..e889bc2 100644 --- a/src/acdc/interface.jl +++ b/src/acdc/interface.jl @@ -1,5 +1,5 @@ #= -ACDC — Accumulated Cutoff Discrepancy Criterion. +ACDC: Accumulated Cutoff Discrepancy Criterion. Robust model selection that measures component-level discrepancy via the "stochastic drivers" framework (Li et al., 2026). The generative process is @@ -19,11 +19,7 @@ per-component discrepancies all fall below a cutoff ``\\rho``. This file holds the model-agnostic core: result types, discrepancy measures, and the loss/selection machinery. Model-specific recovery of the drivers is done by `stochastic_drivers`, whose methods live alongside the model types they -support (for HiddenMarkovModels.jl HMMs, in the package extension). - -This is a diagnostic / model-selection tool run once on a fitted model, not a -hot path, so the code favors clarity over the zero-allocation discipline the -emission `fit!`/`logdensityof` paths follow. +support. =# """ @@ -86,15 +82,14 @@ Recover the stochastic drivers ``\\varepsilon_{n,k}`` for a fitted `model` by inverting its generative process. Returns a [`StochasticDriverResult`](@ref). This is a generic function; methods are defined per model type. The method for -HiddenMarkovModels.jl `AbstractHMM`s is provided by the package extension that -loads with `HiddenMarkovModels`. The fallback below errors for unsupported models. +HiddenMarkovModels.jl `AbstractHMM`s lives in `hmm.jl`. The fallback below errors +for unsupported models. """ function stochastic_drivers(model, data; kwargs...) throw( ArgumentError( "ACDC has no `stochastic_drivers` method for a model of type " * - "$(typeof(model)). Load `HiddenMarkovModels` to enable support for " * - "`AbstractHMM` models, or define a `stochastic_drivers` method.", + "$(typeof(model)). Define a `stochastic_drivers` method to add support.", ), ) end From da25f74560c8a51b1c16aa9258b1d227ca48f2cd Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sun, 28 Jun 2026 10:53:12 -0400 Subject: [PATCH 6/6] Update hmm.jl --- src/acdc/hmm.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/acdc/hmm.jl b/src/acdc/hmm.jl index 752cd50..55aea61 100644 --- a/src/acdc/hmm.jl +++ b/src/acdc/hmm.jl @@ -12,7 +12,7 @@ Recover the ACDC stochastic drivers for a fitted HiddenMarkovModels.jl `hmm`. For each time step the hidden state is sampled from its forward-backward posterior, and that state's emission is inverted to a driver via the probability -integral transform (see [`_emission_to_driver`](@ref)). Component "usage" is the +integral transform (see `_emission_to_driver`). Component "usage" is the posterior expected time spent in each state. # Arguments