GLMs work with ControlledEmission

Project.toml

@@ -6,6 +6,7 @@ authors = ["Ryan Senne"]
 [deps]
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+HiddenMarkovModels = "84ca31d5-effc-45e0-bfda-5a68cd981f47"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
@@ -16,6 +17,7 @@ StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0"
 [compat]
 DensityInterface = "0.4.0"
 Distributions = "0.25.122"
+HiddenMarkovModels = "0.7.1"
 LinearAlgebra = "1.12.0"
 LogExpFunctions = "0.3.29, 1"
 Optim = "1.13.3, 2"

src/EmissionModels.jl

@@ -2,6 +2,7 @@ module EmissionModels
 
 using Distributions: Normal, Bernoulli, Poisson, Chisq
 using DensityInterface
+using HiddenMarkovModels: ControlledEmission
 using LinearAlgebra
 using LogExpFunctions: logaddexp, logsumexp, log1pexp, logistic
 using Optim: optimize, TwiceDifferentiable, Newton, LBFGS, LineSearches

src/glms/glm.jl

@@ -1,20 +1,26 @@
 """
-    AbstractGLM
+    AbstractGLM <: HiddenMarkovModels.ControlledEmission
 
 Abstract type for Generalized Linear Model emission distributions.
 
-GLM subtypes should implement the HiddenMarkovModels.jl interface:
-- `DensityInterface.DensityKind(::YourGLM)` → `HasDensity()`
+Subtyping `ControlledEmission` lets a `Vector` of GLMs serve as the `dists` of a
+`HiddenMarkovModels.ControlledEmissionHMM`. Each concrete GLM implements the
+keyword (`control_seq`) interface internally:
 - `DensityInterface.logdensityof(glm, obs; control_seq)` — log density
 - `Random.rand(rng, glm; control_seq)` — conditional sample
 - `StatsAPI.fit!(glm, obs_seq, weight_seq; control_seq)` — weighted in-place update
 
+and the `ControlledEmission` positional signatures HMM expects — `logdensityof(glm,
+obs, control)`, `rand(rng, glm, control)`, `fit!(glm, obs_seq, control_seq, weights)`
+— are provided as thin adapters at the bottom of this file. `DensityKind` is
+inherited from `ControlledEmission`.
+
 Univariate types (`GaussianGLM`, `BernoulliGLM`, `PoissonGLM`) carry a coefficient
 vector `β` and emit scalar observations. Multivariate variants (`MvGaussianGLM`,
 `MvBernoulliGLM`, `MvPoissonGLM`) carry a coefficient matrix `B` of size `p × k`
 and emit length-`k` observation vectors.
 """
-abstract type AbstractGLM end
+abstract type AbstractGLM <: ControlledEmission end
 
 """
     AbstractPrior
@@ -108,8 +114,6 @@ function GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior)
 end
 GaussianGLM(β::AbstractVector, σ2::Real) = GaussianGLM(β, σ2, NoPrior())
 
-DensityInterface.DensityKind(::GaussianGLM) = DensityInterface.HasDensity()
-
 function DensityInterface.logdensityof(
     reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real}
 )
@@ -500,8 +504,6 @@ function BernoulliGLM(β::AbstractVector, prior::AbstractPrior)
 end
 BernoulliGLM(β::AbstractVector) = BernoulliGLM(β, NoPrior())
 
-DensityInterface.DensityKind(::BernoulliGLM) = DensityInterface.HasDensity()
-
 function DensityInterface.logdensityof(
     glm::BernoulliGLM, y::Integer; control_seq::AbstractVector{<:Real}
 )
@@ -588,8 +590,6 @@ function PoissonGLM(β::AbstractVector, prior::AbstractPrior)
 end
 PoissonGLM(β::AbstractVector) = PoissonGLM(β, NoPrior())
 
-DensityInterface.DensityKind(::PoissonGLM) = DensityInterface.HasDensity()
-
 function DensityInterface.logdensityof(
     glm::PoissonGLM, y::Integer; control_seq::AbstractVector{<:Real}
 )
@@ -709,8 +709,6 @@ function MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix, prior::AbstractPri
 end
 MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix) = MvGaussianGLM(B, Σ, NoPrior())
 
-DensityInterface.DensityKind(::MvGaussianGLM) = DensityInterface.HasDensity()
-
 """
     logdensityof(glm::MvGaussianGLM, y::AbstractVector; control_seq)
 
@@ -960,8 +958,6 @@ function MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior)
 end
 MvBernoulliGLM(B::AbstractMatrix) = MvBernoulliGLM(B, NoPrior())
 
-DensityInterface.DensityKind(::MvBernoulliGLM) = DensityInterface.HasDensity()
-
 function DensityInterface.logdensityof(
     glm::MvBernoulliGLM, y::AbstractVector; control_seq::AbstractVector{<:Real}
 )
@@ -1133,8 +1129,6 @@ function MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior)
 end
 MvPoissonGLM(B::AbstractMatrix) = MvPoissonGLM(B, NoPrior())
 
-DensityInterface.DensityKind(::MvPoissonGLM) = DensityInterface.HasDensity()
-
 function DensityInterface.logdensityof(
     glm::MvPoissonGLM, y::AbstractVector; control_seq::AbstractVector{<:Real}
 )
@@ -1260,3 +1254,54 @@ function StatsAPI.fit!(
     end
     return glm
 end
+
+#= ─── HiddenMarkovModels.ControlledEmission interface ───────────────────────
+   `AbstractGLM <: ControlledEmission`, so a `Vector` of GLMs is a valid `dists`
+   for a `ControlledEmissionHMM`. That HMM drives each emission through the
+   control-aware *positional* signatures below; each `control` is a single
+   timestep's covariate vector — exactly the `control_seq` argument the keyword
+   methods above already consume — and the fit-time `control_seq` is a vector of
+   such vectors (one per timestep). The adapters delegate to the keyword
+   implementations so the actual math has a single source of truth. =#
+
+# Length of one covariate vector for this GLM (the GLM's input dimension `p`).
+_indim(glm::Union{GaussianGLM,BernoulliGLM,PoissonGLM}) = length(glm.β)
+_indim(glm::Union{MvGaussianGLM,MvBernoulliGLM,MvPoissonGLM}) = glm.in_dim
+
+function DensityInterface.logdensityof(
+    glm::AbstractGLM, obs, control::AbstractVector{<:Real}
+)
+    return logdensityof(glm, obs; control_seq=control)
+end
+
+function Random.rand(rng::AbstractRNG, glm::AbstractGLM, control::AbstractVector{<:Real})
+    return rand(rng, glm; control_seq=control)
+end
+
+#= Zero-copy `n×p` design matrix over a length-`n` vector of length-`p` covariate
+   vectors. `ControlledEmissionHMM` hands `fit!` a `control_seq` shaped as a
+   `Vector{<:AbstractVector}` (one covariate vector per timestep), whereas the
+   matrix-based keyword `fit!` implementations want an `n×p` matrix. This presents
+   the former as the latter without copying: `view(M, i, :)` returns the i-th
+   covariate vector directly, so the existing inner loops stay allocation-free. =#
+struct _ControlRowsMatrix{T,V<:AbstractVector{<:AbstractVector}} <: AbstractMatrix{T}
+    rows::V
+    p::Int
+end
+function _ControlRowsMatrix(rows::V, p::Int) where {V<:AbstractVector{<:AbstractVector}}
+    return _ControlRowsMatrix{eltype(eltype(V)),V}(rows, p)
+end
+Base.size(M::_ControlRowsMatrix) = (length(M.rows), M.p)
+Base.@propagate_inbounds Base.getindex(M::_ControlRowsMatrix, i::Int, j::Int) = M.rows[i][j]
+Base.@propagate_inbounds Base.view(M::_ControlRowsMatrix, i::Integer, ::Colon) = M.rows[i]
+
+function StatsAPI.fit!(
+    glm::AbstractGLM,
+    obs_seq::AbstractVector,
+    control_seq::AbstractVector{<:AbstractVector},
+    weights::AbstractVector{<:Real};
+    kwargs...,
+)
+    X = _ControlRowsMatrix(control_seq, _indim(glm))
+    return fit!(glm, obs_seq, weights; control_seq=X, kwargs...)
+end

test/Project.toml

@@ -2,7 +2,9 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+EmissionModels = "1e2dd27c-41a5-43b2-863c-3eddd0c72c67"
 HiddenMarkovModels = "84ca31d5-effc-45e0-bfda-5a68cd981f47"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
@@ -14,4 +16,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 HiddenMarkovModels = "0.7.1"
-JuliaFormatter = "1.0.62"
+JuliaFormatter = "1.0.62"

test/glm/test_controlled_hmm.jl

@@ -0,0 +1,109 @@
+using EmissionModels
+using Distributions
+using HiddenMarkovModels
+using HiddenMarkovModels: ControlledEmission, ControlledEmissionHMM, baum_welch, forward
+using DensityInterface
+using StatsAPI
+using Random
+using LinearAlgebra
+using Test
+
+#= GLMs must subtype `ControlledEmission` so a `Vector` of them is a valid
+   `dists` for a `ControlledEmissionHMM`, and the control-aware positional
+   interface (`logdensityof(d, obs, control)`, `rand(rng, d, control)`,
+   `fit!(d, obs_seq, control_seq, weights)`) must drive inference and learning. =#
+
+@testset "GLMs as ControlledEmissionHMM emissions" begin
+    @testset "subtype relationship" begin
+        for G in (
+            GaussianGLM,
+            BernoulliGLM,
+            PoissonGLM,
+            MvGaussianGLM,
+            MvBernoulliGLM,
+            MvPoissonGLM,
+        )
+            @test G <: ControlledEmission
+        end
+    end
+
+    @testset "positional control-aware interface delegates to keyword methods" begin
+        rng = MersenneTwister(0)
+        p = 3
+        x = vcat(1.0, randn(rng, p - 1))
+
+        pg = PoissonGLM(randn(rng, p) .* 0.2)
+        @test logdensityof(pg, 2, x) == logdensityof(pg, 2; control_seq=x)
+
+        gg = GaussianGLM(randn(rng, p), 1.5)
+        @test logdensityof(gg, 0.7, x) == logdensityof(gg, 0.7; control_seq=x)
+
+        # rand with a fixed rng must match the keyword path
+        @test rand(MersenneTwister(7), pg, x) == rand(MersenneTwister(7), pg; control_seq=x)
+    end
+
+    @testset "fit! via vector-of-vectors control_seq matches matrix fit!" begin
+        rng = MersenneTwister(1)
+        n, p = 200, 3
+        X = [vcat(1.0, randn(rng, p - 1)) for _ in 1:n]
+        Xmat = permutedims(reduce(hcat, X))           # n×p matrix form
+        β_true = [0.5, -0.8, 0.3]
+        y = [rand(rng, Distributions.Poisson(exp(dot(β_true, X[i])))) for i in 1:n]
+        w = ones(n)
+
+        g_pos = PoissonGLM(zeros(p))
+        g_kw = PoissonGLM(zeros(p))
+        fit!(g_pos, y, X, w)                          # positional ControlledEmission path
+        fit!(g_kw, y, w; control_seq=Xmat)            # keyword matrix path
+        @test g_pos.β ≈ g_kw.β rtol = 1e-8
+    end
+
+    @testset "Poisson-GLM ControlledEmissionHMM: sample, forward, baum_welch" begin
+        rng = MersenneTwister(42)
+        p, T = 3, 600
+        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)
+
+        control_seq = [vcat(1.0, randn(rng, p - 1)) for _ in 1:T]
+        obs_seq = rand(rng, hmm, control_seq).obs_seq
+        @test length(obs_seq) == T
+
+        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
+        init0 = [0.5, 0.5]
+        dists0 = [PoissonGLM(zeros(p)), PoissonGLM(zeros(p))]
+        hmm0 = ControlledEmissionHMM(init0, copy(trans), dists0)
+        _, lls = baum_welch(hmm0, obs_seq, control_seq; seq_ends=[T], max_iterations=30)
+        @test all(diff(lls) .>= -1e-6)
+        @test last(lls) >= first(lls)
+    end
+
+    @testset "MvGaussian-GLM ControlledEmissionHMM end-to-end" begin
+        rng = MersenneTwister(123)
+        p, k, T = 2, 2, 500
+        init = [0.5, 0.5]
+        trans = [0.9 0.1; 0.1 0.9]
+        B1 = [1.0 0.0; 0.5 -0.5]
+        B2 = [-1.0 0.5; 0.0 1.0]
+        Σ = Matrix{Float64}(I, k, k)
+        dists = [MvGaussianGLM(B1, copy(Σ)), MvGaussianGLM(B2, copy(Σ))]
+        hmm = ControlledEmissionHMM(init, trans, dists)
+
+        control_seq = [vcat(1.0, randn(rng)) for _ in 1:T]   # length-p = 2 each
+        obs_seq = rand(rng, hmm, control_seq).obs_seq
+        @test length(obs_seq) == T
+        @test length(first(obs_seq)) == k
+
+        dists0 = [
+            MvGaussianGLM(zeros(p, k), Matrix{Float64}(I, k, k)),
+            MvGaussianGLM(zeros(p, k), Matrix{Float64}(I, k, k)),
+        ]
+        hmm0 = ControlledEmissionHMM(init, copy(trans), dists0)
+        _, lls = baum_welch(hmm0, obs_seq, control_seq; seq_ends=[T], max_iterations=30)
+        @test all(diff(lls) .>= -1e-6)
+    end
+end

test/runtests.jl

@@ -21,6 +21,7 @@ using JuliaFormatter
         include("glm/gaussian.jl")
         include("glm/test_bernoulli_poisson.jl")
         include("glm/test_promotion_and_types.jl")
+        include("glm/test_controlled_hmm.jl")
     end
     @testset "Zero-inflated models" begin
         include("zeroinflated/test_poisson.jl")