From a6bd2d113542ec6021818d9abb93716d8faa349c Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 1 Jan 2026 12:54:37 -0500 Subject: [PATCH 01/20] begin GLM additions --- src/EmissionModels.jl | 4 +- src/glms/glm.jl | 31 ++++++++++++ test/glm/gaussian.jl | 111 ++++++++++++++++++++++++++++++++++++++++++ test/runtests.jl | 4 ++ 4 files changed, 149 insertions(+), 1 deletion(-) create mode 100644 test/glm/gaussian.jl diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index 0491f6b..49e05dd 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -1,6 +1,6 @@ module EmissionModels -using Distributions: Poisson, Chisq +using Distributions: Normal, Bernoulli, Poisson, Chisq using DensityInterface using LinearAlgebra using LogExpFunctions: logaddexp, logsumexp @@ -13,10 +13,12 @@ using StatsAPI: fit! include("zeroinflated/poisson.jl") include("multivariate/t.jl") +include("glms/glm.jl") # exports export rand, logdensityof, fit! export PoissonZeroInflated export MultivariateT, MultivariateTDiag +export GaussianGLM end diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 0221454..a24381a 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -10,3 +10,34 @@ GLM subtypes should implement the HiddenMarkovModels.jl interface: - `StatsAPI.fit!(glm, obs_seq, weight_seq)` - Update parameters """ abstract type AbstractGLM end + +mutable struct GaussianGLM{T<:Real} <: AbstractGLM + β::Vector{T} # Coefficients + σ2::T # Variance +end + +DensityInterface.DensityKind(::GaussianGLM) = DensityInterface.HasDensity() + +function DensityInterface.logdensityof(reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real}) + μ = dot(reg.β, control_seq) + return -0.5 * log(2π * reg.σ2) - 0.5 * ((y-μ)^2 / reg.σ2) +end + +function Random.rand(rng::AbstractRNG, reg::GaussianGLM; control_seq::AbstractVector{<:Real}) + return rand(rng, Normal(dot(reg.β, control_seq), sqrt(reg.σ2))) +end + +function StatsAPI.fit!(reg::GaussianGLM, obs_seq::AbstractVector{<:Real}, weights::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}) + # Fit coefficients using weighted least squares + W = Diagonal(weights) + XWX = control_seq' * W * control_seq + XWy = control_seq' * W * obs_seq + reg.β = XWX \ XWy + + # Update variance + residuals = obs_seq .- control_seq * reg.β + reg.σ2 = sum(weights .* (residuals .^ 2)) / sum(weights) + + return reg +end \ No newline at end of file diff --git a/test/glm/gaussian.jl b/test/glm/gaussian.jl new file mode 100644 index 0000000..641576f --- /dev/null +++ b/test/glm/gaussian.jl @@ -0,0 +1,111 @@ + +using Test +using Random +using LinearAlgebra +using Distributions +using DensityInterface +using StableRNGs +using StatsAPI +using EmissionModels + +rng = StableRNG(1234) + +function _synthetic_gaussian_glm(rng, n::Int, p::Int; β::Vector{Float64}, σ2::Float64, weights=:uniform) + @assert length(β) == p + X = hcat(ones(n), randn(rng, n, p-1)) # intercept + random features + μ = X * β + y = rand.(Ref(rng), Normal.(μ, sqrt(σ2))) + w = weights === :uniform ? ones(n) : + weights === :random ? (rand(rng, n) .+ 0.5) : + error("weights must be :uniform or :random") + return X, y, w +end + +_expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) + +@testset "GaussianGLM Basic Functionality" begin + @testset "Constructor" begin + glm = GaussianGLM([0.5, -1.0], 1.0) + @test glm.β == [0.5, -1.0] + @test glm.σ2 == 1.0 + + glm2 = GaussianGLM([1, 2], 3) + @test glm2.β == [1, 2] + @test glm2.σ2 == 3 + end + + @testset "DensityInterface" begin + glm = GaussianGLM([0.5, -1.0], 1.0) + + # Trait + @test DensityKind(glm) == HasDensity() + + X = [1.0 2.0; 1.0 3.0; 1.0 4.0] + μ = X * glm.β + y = rand.(Ref(rng), Normal.(μ, sqrt(glm.σ2))) + + # logdensityof matches closed form + for i in eachindex(y) + x_i = vec(X[i, :]) # ensure Vector + logp = logdensityof(glm, y[i]; control_seq=x_i) + @test isfinite(logp) + + expected = _expected_logpdf(y[i], dot(glm.β, x_i), glm.σ2) + @test logp ≈ expected rtol=1e-12 atol=0.0 + end + + # Higher density near the mean than far away + x0 = vec(X[1, :]) + y_at_mean = dot(glm.β, x0) + y_far = y_at_mean + 10.0 + @test logdensityof(glm, y_at_mean; control_seq=x0) > logdensityof(glm, y_far; control_seq=x0) + end + + @testset "Random sampling" begin + glm = GaussianGLM([0.5, -1.0], 2.0) # variance 2.0 + x = [1.0, 3.0] + n = 10_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + + @test all(s -> s isa Real, samples) + + # Empirical mean should be close to dot(β,x) + μ = dot(glm.β, x) + m = mean(samples) + @test m ≈ μ atol=0.15 + + # Empirical variance should be close to σ2 + v = var(samples) + @test v ≈ glm.σ2 atol=0.25 + end + + @testset "fit! with uniform weights" begin + β_true = [0.3, -1.2] # p=2 + σ2_true = 0.7 + X, y, w = _synthetic_gaussian_glm(rng, 2000, 2; β=β_true, σ2=σ2_true, weights=:uniform) + + glm = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm, y, w; control_seq=X) + + # Coefficients should be close with enough data + @test glm.β ≈ β_true atol=0.08 + @test glm.σ2 ≈ σ2_true atol=0.10 + + @test all(isfinite, glm.β) + @test isfinite(glm.σ2) + @test glm.σ2 > 0 + end + + @testset "fit! with weighted observations" begin + β_true = [1.0, 0.6] + σ2_true = 1.5 + X, y, w = _synthetic_gaussian_glm(rng, 2500, 2; β=β_true, σ2=σ2_true, weights=:random) + + glm = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm, y, w; control_seq=X) + + @test glm.β ≈ β_true atol=0.10 + @test glm.σ2 ≈ σ2_true atol=0.15 + end + +end \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index 1c6a94a..1e58b66 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -13,6 +13,10 @@ using StatsAPI @info "Skipping JET on Julia $VERSION (requires >=1.10 and a release build)" end end + + @testset "GLM Models" begin + include("glm/gaussian.jl") + end @testset "Zero-inflated models" begin include("zeroinflated/test_poisson.jl") end From e693efe8076aeba5b8a615437cb11b55fe8b1155 Mon Sep 17 00:00:00 2001 From: Ryan Senne Date: Sun, 5 Apr 2026 13:42:01 -0400 Subject: [PATCH 02/20] Updates to multivariate t dist. But need to move to an external workspace version. --- src/multivariate/t.jl | 330 +++++++++++++++++++++--------------------- 1 file changed, 166 insertions(+), 164 deletions(-) diff --git a/src/multivariate/t.jl b/src/multivariate/t.jl index 2eff54d..3acdd7b 100644 --- a/src/multivariate/t.jl +++ b/src/multivariate/t.jl @@ -28,12 +28,15 @@ mutable struct MultivariateT{T<:Real} ν::T # Cached values for efficiency - Σ_chol::Cholesky{T,Matrix{T}} # Cholesky decomposition of Σ - logdetΣ::T # log determinant of Σ - dim::Int # Dimension + Σ_chol::Cholesky{T,Matrix{T}} + logdetΣ::T + dim::Int + + # Scratch buffers for fit! (sequential use only — not thread-safe for concurrent fit!) + _diff::Vector{T} + _z::Vector{T} function MultivariateT{T}(μ::Vector{T}, Σ::Matrix{T}, ν::T) where {T<:Real} - # Validation ν > 0 || throw(ArgumentError("ν must be positive, got $ν")) dim = length(μ) @@ -42,7 +45,6 @@ mutable struct MultivariateT{T<:Real} size(Σ) == (dim, dim) || throw(DimensionMismatch("Σ must be $(dim)×$(dim), got $(size(Σ))")) - # Check positive definiteness via Cholesky Σ_chol = try cholesky(Σ) catch @@ -51,11 +53,10 @@ mutable struct MultivariateT{T<:Real} logdetΣ = logdet(Σ_chol) - return new{T}(μ, Σ, ν, Σ_chol, logdetΣ, dim) + return new{T}(μ, Σ, ν, Σ_chol, logdetΣ, dim, zeros(T, dim), zeros(T, dim)) end end -# Outer constructors for convenience MultivariateT(μ::Vector{T}, Σ::Matrix{T}, ν::T) where {T<:Real} = MultivariateT{T}(μ, Σ, ν) function MultivariateT(μ::Vector, Σ::Matrix, ν::Real) @@ -93,11 +94,13 @@ mutable struct MultivariateTDiag{T<:Real} ν::T # Cached values for efficiency - logdetΣ::T # sum(log.(σ²)) - dim::Int # Dimension + logdetΣ::T + dim::Int + + # Scratch buffer for fit! (sequential use only — not thread-safe for concurrent fit!) + _diff::Vector{T} function MultivariateTDiag{T}(μ::Vector{T}, σ²::Vector{T}, ν::T) where {T<:Real} - # Validation ν > 0 || throw(ArgumentError("ν must be positive, got $ν")) dim = length(μ) @@ -111,11 +114,10 @@ mutable struct MultivariateTDiag{T<:Real} logdetΣ = sum(log, σ²) - return new{T}(μ, σ², ν, logdetΣ, dim) + return new{T}(μ, σ², ν, logdetΣ, dim, zeros(T, dim)) end end -# Outer constructors for convenience function MultivariateTDiag(μ::Vector{T}, σ²::Vector{T}, ν::T) where {T<:Real} return MultivariateTDiag{T}(μ, σ², ν) end @@ -132,15 +134,12 @@ DensityInterface.DensityKind(::MultivariateT) = DensityInterface.HasDensity() DensityInterface.DensityKind(::MultivariateTDiag) = DensityInterface.HasDensity() """ - logdensityof(dist::MultivariateT, x::Vector) + logdensityof(dist::MultivariateT, x::AbstractVector) Compute the log probability density of the multivariate t-distribution at `x`. -The log density is: -``` -log p(x) = log Γ((ν+d)/2) - log Γ(ν/2) - (d/2)log(νπ) - (1/2)log|Σ| - - ((ν+d)/2) log(1 + (1/ν)(x-μ)' Σ⁻¹ (x-μ)) -``` +Allocates one vector (the residual). The triangular solve is performed in-place +on that vector to avoid a second allocation. """ function DensityInterface.logdensityof(dist::MultivariateT, x::AbstractVector) length(x) == dist.dim || @@ -149,24 +148,23 @@ function DensityInterface.logdensityof(dist::MultivariateT, x::AbstractVector) d = dist.dim ν = dist.ν - # Compute (x - μ)' Σ⁻¹ (x - μ) efficiently using Cholesky - diff = x .- dist.μ - z = dist.Σ_chol.L \ diff # Solve L z = diff - mahal² = sum(abs2, z) # ||z||² = (x-μ)' Σ⁻¹ (x-μ) + # 1 allocation. ldiv! solves L\(x-μ) in-place, avoiding a second allocation. + diff = x - dist.μ + ldiv!(dist.Σ_chol.L, diff) + mahal² = sum(abs2, diff) - # Log density computation log_norm = loggamma((ν + d) / 2) - loggamma(ν / 2) - (d / 2) * log(ν * π) - dist.logdetΣ / 2 - log_kernel = -((ν + d) / 2) * log1p(mahal² / ν) - - return log_norm + log_kernel + return log_norm - ((ν + d) / 2) * log1p(mahal² / ν) end """ - logdensityof(dist::MultivariateTDiag, x::Vector) + logdensityof(dist::MultivariateTDiag, x::AbstractVector) Compute the log probability density of the diagonal multivariate t-distribution at `x`. + +Zero allocations: Mahalanobis distance is accumulated element-wise inline. """ function DensityInterface.logdensityof(dist::MultivariateTDiag, x::AbstractVector) length(x) == dist.dim || @@ -175,20 +173,15 @@ function DensityInterface.logdensityof(dist::MultivariateTDiag, x::AbstractVecto d = dist.dim ν = dist.ν - # Compute (x - μ)' Σ⁻¹ (x - μ) for diagonal Σ - diff = x .- dist.μ - mahal² = sum(diff[i]^2 / dist.σ²[i] for i in 1:d) + mahal² = sum((x[i] - dist.μ[i])^2 / dist.σ²[i] for i in 1:d) - # Log density computation log_norm = loggamma((ν + d) / 2) - loggamma(ν / 2) - (d / 2) * log(ν * π) - dist.logdetΣ / 2 - log_kernel = -((ν + d) / 2) * log1p(mahal² / ν) - - return log_norm + log_kernel + return log_norm - ((ν + d) / 2) * log1p(mahal² / ν) end -# sampling +# Random sampling """ rand([rng], dist::MultivariateT) @@ -199,15 +192,9 @@ Uses the representation: X = μ + √(ν/U) × Z where Z ~ N(0, Σ) and U ~ χ²(ν) are independent. """ function Random.rand(rng::AbstractRNG, dist::MultivariateT) - # Sample from chi-squared with ν degrees of freedom u = rand(rng, Chisq(dist.ν)) - - # Sample from N(0, I) and transform to N(0, Σ) z = randn(rng, dist.dim) - scaled_noise = dist.Σ_chol.L * z # L z ~ N(0, Σ) where Σ = L L' - - # Apply t-distribution scaling - return dist.μ .+ sqrt(dist.ν / u) .* scaled_noise + return dist.μ .+ sqrt(dist.ν / u) .* (dist.Σ_chol.L * z) end """ @@ -216,15 +203,9 @@ end Generate a random sample from the diagonal multivariate t-distribution. """ function Random.rand(rng::AbstractRNG, dist::MultivariateTDiag) - # Sample from chi-squared with ν degrees of freedom u = rand(rng, Chisq(dist.ν)) - - # Sample from N(0, diag(σ²)) z = randn(rng, dist.dim) - scaled_noise = sqrt.(dist.σ²) .* z - - # Apply t-distribution scaling - return dist.μ .+ sqrt(dist.ν / u) .* scaled_noise + return dist.μ .+ sqrt(dist.ν / u) .* sqrt.(dist.σ²) .* z end """ @@ -268,7 +249,6 @@ function StatsAPI.fit!( isempty(obs_seq) && return dist - # Check dimensions d = dist.dim for (i, obs) in enumerate(obs_seq) length(obs) == d || throw( @@ -276,75 +256,85 @@ function StatsAPI.fit!( ) end - # Filter out zero-weight observations - nonzero_idx = findall(w -> w > 0, weight_seq) - isempty(nonzero_idx) && return dist - - obs_vec = [obs_seq[i] for i in nonzero_idx] - weights = [weight_seq[i] for i in nonzero_idx] - n = length(obs_vec) + # Total weight (skipping zeros avoids inflating the denominator) + weight_sum = zero(eltype(weight_seq)) + for w in weight_seq + w > 0 && (weight_sum += w) + end + iszero(weight_sum) && return dist - # Normalize weights - weight_sum = sum(weights) - weights ./= weight_sum + T = eltype(dist.μ) + n = length(obs_seq) - # EM algorithm + # Allocate working arrays once per fit! call (not per iteration) + posterior_weights = Vector{T}(undef, n) + Σ_acc = zeros(T, d, d) + μ_acc = zeros(T, d) μ_old = copy(dist.μ) Σ_old = copy(dist.Σ) ν_old = dist.ν - for iter in 1:max_iter - # E-step: Compute posterior weights - # w_i = (ν + d) / (ν + δ_i²) where δ_i² is Mahalanobis distance - posterior_weights = zeros(n) - - for i in 1:n - diff = obs_vec[i] .- dist.μ - z = dist.Σ_chol.L \ diff - mahal² = sum(abs2, z) - posterior_weights[i] = (dist.ν + d) / (dist.ν + mahal²) + for _ in 1:max_iter + # E-step + μ M-step accumulation (single pass; uses dist._diff and dist._z) + fill!(μ_acc, zero(T)) + weight_post_sum = zero(T) + + for i in eachindex(obs_seq) + w = weight_seq[i] + if w > 0 + dist._diff .= obs_seq[i] .- dist.μ + ldiv!(dist._z, dist.Σ_chol.L, dist._diff) + pw = T((dist.ν + d) / (dist.ν + sum(abs2, dist._z))) + posterior_weights[i] = pw + wp = T(w / weight_sum) * pw + for j in 1:d + μ_acc[j] += wp * obs_seq[i][j] + end + weight_post_sum += wp + else + posterior_weights[i] = zero(T) + end end - # M-step: Update parameters - # Update μ - weighted_sum = sum(weights[i] * posterior_weights[i] * obs_vec[i] for i in 1:n) - weight_post_sum = sum(weights[i] * posterior_weights[i] for i in 1:n) - dist.μ .= weighted_sum ./ weight_post_sum - - # Update Σ - Σ_new = zeros(d, d) - for i in 1:n - diff = obs_vec[i] .- dist.μ - Σ_new .+= weights[i] * posterior_weights[i] * (diff * diff') - end - Σ_new ./= sum(weights) + dist.μ .= μ_acc ./ weight_post_sum - # Ensure symmetry and positive definiteness - Σ_new .= (Σ_new .+ Σ_new') ./ 2 + # Σ M-step: rank-1 outer product accumulation using BLAS.ger! + fill!(Σ_acc, zero(T)) + for i in eachindex(obs_seq) + weight_seq[i] > 0 || continue + dist._diff .= obs_seq[i] .- dist.μ + wp = T(weight_seq[i] / weight_sum) * posterior_weights[i] + BLAS.ger!(wp, dist._diff, dist._diff, Σ_acc) + end - # Add small regularization if needed - min_eig = minimum(eigvals(Hermitian(Σ_new))) - if min_eig <= 0 - Σ_new .+= (abs(min_eig) + 1e-6) * I + # Symmetrize numerical noise + for j in 1:d, k in 1:j-1 + s = (Σ_acc[j, k] + Σ_acc[k, j]) / 2 + Σ_acc[j, k] = s + Σ_acc[k, j] = s end - dist.Σ .= Σ_new - dist.Σ_chol = cholesky(dist.Σ) + # Update Σ and Cholesky, regularizing if needed + Σ_chol_new = try + cholesky(Symmetric(Σ_acc)) + catch + min_eig = minimum(eigvals(Hermitian(Σ_acc))) + Σ_acc .+= (abs(min_eig) + 1e-6) * I + cholesky(Symmetric(Σ_acc)) + end + dist.Σ .= Σ_acc + dist.Σ_chol = Σ_chol_new dist.logdetΣ = logdet(dist.Σ_chol) - # Update ν (if not fixed) + # ν M-step if !fix_nu - # Use Optim.jl with Newton's method to maximize the Q function w.r.t. ν - # The equation to solve is: - # -ψ(ν/2) + log(ν/2) + 1 + (1/n)Σw_i(log(u_i) - u_i) + ψ((ν+d)/2) - log((ν+d)/2) = 0 - # where u_i are the posterior weights - - avg_log_u_minus_u = sum( - weights[i] * (log(posterior_weights[i]) - posterior_weights[i]) for i in 1:n - ) + avg_log_u_minus_u = zero(T) + for i in eachindex(obs_seq) + weight_seq[i] > 0 || continue + pw = posterior_weights[i] + avg_log_u_minus_u += T(weight_seq[i] / weight_sum) * (log(pw) - pw) + end - # Optimize over log(ν) to ensure ν > 0 - # Let x = log(ν), so ν = exp(x) function objective(x::Vector) ν_val = exp(x[1]) f = @@ -353,7 +343,7 @@ function StatsAPI.fit!( 1 + avg_log_u_minus_u + digamma((ν_val + d) / 2) - log((ν_val + d) / 2) - return f^2 # Minimize squared residual + return f^2 end function gradient!(G, x::Vector) @@ -364,12 +354,10 @@ function StatsAPI.fit!( 1 + avg_log_u_minus_u + digamma((ν_val + d) / 2) - log((ν_val + d) / 2) - # df/dν df_dν = -0.5 * trigamma(ν_val / 2) + 1 / ν_val + 0.5 * trigamma((ν_val + d) / 2) - 1 / (ν_val + d) - # d(f²)/dx = d(f²)/dν * dν/dx = 2f * df/dν * ν (since dν/dx = ν) return G[1] = 2 * f * df_dν * ν_val end @@ -381,30 +369,32 @@ function StatsAPI.fit!( 1 + avg_log_u_minus_u + digamma((ν_val + d) / 2) - log((ν_val + d) / 2) - # df/dν df_dν = -0.5 * trigamma(ν_val / 2) + 1 / ν_val + 0.5 * trigamma((ν_val + d) / 2) - 1 / (ν_val + d) - # d²f/dν² d2f_dν2 = -0.25 * polygamma(2, ν_val / 2) - 1 / ν_val^2 + 0.25 * polygamma(2, (ν_val + d) / 2) + 1 / (ν_val + d)^2 - # d²(f²)/dx² using chain rule: d²(f²)/dx² = [2(df/dν)² + 2f*d²f/dν²]*ν² + 2f*df/dν*ν return H[1, 1] = (2 * df_dν^2 + 2 * f * d2f_dν2) * ν_val^2 + 2 * f * df_dν * ν_val end - # Use Newton's method with analytical derivatives td = TwiceDifferentiable(objective, gradient!, hessian!, [log(dist.ν)]) result = optimize(td, [log(dist.ν)], Newton()) dist.ν = exp(Optim.minimizer(result)[1]) end - # Check convergence - μ_diff = maximum(abs.(dist.μ .- μ_old)) - Σ_diff = maximum(abs.(dist.Σ .- Σ_old)) + # Convergence check (allocation-free) + μ_diff = zero(T) + for j in 1:d + μ_diff = max(μ_diff, abs(dist.μ[j] - μ_old[j])) + end + Σ_diff = zero(T) + for j in 1:d, k in 1:d + Σ_diff = max(Σ_diff, abs(dist.Σ[j, k] - Σ_old[j, k])) + end ν_diff = abs(dist.ν - ν_old) if μ_diff < tol && Σ_diff < tol && (fix_nu || ν_diff < tol) @@ -451,7 +441,6 @@ function StatsAPI.fit!( isempty(obs_seq) && return dist - # Check dimensions d = dist.dim for (i, obs) in enumerate(obs_seq) length(obs) == d || throw( @@ -459,61 +448,74 @@ function StatsAPI.fit!( ) end - # Filter out zero-weight observations - nonzero_idx = findall(w -> w > 0, weight_seq) - isempty(nonzero_idx) && return dist - - obs_vec = [obs_seq[i] for i in nonzero_idx] - weights = [weight_seq[i] for i in nonzero_idx] - n = length(obs_vec) + weight_sum = zero(eltype(weight_seq)) + for w in weight_seq + w > 0 && (weight_sum += w) + end + iszero(weight_sum) && return dist - # Normalize weights - weight_sum = sum(weights) - weights ./= weight_sum + T = eltype(dist.μ) + n = length(obs_seq) - # EM algorithm + # Allocate working arrays once per fit! call + posterior_weights = Vector{T}(undef, n) + σ²_acc = zeros(T, d) + μ_acc = zeros(T, d) μ_old = copy(dist.μ) σ²_old = copy(dist.σ²) ν_old = dist.ν - for iter in 1:max_iter - # E-step: Compute posterior weights - posterior_weights = zeros(n) - - for i in 1:n - diff = obs_vec[i] .- dist.μ - mahal² = sum(diff[j]^2 / dist.σ²[j] for j in 1:d) - posterior_weights[i] = (dist.ν + d) / (dist.ν + mahal²) + for _ in 1:max_iter + # E-step + μ M-step accumulation + fill!(μ_acc, zero(T)) + weight_post_sum = zero(T) + + for i in eachindex(obs_seq) + w = weight_seq[i] + if w > 0 + mahal² = zero(T) + for j in 1:d + dist._diff[j] = obs_seq[i][j] - dist.μ[j] + mahal² += dist._diff[j]^2 / dist.σ²[j] + end + pw = T((dist.ν + d) / (dist.ν + mahal²)) + posterior_weights[i] = pw + wp = T(w / weight_sum) * pw + for j in 1:d + μ_acc[j] += wp * obs_seq[i][j] + end + weight_post_sum += wp + else + posterior_weights[i] = zero(T) + end end - # M-step: Update parameters - # Update μ - weighted_sum = sum(weights[i] * posterior_weights[i] * obs_vec[i] for i in 1:n) - weight_post_sum = sum(weights[i] * posterior_weights[i] for i in 1:n) - dist.μ .= weighted_sum ./ weight_post_sum + dist.μ .= μ_acc ./ weight_post_sum - # Update σ² (diagonal variances) - σ²_new = zeros(d) - for j in 1:d - for i in 1:n - diff_j = obs_vec[i][j] - dist.μ[j] - σ²_new[j] += weights[i] * posterior_weights[i] * diff_j^2 + # σ² M-step (second pass with updated μ) + fill!(σ²_acc, zero(T)) + for i in eachindex(obs_seq) + weight_seq[i] > 0 || continue + wp = T(weight_seq[i] / weight_sum) * posterior_weights[i] + for j in 1:d + diff_j = obs_seq[i][j] - dist.μ[j] + σ²_acc[j] += wp * diff_j^2 end - σ²_new[j] /= sum(weights) - σ²_new[j] = max(σ²_new[j], 1e-8) # Ensure positivity end - - dist.σ² .= σ²_new + for j in 1:d + dist.σ²[j] = max(σ²_acc[j], T(1e-8)) + end dist.logdetΣ = sum(log, dist.σ²) - # Update ν (if not fixed) + # ν M-step if !fix_nu - avg_log_u_minus_u = sum( - weights[i] * (log(posterior_weights[i]) - posterior_weights[i]) for i in 1:n - ) + avg_log_u_minus_u = zero(T) + for i in eachindex(obs_seq) + weight_seq[i] > 0 || continue + pw = posterior_weights[i] + avg_log_u_minus_u += T(weight_seq[i] / weight_sum) * (log(pw) - pw) + end - # Optimize over log(ν) to ensure ν > 0 - # Let x = log(ν), so ν = exp(x) function objective(x::Vector) ν_val = exp(x[1]) f = @@ -522,7 +524,7 @@ function StatsAPI.fit!( 1 + avg_log_u_minus_u + digamma((ν_val + d) / 2) - log((ν_val + d) / 2) - return f^2 # Minimize squared residual + return f^2 end function gradient!(G, x::Vector) @@ -533,12 +535,10 @@ function StatsAPI.fit!( 1 + avg_log_u_minus_u + digamma((ν_val + d) / 2) - log((ν_val + d) / 2) - # df/dν df_dν = -0.5 * trigamma(ν_val / 2) + 1 / ν_val + 0.5 * trigamma((ν_val + d) / 2) - 1 / (ν_val + d) - # d(f²)/dx = d(f²)/dν * dν/dx = 2f * df/dν * ν (since dν/dx = ν) return G[1] = 2 * f * df_dν * ν_val end @@ -550,30 +550,32 @@ function StatsAPI.fit!( 1 + avg_log_u_minus_u + digamma((ν_val + d) / 2) - log((ν_val + d) / 2) - # df/dν df_dν = -0.5 * trigamma(ν_val / 2) + 1 / ν_val + 0.5 * trigamma((ν_val + d) / 2) - 1 / (ν_val + d) - # d²f/dν² d2f_dν2 = -0.25 * polygamma(2, ν_val / 2) - 1 / ν_val^2 + 0.25 * polygamma(2, (ν_val + d) / 2) + 1 / (ν_val + d)^2 - # d²(f²)/dx² using chain rule: d²(f²)/dx² = [2(df/dν)² + 2f*d²f/dν²]*ν² + 2f*df/dν*ν return H[1, 1] = (2 * df_dν^2 + 2 * f * d2f_dν2) * ν_val^2 + 2 * f * df_dν * ν_val end - # Use Newton's method with analytical derivatives td = TwiceDifferentiable(objective, gradient!, hessian!, [log(dist.ν)]) result = optimize(td, [log(dist.ν)], Newton()) dist.ν = exp(Optim.minimizer(result)[1]) end - # Check convergence - μ_diff = maximum(abs.(dist.μ .- μ_old)) - σ²_diff = maximum(abs.(dist.σ² .- σ²_old)) + # Convergence check + μ_diff = zero(T) + for j in 1:d + μ_diff = max(μ_diff, abs(dist.μ[j] - μ_old[j])) + end + σ²_diff = zero(T) + for j in 1:d + σ²_diff = max(σ²_diff, abs(dist.σ²[j] - σ²_old[j])) + end ν_diff = abs(dist.ν - ν_old) if μ_diff < tol && σ²_diff < tol && (fix_nu || ν_diff < tol) From 6e5a193395316457ef921584ffdff7f11015a11b Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sat, 2 May 2026 00:13:31 -0400 Subject: [PATCH 03/20] Add Bernoulli/Poisson GLMs and priors Introduce an AbstractPrior hierarchy (NoPrior, RidgePrior) with neglogprior, neglogprior_grad!, and neglogprior_hess! helpers. Add BernoulliGLM and PoissonGLM types with DensityInterface, Random.rand, and StatsAPI.fit! implementations that minimize the weighted negative log-posterior via Optim (Newton) and compose with priors. Extend GaussianGLM to carry a prior and apply the prior Hessian in weighted least squares (Ridge support). Export new symbols and pull in extra LogExpFunctions helpers (log1pexp, logistic). Add comprehensive tests for priors, BernoulliGLM, and PoissonGLM and include the new test file in runtests.jl. --- src/EmissionModels.jl | 6 +- src/glms/glm.jl | 335 +++++++++++++++++++++++++++-- test/glm/gaussian.jl | 24 +++ test/glm/test_bernoulli_poisson.jl | 273 +++++++++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 624 insertions(+), 15 deletions(-) create mode 100644 test/glm/test_bernoulli_poisson.jl diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index 49e05dd..58159c2 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -3,7 +3,7 @@ module EmissionModels using Distributions: Normal, Bernoulli, Poisson, Chisq using DensityInterface using LinearAlgebra -using LogExpFunctions: logaddexp, logsumexp +using LogExpFunctions: logaddexp, logsumexp, log1pexp, logistic using Optim: optimize, TwiceDifferentiable, Newton, LBFGS, LineSearches using Optim using Random @@ -19,6 +19,8 @@ include("glms/glm.jl") export rand, logdensityof, fit! export PoissonZeroInflated export MultivariateT, MultivariateTDiag -export GaussianGLM +export GaussianGLM, BernoulliGLM, PoissonGLM +export AbstractPrior, NoPrior, RidgePrior +export neglogprior, neglogprior_grad!, neglogprior_hess! end diff --git a/src/glms/glm.jl b/src/glms/glm.jl index a24381a..d977cfb 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -4,40 +4,349 @@ Abstract type for Generalized Linear Model emission distributions. GLM subtypes should implement the HiddenMarkovModels.jl interface: -- `DensityInterface.DensityKind(::YourGLM)` - Declare HasDensity() -- `DensityInterface.logdensityof(glm, obs)` - Compute log density -- `Random.rand(rng, glm, X)` - Generate samples given covariates -- `StatsAPI.fit!(glm, obs_seq, weight_seq)` - Update parameters +- `DensityInterface.DensityKind(::YourGLM)` → `HasDensity()` +- `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 """ abstract type AbstractGLM end -mutable struct GaussianGLM{T<:Real} <: AbstractGLM - β::Vector{T} # Coefficients - σ2::T # Variance +""" + AbstractPrior + +Supertype for priors on GLM coefficients β. + +Subtypes must implement: +- `neglogprior(prior, β)` → scalar negative log-prior (up to constant) +- `neglogprior_grad!(prior, g, β)` → accumulate ∂(-log p(β))/∂β into `g` +- `neglogprior_hess!(prior, H, β)` → accumulate ∂²(-log p(β))/∂β² into `H` + +The gradient and Hessian methods accumulate (+=) rather than overwrite so +multiple priors can compose additively. +""" +abstract type AbstractPrior end + +""" + NoPrior <: AbstractPrior + +Flat (improper uniform) prior — no regularization. +""" +struct NoPrior <: AbstractPrior end + +neglogprior(::NoPrior, β) = zero(eltype(β)) +neglogprior_grad!(::NoPrior, g, β) = nothing +neglogprior_hess!(::NoPrior, H, β) = nothing + +""" + RidgePrior{T<:Real} <: AbstractPrior + +Isotropic Gaussian prior β ~ N(0, (1/λ)I), equivalent to L2 regularization. +Contributes 0.5 λ ‖β‖² to the negative log-posterior. +""" +struct RidgePrior{T<:Real} <: AbstractPrior + λ::T +end + +neglogprior(p::RidgePrior, β) = oftype(p.λ, 0.5) * p.λ * dot(β, β) + +function neglogprior_grad!(p::RidgePrior, g, β) + for i in eachindex(g, β) + g[i] += p.λ * β[i] + end end +function neglogprior_hess!(p::RidgePrior, H, β) + for i in eachindex(β) + H[i, i] += p.λ + end +end + +""" + GaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + +Linear regression emission with Gaussian noise. + +E[Y|x] = xᵀβ, Var[Y|x] = σ². `fit!` uses weighted least squares; a +`RidgePrior(λ)` augments the normal equations with `λI`, which is the +exact MAP solution for a Gaussian prior β ~ N(0, (1/λ)I). + +# Fields +- `β`: coefficient vector (length p) +- `σ2`: noise variance +- `prior`: regularization prior (default `NoPrior()`) +""" +mutable struct GaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + β::Vector{T} + σ2::T + prior::P +end + +GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:Real} = + GaussianGLM{T, NoPrior}(Vector{T}(β), σ2, NoPrior()) + +GaussianGLM(β::AbstractVector{T}, σ2::T, prior::P) where {T<:Real, P<:AbstractPrior} = + GaussianGLM{T, P}(Vector{T}(β), σ2, prior) + +# Convenience: promote β eltype and σ2 together +GaussianGLM(β::AbstractVector, σ2::Real) = + GaussianGLM(float.(β), float(σ2)) + +GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior) = + GaussianGLM(float.(β), float(σ2), prior) + DensityInterface.DensityKind(::GaussianGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof(reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real}) μ = dot(reg.β, control_seq) - return -0.5 * log(2π * reg.σ2) - 0.5 * ((y-μ)^2 / reg.σ2) + return -0.5 * log(2π * reg.σ2) - 0.5 * ((y - μ)^2 / reg.σ2) end function Random.rand(rng::AbstractRNG, reg::GaussianGLM; control_seq::AbstractVector{<:Real}) return rand(rng, Normal(dot(reg.β, control_seq), sqrt(reg.σ2))) end -function StatsAPI.fit!(reg::GaussianGLM, obs_seq::AbstractVector{<:Real}, weights::AbstractVector{<:Real}; - control_seq::AbstractMatrix{<:Real}) - # Fit coefficients using weighted least squares +function StatsAPI.fit!( + reg::GaussianGLM, + obs_seq::AbstractVector{<:Real}, + weights::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, +) W = Diagonal(weights) XWX = control_seq' * W * control_seq XWy = control_seq' * W * obs_seq + + #= For RidgePrior(λ), the MAP normal equations are (X'WX + λI)β = X'Wy. + neglogprior_hess! accumulates λI into XWX in-place, so other prior + types compose here automatically. =# + neglogprior_hess!(reg.prior, XWX, reg.β) + reg.β = XWX \ XWy - # Update variance residuals = obs_seq .- control_seq * reg.β reg.σ2 = sum(weights .* (residuals .^ 2)) / sum(weights) return reg -end \ No newline at end of file +end + +""" + BernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + +Logistic-regression emission for binary (0/1) observations. + +P(Y=1|x) = σ(xᵀβ) where σ is the logistic function. `fit!` minimizes the +weighted negative log-posterior via Optim Newton, so any `AbstractPrior` +composes without changes to the solver. + +# Fields +- `β`: coefficient vector (length p) +- `prior`: regularization prior (default `NoPrior()`) +""" +mutable struct BernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + β::Vector{T} + prior::P +end + +BernoulliGLM(β::AbstractVector{T}) where {T<:Real} = + BernoulliGLM{T, NoPrior}(Vector{T}(β), NoPrior()) + +BernoulliGLM(β::AbstractVector{T}, prior::P) where {T<:Real, P<:AbstractPrior} = + BernoulliGLM{T, P}(Vector{T}(β), prior) + +DensityInterface.DensityKind(::BernoulliGLM) = DensityInterface.HasDensity() + +function DensityInterface.logdensityof( + glm::BernoulliGLM, + y::Integer; + control_seq::AbstractVector{<:Real}, +) + (y == 0 || y == 1) || return oftype(dot(glm.β, control_seq), -Inf) + η = dot(glm.β, control_seq) + return y == 1 ? -log1pexp(-η) : -log1pexp(η) +end + +function Random.rand(rng::AbstractRNG, glm::BernoulliGLM; control_seq::AbstractVector{<:Real}) + return rand(rng, Bernoulli(logistic(dot(glm.β, control_seq)))) +end + +""" + fit!(glm::BernoulliGLM, obs_seq, weight_seq; control_seq, optim_opts) + +Minimize the weighted negative log-posterior via Optim Newton. + +The objective is Σᵢ wᵢ · ℓ(β; yᵢ, xᵢ) + neglogprior(prior, β) where ℓ is the +Bernoulli log-likelihood. Gradient and Hessian are computed analytically; the +Hessian is the standard logistic regression Fisher information weighted by wᵢ, +plus the prior's Hessian contribution. +""" +function StatsAPI.fit!( + glm::BernoulliGLM{T}, + obs_seq::AbstractVector, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, + optim_opts::Optim.Options=Optim.Options(), +) where {T} + n, p = size(control_seq) + length(obs_seq) == n || + throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || + throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(glm.β) == p || + throw(DimensionMismatch("β length $(length(glm.β)) ≠ control_seq columns $p")) + + prior = glm.prior # avoid closure over glm to keep captures minimal + + function neglogpost(β) + nll = zero(T) + for i in 1:n + x_i = view(control_seq, i, :) + η_i = dot(β, x_i) + #= log1pexp is numerically safe for both signs of η =# + nll += weight_seq[i] * (obs_seq[i] == 1 ? log1pexp(-η_i) : log1pexp(η_i)) + end + return nll + neglogprior(prior, β) + end + + function neglogpost_grad!(g, β) + fill!(g, zero(T)) + for i in 1:n + x_i = view(control_seq, i, :) + r_i = weight_seq[i] * (logistic(dot(β, x_i)) - obs_seq[i]) + for j in 1:p + g[j] += r_i * x_i[j] + end + end + neglogprior_grad!(prior, g, β) + end + + function neglogpost_hess!(H, β) + fill!(H, zero(T)) + for i in 1:n + x_i = view(control_seq, i, :) + μ_i = logistic(dot(β, x_i)) + W_i = weight_seq[i] * μ_i * (one(T) - μ_i) + for j in 1:p + for k in 1:p + H[j, k] += W_i * x_i[j] * x_i[k] + end + end + end + neglogprior_hess!(prior, H, β) + end + + td = TwiceDifferentiable(neglogpost, neglogpost_grad!, neglogpost_hess!, glm.β) + result = optimize(td, glm.β, Newton(), optim_opts) + copyto!(glm.β, Optim.minimizer(result)) + + return glm +end + +""" + PoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + +Log-linear (Poisson) regression emission for count observations. + +E[Y|x] = exp(xᵀβ). `fit!` minimizes the weighted negative log-posterior via +Optim Newton. Any `AbstractPrior` composes without changes to the solver. + +# Fields +- `β`: coefficient vector (length p) +- `prior`: regularization prior (default `NoPrior()`) +""" +mutable struct PoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + β::Vector{T} + prior::P +end + +PoissonGLM(β::AbstractVector{T}) where {T<:Real} = + PoissonGLM{T, NoPrior}(Vector{T}(β), NoPrior()) + +PoissonGLM(β::AbstractVector{T}, prior::P) where {T<:Real, P<:AbstractPrior} = + PoissonGLM{T, P}(Vector{T}(β), prior) + +DensityInterface.DensityKind(::PoissonGLM) = DensityInterface.HasDensity() + +function DensityInterface.logdensityof( + glm::PoissonGLM, + y::Integer; + control_seq::AbstractVector{<:Real}, +) + y >= 0 || return oftype(dot(glm.β, control_seq), -Inf) + η = dot(glm.β, control_seq) + return y * η - exp(η) - logfactorial(y) +end + +function Random.rand(rng::AbstractRNG, glm::PoissonGLM; control_seq::AbstractVector{<:Real}) + return rand(rng, Poisson(exp(dot(glm.β, control_seq)))) +end + +""" + fit!(glm::PoissonGLM, obs_seq, weight_seq; control_seq, optim_opts) + +Minimize the weighted negative log-posterior via Optim Newton. + +The objective is Σᵢ wᵢ · ℓ(β; yᵢ, xᵢ) + neglogprior(prior, β) where ℓ is +the Poisson log-likelihood. The Hessian is the Fisher information weighted by +wᵢ (Hessian of Poisson negloglik = Σ wᵢμᵢ xᵢxᵢᵀ), plus the prior contribution. +""" +function StatsAPI.fit!( + glm::PoissonGLM{T}, + obs_seq::AbstractVector, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, + optim_opts::Optim.Options=Optim.Options(), +) where {T} + n, p = size(control_seq) + length(obs_seq) == n || + throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || + throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(glm.β) == p || + throw(DimensionMismatch("β length $(length(glm.β)) ≠ control_seq columns $p")) + + prior = glm.prior + + function neglogpost(β) + nll = zero(T) + for i in 1:n + x_i = view(control_seq, i, :) + η_i = dot(β, x_i) + μ_i = exp(η_i) + nll += weight_seq[i] * (μ_i - obs_seq[i] * η_i) + end + return nll + neglogprior(prior, β) + end + + function neglogpost_grad!(g, β) + fill!(g, zero(T)) + for i in 1:n + x_i = view(control_seq, i, :) + η_i = clamp(dot(β, x_i), -500, 500) + r_i = weight_seq[i] * (exp(η_i) - obs_seq[i]) + for j in 1:p + g[j] += r_i * x_i[j] + end + end + neglogprior_grad!(prior, g, β) + end + + function neglogpost_hess!(H, β) + fill!(H, zero(T)) + for i in 1:n + x_i = view(control_seq, i, :) + η_i = clamp(dot(β, x_i), -500, 500) + W_i = weight_seq[i] * exp(η_i) + for j in 1:p + for k in 1:p + H[j, k] += W_i * x_i[j] * x_i[k] + end + end + end + neglogprior_hess!(prior, H, β) + end + + td = TwiceDifferentiable(neglogpost, neglogpost_grad!, neglogpost_hess!, glm.β) + result = optimize(td, glm.β, Newton(), optim_opts) + copyto!(glm.β, Optim.minimizer(result)) + + return glm +end diff --git a/test/glm/gaussian.jl b/test/glm/gaussian.jl index 641576f..f0fb42d 100644 --- a/test/glm/gaussian.jl +++ b/test/glm/gaussian.jl @@ -108,4 +108,28 @@ _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) @test glm.σ2 ≈ σ2_true atol=0.15 end + @testset "Constructor with prior" begin + glm = GaussianGLM([0.0, 0.0], 1.0) + @test glm.prior isa NoPrior + + glm2 = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(2.0)) + @test glm2.prior isa RidgePrior + @test glm2.prior.λ == 2.0 + end + + @testset "fit! with RidgePrior shrinks toward zero" begin + β_true = [3.0, -3.0] + σ2_true = 1.0 + X, y, w = _synthetic_gaussian_glm(rng, 300, 2; β=β_true, σ2=σ2_true) + + glm_noprior = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm_noprior, y, w; control_seq=X) + + glm_ridge = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(10.0)) + fit!(glm_ridge, y, w; control_seq=X) + + @test norm(glm_ridge.β) < norm(glm_noprior.β) + @test all(isfinite, glm_ridge.β) + end + end \ No newline at end of file diff --git a/test/glm/test_bernoulli_poisson.jl b/test/glm/test_bernoulli_poisson.jl new file mode 100644 index 0000000..734a8ee --- /dev/null +++ b/test/glm/test_bernoulli_poisson.jl @@ -0,0 +1,273 @@ +using EmissionModels +using Test +using Random +using DensityInterface +using StatsAPI +using LinearAlgebra +using Distributions: Bernoulli, Poisson, logpdf + + +function _sigmoid(η) + return one(η) / (one(η) + exp(-η)) +end + +function _synthetic_bernoulli(rng, n, β_true; weights=:uniform) + p = length(β_true) + X = hcat(ones(n), randn(rng, n, p - 1)) + η = X * β_true + y = Int[rand(rng) < _sigmoid(η[i]) ? 1 : 0 for i in 1:n] + w = weights === :uniform ? ones(n) : rand(rng, n) .+ 0.5 + return X, y, w +end + +function _synthetic_poisson(rng, n, β_true; weights=:uniform) + p = length(β_true) + X = hcat(ones(n), randn(rng, n, p - 1)) + η = X * β_true + y = Int[rand(rng, Poisson(exp(η[i]))) for i in 1:n] + w = weights === :uniform ? ones(n) : rand(rng, n) .+ 0.5 + return X, y, w +end + +@testset "AbstractPrior" begin + @testset "NoPrior" begin + prior = NoPrior() + β = [1.0, -2.0, 0.5] + g = zeros(3) + H = zeros(3, 3) + + @test neglogprior(prior, β) == 0.0 + neglogprior_grad!(prior, g, β) + @test g == zeros(3) + neglogprior_hess!(prior, H, β) + @test H == zeros(3, 3) + end + + @testset "RidgePrior" begin + prior = RidgePrior(2.0) + β = [1.0, -1.0] + g = zeros(2) + H = zeros(2, 2) + + @test neglogprior(prior, β) ≈ 0.5 * 2.0 * dot(β, β) + + neglogprior_grad!(prior, g, β) + @test g ≈ 2.0 .* β + + neglogprior_hess!(prior, H, β) + @test H ≈ 2.0 * I(2) + end + + @testset "RidgePrior accumulates into existing g and H" begin + prior = RidgePrior(1.0) + β = [2.0, 3.0] + g = [10.0, 20.0] + H = [1.0 0.0; 0.0 1.0] + + neglogprior_grad!(prior, g, β) + @test g ≈ [12.0, 23.0] + + neglogprior_hess!(prior, H, β) + @test H ≈ [2.0 0.0; 0.0 2.0] + end +end + +@testset "BernoulliGLM" begin + rng = Random.MersenneTwister(42) + + @testset "Constructor" begin + glm = BernoulliGLM([0.5, -1.0]) + @test glm.β == [0.5, -1.0] + @test glm.prior isa NoPrior + + glm2 = BernoulliGLM([0.5, -1.0], RidgePrior(1.0)) + @test glm2.prior isa RidgePrior + end + + @testset "DensityInterface" begin + glm = BernoulliGLM([0.0, 1.0]) + @test DensityKind(glm) == HasDensity() + + x = [1.0, 0.5] + η = dot(glm.β, x) + μ = _sigmoid(η) + + @test logdensityof(glm, 1; control_seq=x) ≈ log(μ) rtol=1e-10 + @test logdensityof(glm, 0; control_seq=x) ≈ log(1 - μ) rtol=1e-10 + @test logdensityof(glm, 2; control_seq=x) == -Inf + @test logdensityof(glm, -1; control_seq=x) == -Inf + end + + @testset "logdensityof large η (numerical stability)" begin + glm = BernoulliGLM([50.0]) + x = [1.0] + @test isfinite(logdensityof(glm, 1; control_seq=x)) + @test isfinite(logdensityof(glm, 0; control_seq=x)) + + glm2 = BernoulliGLM([-50.0]) + @test isfinite(logdensityof(glm2, 1; control_seq=x)) + @test isfinite(logdensityof(glm2, 0; control_seq=x)) + end + + @testset "Random sampling" begin + glm = BernoulliGLM([0.0, 2.0]) + x = [1.0, 0.0] + n = 5_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + + @test all(s -> s == 0 || s == 1, samples) + @test mean(samples) ≈ _sigmoid(dot(glm.β, x)) atol=0.05 + end + + @testset "fit! uniform weights" begin + β_true = [0.3, -1.2] + X, y, w = _synthetic_bernoulli(rng, 2000, β_true) + + glm = BernoulliGLM(zeros(2)) + fit!(glm, y, w; control_seq=X) + + @test glm.β ≈ β_true atol=0.15 + @test all(isfinite, glm.β) + end + + @testset "fit! random weights" begin + β_true = [1.0, 0.5] + X, y, w = _synthetic_bernoulli(rng, 2000, β_true; weights=:random) + + glm = BernoulliGLM(zeros(2)) + fit!(glm, y, w; control_seq=X) + + @test glm.β ≈ β_true atol=0.20 + end + + @testset "fit! with RidgePrior shrinks toward zero" begin + β_true = [2.0, -2.0] + X, y, w = _synthetic_bernoulli(rng, 500, β_true) + + glm_noprior = BernoulliGLM(zeros(2)) + fit!(glm_noprior, y, w; control_seq=X) + + glm_ridge = BernoulliGLM(zeros(2), RidgePrior(10.0)) + fit!(glm_ridge, y, w; control_seq=X) + + # Ridge should shrink coefficients toward zero + @test norm(glm_ridge.β) < norm(glm_noprior.β) + @test all(isfinite, glm_ridge.β) + end + + @testset "fit! DimensionMismatch" begin + glm = BernoulliGLM(zeros(2)) + X = ones(5, 2) + @test_throws DimensionMismatch fit!(glm, ones(Int, 4), ones(5); control_seq=X) + @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(4); control_seq=X) + @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(5); control_seq=ones(5, 3)) + end + + @testset "fit! all-zero observations" begin + X = hcat(ones(50), randn(rng, 50)) + y = zeros(Int, 50) + w = ones(50) + + glm = BernoulliGLM(zeros(2)) + fit!(glm, y, w; control_seq=X) + + # All y=0 ⟹ intercept should be strongly negative + @test glm.β[1] < -2.0 + @test all(isfinite, glm.β) + end +end + + +@testset "PoissonGLM" begin + rng = Random.MersenneTwister(99) + + @testset "Constructor" begin + glm = PoissonGLM([1.0, 0.5]) + @test glm.β == [1.0, 0.5] + @test glm.prior isa NoPrior + + glm2 = PoissonGLM([1.0, 0.5], RidgePrior(0.5)) + @test glm2.prior isa RidgePrior + end + + @testset "DensityInterface" begin + glm = PoissonGLM([0.5, 0.2]) + @test DensityKind(glm) == HasDensity() + + x = [1.0, 1.0] + η = dot(glm.β, x) + μ = exp(η) + + for k in 0:5 + @test logdensityof(glm, k; control_seq=x) ≈ logpdf(Poisson(μ), k) rtol=1e-10 + end + + @test logdensityof(glm, -1; control_seq=x) == -Inf + end + + @testset "Random sampling" begin + glm = PoissonGLM([1.5]) + x = [1.0] + n = 5_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + + @test all(s -> s >= 0, samples) + @test mean(samples) ≈ exp(glm.β[1]) atol=0.2 + end + + @testset "fit! uniform weights" begin + β_true = [1.0, 0.4] + X, y, w = _synthetic_poisson(rng, 2000, β_true) + + glm = PoissonGLM(zeros(2)) + fit!(glm, y, w; control_seq=X) + + @test glm.β ≈ β_true atol=0.15 + @test all(isfinite, glm.β) + end + + @testset "fit! random weights" begin + β_true = [0.5, -0.3] + X, y, w = _synthetic_poisson(rng, 2000, β_true; weights=:random) + + glm = PoissonGLM(zeros(2)) + fit!(glm, y, w; control_seq=X) + + @test glm.β ≈ β_true atol=0.15 + end + + @testset "fit! with RidgePrior shrinks toward zero" begin + β_true = [2.0, -1.5] + X, y, w = _synthetic_poisson(rng, 500, β_true) + + glm_noprior = PoissonGLM(zeros(2)) + fit!(glm_noprior, y, w; control_seq=X) + + glm_ridge = PoissonGLM(zeros(2), RidgePrior(10.0)) + fit!(glm_ridge, y, w; control_seq=X) + + @test norm(glm_ridge.β) < norm(glm_noprior.β) + @test all(isfinite, glm_ridge.β) + end + + @testset "fit! DimensionMismatch" begin + glm = PoissonGLM(zeros(2)) + X = ones(5, 2) + @test_throws DimensionMismatch fit!(glm, ones(Int, 4), ones(5); control_seq=X) + @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(4); control_seq=X) + @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(5); control_seq=ones(5, 3)) + end + + @testset "fit! all-zero counts" begin + X = hcat(ones(50), randn(rng, 50)) + y = zeros(Int, 50) + w = ones(50) + + glm = PoissonGLM([0.0, 0.0]) + fit!(glm, y, w; control_seq=X) + + # All y=0 ⟹ intercept (log mean) should be strongly negative + @test glm.β[1] < -1.0 + @test all(isfinite, glm.β) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 1e58b66..55c5e4e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -16,6 +16,7 @@ using StatsAPI @testset "GLM Models" begin include("glm/gaussian.jl") + include("glm/test_bernoulli_poisson.jl") end @testset "Zero-inflated models" begin include("zeroinflated/test_poisson.jl") From 04aa06d85098f205d50584ef04d00ed7e0724fbe Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sat, 2 May 2026 09:24:07 -0400 Subject: [PATCH 04/20] Add the beginnings of documentation --- README.md | 89 ++++++++++++++++++++++++++++++--------- docs/make.jl | 8 +++- docs/src/custom.md | 82 ++++++++++++++++++++++++++++++++++++ docs/src/distributions.md | 83 ++++++++++++++++++++++++++++++++++++ docs/src/glm.md | 53 +++++++++++++++++++++++ docs/src/index.md | 46 ++++++++++++++++---- docs/src/priors.md | 73 ++++++++++++++++++++++++++++++++ 7 files changed, 404 insertions(+), 30 deletions(-) create mode 100644 docs/src/custom.md create mode 100644 docs/src/distributions.md create mode 100644 docs/src/glm.md create mode 100644 docs/src/priors.md diff --git a/README.md b/README.md index 314adda..f055f6b 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -# EmissionModels +# EmissionModels.jl [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://rsenne.github.io/EmissionModels.jl/stable/) [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://rsenne.github.io/EmissionModels.jl/dev/) @@ -8,40 +8,89 @@ [![JET](https://img.shields.io/badge/%F0%9F%9B%A9%EF%B8%8F_tested_with-JET.jl-233f9a)](https://github.com/aviatesk/JET.jl) [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/JuliaDiff/BlueStyle) -## What purpose does this package serve? +A Julia package providing emission models for [HiddenMarkovModels.jl](https://github.com/baggepinn/HiddenMarkovModels.jl). It supplies ready-to-use distributions that describe how observations are generated conditioned on the HMM's latent states. -HiddenMarkovModels.jl is completely *generic*. This means users can have models that emit generic julia objects. This level of expressiveness is allowed through the ability to create custom emission models (i.e., the models that describe how observvations are generated conditioned on the current latent state of an HMM). HiddenMarkovModels.jl expects emission types to implement a small set of methods. +## Quick start ```julia -Random.rand(rng::AbstractRNG, dist::EmissionModel) -DensityInterface.DensityKind(::EmissionModel) = HasDensity() -DensityInterface.logdensityof(dist::EmissionModel, obs) -StatsAPI.fit!(dist::EmissionModel, obs_seq, weight_seq) +using Pkg +Pkg.add("HiddenMarkovModels") +Pkg.add(url="https://github.com/rsenne/EmissionModels.jl") + +using EmissionModels +using HiddenMarkovModels + +# Create an emission model +dist = PoissonZeroInflated(5.0, 0.3) + +# Sample, evaluate densities, or fit to data +x = rand(dist) +logp = logdensityof(dist, x) +fit!(dist, observations, weights) +``` + +## Distribution models + +All types implement the `HiddenMarkovModels` emission interface (`rand`, `logdensityof`, `fit!`). + +### Count data + +| Type | Description | +|------|-------------| +| `PoissonZeroInflated(λ, π)` | Zero-inflated Poisson for excess zeros in count data. | + +### Multivariate continuous + +| Type | Description | +|------|-------------| +| `MultivariateT(μ, Σ, ν)` | Full-covariance multivariate Student's t. | +| `MultivariateTDiag(μ, σ², ν)` | Diagonal-covariance multivariate Student's t. | + +### GLM emissions (observation depends on a control vector) + +| Type | Description | +|------|-------------| +| `GaussianGLM(β, σ²)` | Linear regression with Gaussian noise. | +| `BernoulliGLM(β)` | Logistic regression for binary data. | +| `PoissonGLM(β)` | Log-linear regression for count data. | + +GLM types support regularization via priors: + +```julia +using EmissionModels: RidgePrior + +β = zeros(3) +glm = GaussianGLM(β, 1.0, RidgePrior(0.5)) # L2 regularization ``` -This package supplies many models that already implement those methods and are thus ready-to-use with HiddenMarkovModels.jl. +Each GLM is fit via `fit!(glm, y, w; control_seq=X)`, where `control_seq` (design matrix `X`) maps latent states to the regression covariates. + +## Creating custom emission models + +`HiddenMarkovModels.jl` accepts any type that implements the following interface: + +```julia +Random.rand(rng::AbstractRNG, dist::MyEmission) +DensityInterface.DensityKind(::MyEmission) # return HasDensity() +DensityInterface.logdensityof(dist::MyEmission, obs) +StatsAPI.fit!(dist::MyEmission, obs_seq, weight_seq) +``` + +See the [documentation](https://rsenne.github.io/EmissionModels.jl/dev/) for details. ## Installation -This package is not yet registered on the Julia REPL. To add directly from GitHub (latest main): +EmissionModels.jl is not yet registered. Install from GitHub: ```julia +using Pkg Pkg.add(url="https://github.com/rsenne/EmissionModels.jl") ``` ## Contributing -Contributions are very welcome. Suggested ways to help: - -- Open an issue for bugs or feature requests. -- Submit a pull request with tests when adding features or fixing bugs. -- Improve examples and documentation under `docs/`. - -When contributing, please follow the repository coding style and add tests for new behavior. +Contributions are welcome. Please follow the [Julia Blue Style](https://github.com/JuliaDiff/BlueStyle) and add tests for new behavior. Pull requests and issues are appreciated. ## License -This project is licensed under the terms in the `LICENSE` file in the repository root. - ---- -*Thank you for using EmissionModels.jl — feedback and contributions appreciated.* +EmissionModels.jl is licensed under the terms of the `LICENSE` file. diff --git a/docs/make.jl b/docs/make.jl index 6b78852..d6cac3d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -12,7 +12,13 @@ makedocs(; edit_link="main", assets=String[], ), - pages=["Home" => "index.md"], + pages=[ + "Home" => "index.md", + "Distributions" => "distributions.md", + "GLM Emissions" => "glm.md", + "Priors" => "priors.md", + "Custom Emission Models" => "custom.md", + ], ) deploydocs(; repo="github.com/rsenne/EmissionModels.jl", devbranch="main") diff --git a/docs/src/custom.md b/docs/src/custom.md new file mode 100644 index 0000000..5f79e0e --- /dev/null +++ b/docs/src/custom.md @@ -0,0 +1,82 @@ +# Custom Emission Models + +To create a custom emission model that works with `HiddenMarkovModels.jl`, your type must implement the following methods. + +## Required interface + +### 1. Sample from the emission + +```julia +Random.rand([rng::AbstractRNG,] dist::MyEmission) +Random.rand([rng::AbstractRNG,] dist::MyEmission, n::Integer) # n samples +``` + +The first form returns a single sample. The second form (optional) samples `n` i.i.d. draws. + +### 2. Declare the density interface + +```julia +DensityInterface.DensityKind(::MyEmission) = DensityInterface.HasDensity() +``` + +Return `HasDensity()` if `logdensityof` is implemented; otherwise return `NotAdjointDensity()`. + +### 3. Evaluate log-density / log-probability-mass + +```julia +DensityInterface.logdensityof(dist::MyEmission, obs) +``` + +This is called internally during the forward-backward algorithm. + +### 4. Parameter estimation + +```julia +StatsAPI.fit!(dist::MyEmission, obs_seq::AbstractVector, weight_seq::AbstractVector) +``` + +Optionally accept keyword arguments: + +```julia +StatsAPI.fit!(dist::MyEmission, obs_seq, weight_seq; max_iter=100, tol=1e-6) +``` + +## Optional: Conditional emissions (with control vectors) + +If your model depends on a design vector, add a `control` keyword to `logdensityof` and `rand`: + +```julia +DensityInterface.logdensityof(dist::MyEmission, obs; control=nothing) +Random.rand([rng,] dist::MyEmission; control=nothing) +``` + +## Example: Bernoulli emission + +```julia +using Random, DensityInterface, StatsAPI + +mutable struct MyBernoulli{T<:Real} + p::T # probability of success +end + +DensityInterface.DensityKind(::MyBernoulli) = DensityInterface.HasDensity() + +function DensityInterface.logdensityof(b::MyBernoulli, obs::Integer) + obs ∈ (0, 1) || return oftype(b.p, -Inf) + return obs * log(b.p) + (one(obs) - obs) * log(one(b.p) - b.p) +end + +function Random.rand(rng::AbstractRNG, b::MyBernoulli) + rand(rng) < b.p ? 1 : 0 +end + +# Fitting: MLE for Bernoulli is the sample mean +function StatsAPI.fit!(b::MyBernoulli, obs_seq, weight_seq) + total_w = sum(weight_seq) + b.p = dot(weight_seq, obs_seq) / total_w + return b +end + +# Use in an HMM: +emission = MyBernoulli(0.5) +``` diff --git a/docs/src/distributions.md b/docs/src/distributions.md new file mode 100644 index 0000000..6bbca5c --- /dev/null +++ b/docs/src/distributions.md @@ -0,0 +1,83 @@ +# Distributions + +This page documents the probability distributions provided by EmissionModels.jl. Each type implements the `HiddenMarkovModels` emission interface: + +```julia +Random.rand(rng::AbstractRNG, dist::MyEmission) +DensityInterface.DensityKind(::MyEmission) = DensityInterface.HasDensity() +DensityInterface.logdensityof(dist::MyEmission, obs) +StatsAPI.fit!(dist::MyEmission, obs_seq, weight_seq) +``` + +## Count data + +### `PoissonZeroInflated(λ, π)` + +A zero-inflated Poisson distribution for count observations with excess zeros. + +| Field | Type | Description | +|-------|------|-------------| +| `λ` | `T` | Rate parameter of the Poisson component (``λ > 0``). | +| `π` | `T` | Probability of structural (extra) zeros (``0 ≤ π ≤ 1``). | + +**Model:** ``P(X = k) = π + (1-π)\exp(-λ)`` for ``k = 0``, and ``P(X = k) = (1-π) \cdot \frac{λ^k e^{-λ}}{k!}`` for ``k > 0``. + +**Example:** + +```julia +dist = PoissonZeroInflated(5.0, 0.3) +x = rand(dist) # sample +lp = logdensityof(dist, 0) # log-density at k = 0 +fit!(dist, [1, 0, 3, 0], [1.0, 1.0, 1.0, 1.0]) +``` + +## Multivariate continuous + +### `MultivariateT(μ, Σ, ν)` + +Full-covariance multivariate Student's t-distribution. + +| Field | Description | +|-------|-------------| +| `μ` | Location vector (mean for ``ν > 1``). | +| `Σ` | Positive-definite scale matrix. | +| `ν` | Degrees of freedom (> 0). | + +**Example:** + +```julia +μ = [0.0, 1.0] +Σ = [1.0 0.5; 0.5 2.0] +dist = MultivariateT(μ, Σ, ν=5.0) +x = rand(dist) +lp = logdensityof(dist, x) +``` + +### `MultivariateTDiag(μ, σ², ν)` + +Diagonal-covariance multivariate Student's t-distribution (more efficient for high dimensions). + +| Field | Description | +|-------|-------------| +| `μ` | Location vector. | +| `σ²` | Vector of diagonal variances (all positive). | +| `ν` | Degrees of freedom (> 0). | + +**Example:** + +```julia +dist = MultivariateTDiag([0.0, 0.0], [1.0, 2.0], ν=5.0) +``` + +## Parameter estimation (`fit!`) + +All multivariate t-distributions are fit via a weighted EM algorithm: + +```julia +fit!(dist, obs_seq, weight_seq; max_iter=100, tol=1e-6, fix_nu=false) +``` + +- The E-step posterior weights depend on the current Mahalanobis distance. +- The M-step updates ``μ`` and ``Σ`` (and optionally ``ν``). +- Pass `fix_nu=true` to keep the degrees of freedom fixed during fitting. +- `μ` and ``Σ`` are re-regularized on-the-fly if the Cholesky factorization fails. diff --git a/docs/src/glm.md b/docs/src/glm.md new file mode 100644 index 0000000..9ad53c4 --- /dev/null +++ b/docs/src/glm.md @@ -0,0 +1,53 @@ +# GLM Emissions + +Generalized linear model (GLM) emission types model observations that depend on an external **control vector** (also called a design or feature vector ``x``). The parameter ``β`` maps the control vector to the latent linear predictor, and a link function transforms this into the emission parameters. + +## Provided GLM types + +### `GaussianGLM(β, σ²)` + +Linear regression with Gaussian noise. + +``f(xᵀβ) = xᵀβ`` (identity link) + +``Y \mid x \sim \mathcal{N}(xᵀβ, \ σ²)`` + +```julia +glm = GaussianGLM(zeros(3), 1.0) +# or with a prior: +glm = GaussianGLM(zeros(3), 1.0, RidgePrior(0.5)) +``` + +### `BernoulliGLM(β)` + +Logistic regression for binary observations ``y ∈ {0, 1}``. + +``f(xᵀβ) = \frac{1}{1 + \exp(-xᵀβ)}`` (logistic / sigmoid) + +``Y \mid x \sim \text{Bernoulli}(f(xᵀβ))`` + +```julia +glm = BernoulliGLM(zeros(3)) +``` + +### `PoissonGLM(β)` + +Log-linear regression for count observations ``Y ∈ \{0, 1, \ldots\}``. + +``f(xᵀβ) = \exp(xᵀβ)`` + +``Y \mid x \sim \text{Poisson}(\exp(xᵀβ))`` + +```julia +glm = PoissonGLM(zeros(3)) +``` + +## Fitting GLM emissions + +GLM emissions are fit by minimizing the negative log-posterior (or weighted negative log-likelihood when no prior is used): + +``\ell(β) = -\sum_{i=1}^{n} w_i \, \log p(y_i \mid x_i, β) + \text{neglogprior}(β)`` + +```julia +fit!(glm, y_seq, w_seq; control_seq=X, optim_opts=Optim.Options()) +``` diff --git a/docs/src/index.md b/docs/src/index.md index de480bc..0001a3d 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,14 +1,42 @@ -```@meta -CurrentModule = EmissionModels -``` - # EmissionModels -Documentation for [EmissionModels](https://github.com/rsenne/EmissionModels.jl). +## What is EmissionModels? -```@index -``` +EmissionModels.jl is a Julia package that provides ready-to-use **emission models** for [HiddenMarkovModels.jl](https://github.com/baggepinn/HiddenMarkovModels.jl). An *emission model* (or *emission distribution*) describes the probabilistic relationship between the latent state of a hidden Markov model and its observable output. + +This package supplies a collection of well-tested distributions — from count and multivariate continuous models to generalized linear models (GLMs) — so you can get started with HMM modeling without writing custom emission code from scratch. + +## Quick start + +```julia +using EmissionModels +using HiddenMarkovModels + +# --- A simple count emission model --- +emission = PoissonZeroInflated(5.0, 0.3) -```@autodocs -Modules = [EmissionModels] +# Sample +x = rand(emission) + +# Evaluate log-density +lp = logdensityof(emission, x) + +# --- Use in an HMM (example) --- +# Transition matrix (2 states → 2 states) +T = [0.8 0.2; 0.3 0.7] + +# Two different emission models for the two states +emissions = [PoissonZeroInflated(5.0, 0.3), PoissonZeroInflated(20.0, 0.1)] + +hmm = HiddenMarkModel(T, emissions) + +# Fit to observed sequence (uses Viterbi / forward-backward under the hood) +fit!(hmm, observations) ``` + +## Where to go from here + +- **[Distributions](distributions.md)** — Browse all available count, multivariate, and GLM emission models. +- **[GLM Emissions](glm.md)** — Learn about regression-based emissions that take a control (design) vector. +- **[Priors](priors.md)** — Add regularization to GLM emissions. +- **[Custom Emission Models](custom.md)** — Build your own emission types and plug them into the HMM framework. diff --git a/docs/src/priors.md b/docs/src/priors.md new file mode 100644 index 0000000..458916d --- /dev/null +++ b/docs/src/priors.md @@ -0,0 +1,73 @@ +# Priors + +## Available priors + +### `RidgePrior(λ)` + +Implements L2 (ridge) penalization ``\frac{λ}{2} \|\beta\|_2^2``. + +```julia +using EmissionModels: RidgePrior + +β = zeros(3) +glm = GaussianGLM(β, 1.0, RidgePrior(0.5)) # ridge with λ = 0.5 +``` + +## Implementing a custom prior + +Every prior must conform to the following interface: + +```julia +# Required methods + +"""Return +neglog prior value (no penalty => 0).""" +neglogprior(prior::MyPrior, β) + +"""Return + neglog prior gradient (in-place update to `g`).""" +neglogprior_grad!(prior::MyPrior, g, β) + +"""Return + neglog prior hessian (in-place update to `H`).""" +neglogprior_hess!(prior::MyPrior, H, β) +``` + +### Example: Lasso prior (L1 penalty) + +L1 penalties are non-differentiable at 0, so they only participate in the gradient. + +```julia +struct LassoPrior{T} + λ::T +end + +neglogprior(p::LassoPrior, β) = p.λ * sum(abs, β) + +function neglogprior_grad!(p::LassoPrior, g, β) + @inbounds for j in eachindex(β) + g[j] += p.λ * sign(β[j]) + end +end + +function neglogprior_hess!(p::LassoPrior, H, β) + # Lasso is not twice differentiable at β == 0. + # Use the subdifferential (zero everywhere except β == 0). + @inbounds for j in eachindex(β) + if β[j] == 0 + H[j, j] += 1e6 # effectively ∞ at 0 + end + end +end + +# Usage: +glm = GaussianGLM(zeros(3), 1.0, LassoPrior(0.5)) +``` + +## Composing priors + +Priors can be composed with `ComposedPrior`: + +```julia +prior = ComposedPrior(RidgePrior(0.5), LassoPrior(0.1)) +glm = GaussianGLM(zeros(5), 1.0, prior) +``` + +The composed prior simply sums the individual penalties, gradients, and Hessians. From 5120189a4c0fafde8d46c4290d7cdcb6e18da8b7 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Fri, 8 May 2026 22:44:45 -0400 Subject: [PATCH 05/20] Add Mv support, reduce allocatiosn greatly. Expand tests to allocations --- src/EmissionModels.jl | 1 + src/glms/glm.jl | 1015 ++++++++++++++++++++++++---- src/multivariate/t.jl | 17 +- test/allocations.jl | 192 ++++++ test/glm/gaussian.jl | 147 +++- test/glm/test_bernoulli_poisson.jl | 259 +++++++ test/runtests.jl | 4 + 7 files changed, 1504 insertions(+), 131 deletions(-) create mode 100644 test/allocations.jl diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index 58159c2..e059dcc 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -20,6 +20,7 @@ export rand, logdensityof, fit! export PoissonZeroInflated export MultivariateT, MultivariateTDiag export GaussianGLM, BernoulliGLM, PoissonGLM +export MvGaussianGLM, MvBernoulliGLM, MvPoissonGLM export AbstractPrior, NoPrior, RidgePrior export neglogprior, neglogprior_grad!, neglogprior_hess! diff --git a/src/glms/glm.jl b/src/glms/glm.jl index d977cfb..4e04498 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -8,6 +8,11 @@ GLM subtypes should implement the HiddenMarkovModels.jl interface: - `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 + +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 @@ -41,7 +46,9 @@ neglogprior_hess!(::NoPrior, H, β) = nothing RidgePrior{T<:Real} <: AbstractPrior Isotropic Gaussian prior β ~ N(0, (1/λ)I), equivalent to L2 regularization. -Contributes 0.5 λ ‖β‖² to the negative log-posterior. +Contributes 0.5 λ ‖β‖² to the negative log-posterior. For multivariate GLMs +the same prior is applied independently to each column of the coefficient +matrix `B`, giving a Frobenius-norm penalty 0.5 λ ‖B‖_F². """ struct RidgePrior{T<:Real} <: AbstractPrior λ::T @@ -106,36 +113,306 @@ function Random.rand(rng::AbstractRNG, reg::GaussianGLM; control_seq::AbstractVe end function StatsAPI.fit!( - reg::GaussianGLM, + reg::GaussianGLM{T}, obs_seq::AbstractVector{<:Real}, weights::AbstractVector{<:Real}; control_seq::AbstractMatrix{<:Real}, -) - W = Diagonal(weights) - XWX = control_seq' * W * control_seq - XWy = control_seq' * W * obs_seq +) where {T<:Real} + n, p = size(control_seq) + length(obs_seq) == n || + throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) + length(weights) == n || + throw(DimensionMismatch("weights length $(length(weights)) ≠ control_seq rows $n")) + length(reg.β) == p || + throw(DimensionMismatch("β length $(length(reg.β)) ≠ control_seq columns $p")) + + XWX = zeros(T, p, p) + XWy = zeros(T, p) + wsum = zero(T) + + for i in 1:n + w = T(weights[i]) + wsum += w + x_i = view(control_seq, i, :) + y_i = T(obs_seq[i]) + for a in 1:p + xa = x_i[a] + wxa = w * xa + for b in 1:p + XWX[a, b] += wxa * x_i[b] + end + XWy[a] += wxa * y_i + end + end - #= For RidgePrior(λ), the MAP normal equations are (X'WX + λI)β = X'Wy. - neglogprior_hess! accumulates λI into XWX in-place, so other prior - types compose here automatically. =# + #= RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. =# neglogprior_hess!(reg.prior, XWX, reg.β) - reg.β = XWX \ XWy + F = cholesky!(Symmetric(XWX)) + copyto!(reg.β, XWy) + ldiv!(F, reg.β) - residuals = obs_seq .- control_seq * reg.β - reg.σ2 = sum(weights .* (residuals .^ 2)) / sum(weights) + sw_r2 = zero(T) + for i in 1:n + x_i = view(control_seq, i, :) + r_i = T(obs_seq[i]) - dot(reg.β, x_i) + sw_r2 += T(weights[i]) * r_i * r_i + end + reg.σ2 = sw_r2 / wsum return reg end +#= Hand-rolled Newton solver shared by Bernoulli and Poisson GLM fits. + Avoids Optim's TwiceDifferentiable + closure-object allocations: workspace + buffers (g, H, Δ) are passed in by the caller, who allocates them once and + reuses across columns in the multivariate case. =# + +#= Lazy column accessor: indexing into a Vector{Vector{T}} along output dim j + without copying. Used by the multivariate GLM fits to avoid an n-sized + y-buffer per column. =# +struct _ColumnElementView{T,V<:AbstractVector{<:AbstractVector}} <: AbstractVector{T} + seq::V + j::Int +end +_ColumnElementView(seq::V, j::Int) where {V<:AbstractVector{<:AbstractVector}} = + _ColumnElementView{eltype(eltype(V)), V}(seq, j) +Base.size(c::_ColumnElementView) = (length(c.seq),) +Base.IndexStyle(::Type{<:_ColumnElementView}) = IndexLinear() +Base.@propagate_inbounds Base.getindex(c::_ColumnElementView, i::Integer) = c.seq[i][c.j] + +function _bernoulli_loss( + β::AbstractVector{T}, + y::AbstractVector, + w::AbstractVector{<:Real}, + X::AbstractMatrix{<:Real}, + prior::AbstractPrior, +) where {T<:Real} + n, _ = size(X) + nll = zero(T) + for i in 1:n + x_i = view(X, i, :) + η_i = dot(β, x_i) + nll += T(w[i]) * (y[i] == 1 ? log1pexp(-η_i) : log1pexp(η_i)) + end + return nll + neglogprior(prior, β) +end + +function _bernoulli_gh!( + g::AbstractVector{T}, + H::AbstractMatrix{T}, + β::AbstractVector{T}, + y::AbstractVector, + w::AbstractVector{<:Real}, + X::AbstractMatrix{<:Real}, + prior::AbstractPrior, +) where {T<:Real} + n, p = size(X) + fill!(g, zero(T)) + fill!(H, zero(T)) + for i in 1:n + x_i = view(X, i, :) + wi = T(w[i]) + μ_i = logistic(dot(β, x_i)) + r_i = wi * (μ_i - T(y[i])) + W_i = wi * μ_i * (one(T) - μ_i) + for a in 1:p + xa = x_i[a] + g[a] += r_i * xa + wxa = W_i * xa + for b in 1:p + H[a, b] += wxa * x_i[b] + end + end + end + neglogprior_grad!(prior, g, β) + neglogprior_hess!(prior, H, β) + return nothing +end + +function _poisson_loss( + β::AbstractVector{T}, + y::AbstractVector, + w::AbstractVector{<:Real}, + X::AbstractMatrix{<:Real}, + prior::AbstractPrior, +) where {T<:Real} + n, _ = size(X) + nll = zero(T) + for i in 1:n + x_i = view(X, i, :) + η_i = clamp(dot(β, x_i), T(-500), T(500)) + nll += T(w[i]) * (exp(η_i) - T(y[i]) * η_i) + end + return nll + neglogprior(prior, β) +end + +function _poisson_gh!( + g::AbstractVector{T}, + H::AbstractMatrix{T}, + β::AbstractVector{T}, + y::AbstractVector, + w::AbstractVector{<:Real}, + X::AbstractMatrix{<:Real}, + prior::AbstractPrior, +) where {T<:Real} + n, p = size(X) + fill!(g, zero(T)) + fill!(H, zero(T)) + for i in 1:n + x_i = view(X, i, :) + wi = T(w[i]) + η_i = clamp(dot(β, x_i), T(-500), T(500)) + eη = exp(η_i) + r_i = wi * (eη - T(y[i])) + W_i = wi * eη + for a in 1:p + xa = x_i[a] + g[a] += r_i * xa + wxa = W_i * xa + for b in 1:p + H[a, b] += wxa * x_i[b] + end + end + end + neglogprior_grad!(prior, g, β) + neglogprior_hess!(prior, H, β) + return nothing +end + +#= Newton with backtracking line search. `loss` and `gh!` are concrete callables + (one per family). With type-stable inputs the JIT specializes and the inner + loop is allocation-free aside from `cholesky!`'s internal work. =# +function _newton_solve!( + β::AbstractVector{T}, + g::AbstractVector{T}, + H::AbstractMatrix{T}, + Δ::AbstractVector{T}, + loss::F1, + gh!::F2, + y::AbstractVector, + w::AbstractVector{<:Real}, + X::AbstractMatrix{<:Real}, + prior::AbstractPrior; + max_iter::Int=50, + gtol::Real=1e-8, + max_backtrack::Int=20, + ridge::Real=1e-12, +) where {T<:Real, F1, F2} + p = length(β) + f_curr = loss(β, y, w, X, prior) + + for _ in 1:max_iter + gh!(g, H, β, y, w, X, prior) + + gnorm = zero(T) + for j in 1:p + ag = abs(g[j]) + ag > gnorm && (gnorm = ag) + end + gnorm < gtol && break + + # Tiny ridge for numerical PD (Bernoulli/Poisson Hessians are PSD and + # PD when X has full column rank; this guards against rare singular cases). + for j in 1:p + H[j, j] += T(ridge) + end + + F = cholesky!(Symmetric(H, :L)) + copyto!(Δ, g) + ldiv!(F, Δ) + + # Backtracking line search: halve α until the loss decreases. + α = one(T) + success = false + f_new = f_curr + for _ in 0:max_backtrack + for j in 1:p + β[j] -= α * Δ[j] + end + f_new = loss(β, y, w, X, prior) + if isfinite(f_new) && f_new < f_curr + success = true + break + end + for j in 1:p + β[j] += α * Δ[j] + end + α /= 2 + end + success || break + f_curr = f_new + end + return β +end + +#= Internal helpers used by univariate `BernoulliGLM.fit!` / `PoissonGLM.fit!`. + Allocate the Newton workspace (g, H, Δ) — three small vectors / one matrix. =# +function _fit_bernoulli_glm!( + β::Vector{T}, + obs_seq::AbstractVector, + weight_seq::AbstractVector{<:Real}, + control_seq::AbstractMatrix{<:Real}, + prior::AbstractPrior; + max_iter::Int=50, + gtol::Real=1e-8, + max_backtrack::Int=20, +) where {T<:Real} + n, p = size(control_seq) + length(obs_seq) == n || + throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || + throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(β) == p || + throw(DimensionMismatch("β length $(length(β)) ≠ control_seq columns $p")) + + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) + return _newton_solve!( + β, g, H, Δ, _bernoulli_loss, _bernoulli_gh!, + obs_seq, weight_seq, control_seq, prior; + max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + ) +end + +function _fit_poisson_glm!( + β::Vector{T}, + obs_seq::AbstractVector, + weight_seq::AbstractVector{<:Real}, + control_seq::AbstractMatrix{<:Real}, + prior::AbstractPrior; + max_iter::Int=50, + gtol::Real=1e-8, + max_backtrack::Int=20, +) where {T<:Real} + n, p = size(control_seq) + length(obs_seq) == n || + throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || + throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(β) == p || + throw(DimensionMismatch("β length $(length(β)) ≠ control_seq columns $p")) + + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) + return _newton_solve!( + β, g, H, Δ, _poisson_loss, _poisson_gh!, + obs_seq, weight_seq, control_seq, prior; + max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + ) +end + """ BernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM Logistic-regression emission for binary (0/1) observations. P(Y=1|x) = σ(xᵀβ) where σ is the logistic function. `fit!` minimizes the -weighted negative log-posterior via Optim Newton, so any `AbstractPrior` -composes without changes to the solver. +weighted negative log-posterior via a hand-rolled Newton solver with +backtracking line search, so any `AbstractPrior` composes without changes +to the solver. # Fields - `β`: coefficient vector (length p) @@ -169,74 +446,27 @@ function Random.rand(rng::AbstractRNG, glm::BernoulliGLM; control_seq::AbstractV end """ - fit!(glm::BernoulliGLM, obs_seq, weight_seq; control_seq, optim_opts) - -Minimize the weighted negative log-posterior via Optim Newton. + fit!(glm::BernoulliGLM, obs_seq, weight_seq; + control_seq, max_iter=50, gtol=1e-8, max_backtrack=20) -The objective is Σᵢ wᵢ · ℓ(β; yᵢ, xᵢ) + neglogprior(prior, β) where ℓ is the -Bernoulli log-likelihood. Gradient and Hessian are computed analytically; the -Hessian is the standard logistic regression Fisher information weighted by wᵢ, -plus the prior's Hessian contribution. +Minimize the weighted negative log-posterior via a hand-rolled Newton solver +with backtracking line search. The objective is +`Σᵢ wᵢ · ℓ(β; yᵢ, xᵢ) + neglogprior(prior, β)` where `ℓ` is the Bernoulli +log-likelihood; gradient and Hessian are analytic. """ function StatsAPI.fit!( - glm::BernoulliGLM{T}, + glm::BernoulliGLM, obs_seq::AbstractVector, weight_seq::AbstractVector{<:Real}; control_seq::AbstractMatrix{<:Real}, - optim_opts::Optim.Options=Optim.Options(), -) where {T} - n, p = size(control_seq) - length(obs_seq) == n || - throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) - length(glm.β) == p || - throw(DimensionMismatch("β length $(length(glm.β)) ≠ control_seq columns $p")) - - prior = glm.prior # avoid closure over glm to keep captures minimal - - function neglogpost(β) - nll = zero(T) - for i in 1:n - x_i = view(control_seq, i, :) - η_i = dot(β, x_i) - #= log1pexp is numerically safe for both signs of η =# - nll += weight_seq[i] * (obs_seq[i] == 1 ? log1pexp(-η_i) : log1pexp(η_i)) - end - return nll + neglogprior(prior, β) - end - - function neglogpost_grad!(g, β) - fill!(g, zero(T)) - for i in 1:n - x_i = view(control_seq, i, :) - r_i = weight_seq[i] * (logistic(dot(β, x_i)) - obs_seq[i]) - for j in 1:p - g[j] += r_i * x_i[j] - end - end - neglogprior_grad!(prior, g, β) - end - - function neglogpost_hess!(H, β) - fill!(H, zero(T)) - for i in 1:n - x_i = view(control_seq, i, :) - μ_i = logistic(dot(β, x_i)) - W_i = weight_seq[i] * μ_i * (one(T) - μ_i) - for j in 1:p - for k in 1:p - H[j, k] += W_i * x_i[j] * x_i[k] - end - end - end - neglogprior_hess!(prior, H, β) - end - - td = TwiceDifferentiable(neglogpost, neglogpost_grad!, neglogpost_hess!, glm.β) - result = optimize(td, glm.β, Newton(), optim_opts) - copyto!(glm.β, Optim.minimizer(result)) - + max_iter::Int=50, + gtol::Real=1e-8, + max_backtrack::Int=20, +) + _fit_bernoulli_glm!( + glm.β, obs_seq, weight_seq, control_seq, glm.prior; + max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + ) return glm end @@ -246,7 +476,8 @@ end Log-linear (Poisson) regression emission for count observations. E[Y|x] = exp(xᵀβ). `fit!` minimizes the weighted negative log-posterior via -Optim Newton. Any `AbstractPrior` composes without changes to the solver. +a hand-rolled Newton solver with backtracking line search. Any `AbstractPrior` +composes without changes to the solver. # Fields - `β`: coefficient vector (length p) @@ -280,73 +511,611 @@ function Random.rand(rng::AbstractRNG, glm::PoissonGLM; control_seq::AbstractVec end """ - fit!(glm::PoissonGLM, obs_seq, weight_seq; control_seq, optim_opts) - -Minimize the weighted negative log-posterior via Optim Newton. + fit!(glm::PoissonGLM, obs_seq, weight_seq; + control_seq, max_iter=50, gtol=1e-8, max_backtrack=20) -The objective is Σᵢ wᵢ · ℓ(β; yᵢ, xᵢ) + neglogprior(prior, β) where ℓ is -the Poisson log-likelihood. The Hessian is the Fisher information weighted by -wᵢ (Hessian of Poisson negloglik = Σ wᵢμᵢ xᵢxᵢᵀ), plus the prior contribution. +Minimize the weighted negative log-posterior via a hand-rolled Newton solver +with backtracking line search. The objective is +`Σᵢ wᵢ · ℓ(β; yᵢ, xᵢ) + neglogprior(prior, β)` where `ℓ` is the Poisson +log-likelihood; gradient and Hessian are analytic. """ function StatsAPI.fit!( - glm::PoissonGLM{T}, + glm::PoissonGLM, obs_seq::AbstractVector, weight_seq::AbstractVector{<:Real}; control_seq::AbstractMatrix{<:Real}, - optim_opts::Optim.Options=Optim.Options(), -) where {T} + max_iter::Int=50, + gtol::Real=1e-8, + max_backtrack::Int=20, +) + _fit_poisson_glm!( + glm.β, obs_seq, weight_seq, control_seq, glm.prior; + max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + ) + return glm +end + +""" + MvGaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + +Multivariate linear regression emission with shared full covariance. + +For input x ∈ ℝᵖ the conditional distribution of y ∈ ℝᵏ is +`y | x ~ MvNormal(Bᵀ x, Σ)`, where `B` is `p × k` and `Σ` is `k × k`. + +`fit!` is the closed-form weighted multivariate least squares update: +`B = (XᵀWX + λI) \\ XᵀWY` and `Σ = (Y − XB)ᵀ W (Y − XB) / Σwᵢ`. A +`RidgePrior(λ)` adds `λI` to `XᵀWX`, equivalent to a per-column Gaussian +prior on `B` (Frobenius-norm penalty). + +# Fields +- `B`: coefficient matrix of size `p × k` +- `Σ`: residual covariance of size `k × k` +- `prior`: regularization prior on `B` (default `NoPrior()`) +""" +#= Cholesky uplo: we always store the LOWER triangle. `getproperty(::Cholesky, :L)` + only returns a wrapper without copying when uplo == 'L'; with the upper-stored + default, every `.L` access does `copy(factors')`, costing ~k² floats per call. =# + +mutable struct MvGaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + B::Matrix{T} + Σ::Matrix{T} + prior::P + + Σ_chol::Cholesky{T, Matrix{T}} + logdetΣ::T + in_dim::Int + out_dim::Int + + #= Scratch for hot paths. Sequential use only — not safe for concurrent + calls on the same instance. HMM E-steps run forward/backward serially + per state, so each emission has its own scratch and this is fine. =# + _diff::Vector{T} + _z::Vector{T} +end + +function MvGaussianGLM( + B::AbstractMatrix{T}, + Σ::AbstractMatrix{T}, + prior::P, +) where {T<:Real, P<:AbstractPrior} + p, k = size(B) + p > 0 || throw(ArgumentError("B must have at least one row")) + k > 0 || throw(ArgumentError("B must have at least one column")) + size(Σ) == (k, k) || + throw(DimensionMismatch("Σ must be $(k)×$(k) for B with $(k) columns, got $(size(Σ))")) + + Σ_chol = try + cholesky(Symmetric(Σ, :L)) + catch + throw(ArgumentError("Σ must be positive definite")) + end + + return MvGaussianGLM{T, P}( + Matrix{T}(B), Matrix{T}(Σ), prior, + Σ_chol, logdet(Σ_chol), p, k, + Vector{T}(undef, k), Vector{T}(undef, k), + ) +end + +MvGaussianGLM(B::AbstractMatrix{T}, Σ::AbstractMatrix{T}) where {T<:Real} = + MvGaussianGLM(B, Σ, NoPrior()) + +function MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix) + T = promote_type(eltype(B), eltype(Σ)) + Tf = float(T) + return MvGaussianGLM(Matrix{Tf}(B), Matrix{Tf}(Σ), NoPrior()) +end + +function MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix, prior::AbstractPrior) + T = promote_type(eltype(B), eltype(Σ)) + Tf = float(T) + return MvGaussianGLM(Matrix{Tf}(B), Matrix{Tf}(Σ), prior) +end + +DensityInterface.DensityKind(::MvGaussianGLM) = DensityInterface.HasDensity() + +""" + logdensityof(glm::MvGaussianGLM, y::AbstractVector; control_seq) + +Log density of `y ∈ ℝᵏ` under the conditional MvNormal model. Zero allocation: +the residual buffer `_diff` is pre-allocated in the struct. +""" +function DensityInterface.logdensityof( + glm::MvGaussianGLM, + y::AbstractVector; + control_seq::AbstractVector{<:Real}, +) + length(y) == glm.out_dim || + throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) + length(control_seq) == glm.in_dim || + throw(DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", + )) + + k = glm.out_dim + p = glm.in_dim + T = eltype(glm.B) + diff = glm._diff + + #= μ = Bᵀ x and diff = y − μ, fused as a single loop to avoid the + Adjoint*Vector temporary that BLAS would otherwise allocate. =# + for j in 1:k + sj = zero(T) + for r in 1:p + sj += glm.B[r, j] * control_seq[r] + end + diff[j] = T(y[j]) - sj + end + ldiv!(glm.Σ_chol.L, diff) + mahal² = zero(T) + for j in 1:k + mahal² += diff[j] * diff[j] + end + + return -k / 2 * log(2π) - glm.logdetΣ / 2 - mahal² / 2 +end + +function Random.rand( + rng::AbstractRNG, + glm::MvGaussianGLM{T}; + control_seq::AbstractVector{<:Real}, +) where {T<:Real} + out = Vector{T}(undef, glm.out_dim) + rand!(rng, glm, out; control_seq=control_seq) + return out +end + +""" + rand!(rng, glm::MvGaussianGLM, out; control_seq) + +In-place sample. `out` must be a length-`out_dim` `AbstractVector{T}`. +Zero allocation. +""" +function Random.rand!( + rng::AbstractRNG, + glm::MvGaussianGLM{T}, + out::AbstractVector; + control_seq::AbstractVector{<:Real}, +) where {T<:Real} + length(out) == glm.out_dim || + throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) + length(control_seq) == glm.in_dim || + throw(DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", + )) + + k = glm.out_dim + p = glm.in_dim + z = glm._z + randn!(rng, z) + + #= out = Bᵀ x (manual loop avoids the Adjoint*Vector temporary) =# + for j in 1:k + sj = zero(T) + for r in 1:p + sj += glm.B[r, j] * control_seq[r] + end + out[j] = sj + end + #= out += L z via in-place 5-arg mul! =# + mul!(out, glm.Σ_chol.L, z, true, true) + return out +end + +""" + fit!(glm::MvGaussianGLM, obs_seq, weight_seq; control_seq) + +Closed-form weighted multivariate WLS update. Each observation `obs_seq[i]` +must be a length-`k` vector. `RidgePrior(λ)` augments the normal equations +with `λI`. +""" +function StatsAPI.fit!( + glm::MvGaussianGLM{T}, + obs_seq::AbstractVector{<:AbstractVector}, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, +) where {T<:Real} n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) length(weight_seq) == n || throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) - length(glm.β) == p || - throw(DimensionMismatch("β length $(length(glm.β)) ≠ control_seq columns $p")) - - prior = glm.prior - - function neglogpost(β) - nll = zero(T) - for i in 1:n - x_i = view(control_seq, i, :) - η_i = dot(β, x_i) - μ_i = exp(η_i) - nll += weight_seq[i] * (μ_i - obs_seq[i] * η_i) + p == glm.in_dim || + throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) + + k = glm.out_dim + + 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. =# + for i in 1:n + obs_i = obs_seq[i] + length(obs_i) == k || + throw(DimensionMismatch( + "obs_seq[$i] length $(length(obs_i)) ≠ out_dim $k", + )) + w = T(weight_seq[i]) + wsum += w + x_i = view(control_seq, i, :) + for a in 1:p + xa = x_i[a] + wxa = w * xa + for b in 1:p + XWX[a, b] += wxa * x_i[b] + end + for j in 1:k + XWY[a, j] += wxa * T(obs_i[j]) + end end - return nll + neglogprior(prior, β) end - function neglogpost_grad!(g, β) - fill!(g, zero(T)) - for i in 1:n - x_i = view(control_seq, i, :) - η_i = clamp(dot(β, x_i), -500, 500) - r_i = weight_seq[i] * (exp(η_i) - obs_seq[i]) - for j in 1:p - g[j] += r_i * x_i[j] + #= Per-column ridge: λI added to XᵀWX gives B = (XᵀWX + λI) \ XᵀWY, + the joint MAP for independent N(0, (1/λ)I) priors on each column of B. + Pass a length-p view so RidgePrior loops over rows of XWX, not p*k. =# + neglogprior_hess!(glm.prior, XWX, view(glm.B, :, 1)) + + F = try + cholesky(Symmetric(XWX, :L)) + catch + min_eig = minimum(eigvals(Hermitian(XWX))) + XWX .+= (abs(min_eig) + T(1e-6)) * I + cholesky(Symmetric(XWX, :L)) + end + copyto!(glm.B, XWY) + ldiv!(F, glm.B) + + #= Σ M-step: Σ = Σᵢ wᵢ (yᵢ - Bᵀxᵢ)(yᵢ - Bᵀxᵢ)ᵀ / Σwᵢ. Computes residuals + on the fly into a length-k buffer, then accumulates outer products into + Σ_new. Avoids the full n×k residual matrix. =# + r = Vector{T}(undef, k) + Σ_new = zeros(T, k, k) + for i in 1:n + obs_i = obs_seq[i] + w = T(weight_seq[i]) + x_i = view(control_seq, i, :) + for j in 1:k + rj = T(obs_i[j]) + for a in 1:p + rj -= glm.B[a, j] * x_i[a] end + r[j] = rj end - neglogprior_grad!(prior, g, β) - end - - function neglogpost_hess!(H, β) - fill!(H, zero(T)) - for i in 1:n - x_i = view(control_seq, i, :) - η_i = clamp(dot(β, x_i), -500, 500) - W_i = weight_seq[i] * exp(η_i) - for j in 1:p - for k in 1:p - H[j, k] += W_i * x_i[j] * x_i[k] + for j in 1:k + wrj = w * r[j] + for l in 1:j + v = wrj * r[l] + Σ_new[l, j] += v + if l != j + Σ_new[j, l] += v end end end - neglogprior_hess!(prior, H, β) end + Σ_new ./= wsum + + Σ_chol_new = try + cholesky(Symmetric(Σ_new, :L)) + catch + min_eig = minimum(eigvals(Hermitian(Σ_new))) + Σ_new .+= (abs(min_eig) + T(1e-6)) * I + cholesky(Symmetric(Σ_new, :L)) + end + + copyto!(glm.Σ, Σ_new) + glm.Σ_chol = Σ_chol_new + glm.logdetΣ = logdet(glm.Σ_chol) + + return glm +end + +""" + MvBernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + +Multivariate Bernoulli emission as `k` independent logistic regressions. + +For input x ∈ ℝᵖ the conditional distribution of y ∈ {0,1}ᵏ factorizes as +`P(y|x) = ∏ⱼ Bernoulli(σ(B[:,j]ᵀ x))`. Coefficients live in a `p × k` matrix +`B`. `fit!` runs an independent weighted Newton fit per column. + +# Fields +- `B`: coefficient matrix of size `p × k` +- `prior`: regularization prior applied per column (default `NoPrior()`) +""" +mutable struct MvBernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + B::Matrix{T} + prior::P + in_dim::Int + out_dim::Int +end + +function MvBernoulliGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real, P<:AbstractPrior} + p, k = size(B) + p > 0 || throw(ArgumentError("B must have at least one row")) + k > 0 || throw(ArgumentError("B must have at least one column")) + return MvBernoulliGLM{T, P}(Matrix{T}(B), prior, p, k) +end + +MvBernoulliGLM(B::AbstractMatrix{T}) where {T<:Real} = MvBernoulliGLM(B, NoPrior()) + +MvBernoulliGLM(B::AbstractMatrix) = MvBernoulliGLM(float.(B), NoPrior()) + +MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior) = + MvBernoulliGLM(float.(B), prior) + +DensityInterface.DensityKind(::MvBernoulliGLM) = DensityInterface.HasDensity() + +function DensityInterface.logdensityof( + glm::MvBernoulliGLM, + y::AbstractVector; + control_seq::AbstractVector{<:Real}, +) + length(y) == glm.out_dim || + throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) + length(control_seq) == glm.in_dim || + throw(DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", + )) + + Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) + lp = zero(Tη) + for j in 1:glm.out_dim + yj = y[j] + (yj == 0 || yj == 1) || return Tη(-Inf) + η = zero(Tη) + for r in 1:glm.in_dim + η += glm.B[r, j] * control_seq[r] + end + lp += yj == 1 ? -log1pexp(-η) : -log1pexp(η) + end + return lp +end + +function Random.rand( + rng::AbstractRNG, + glm::MvBernoulliGLM; + control_seq::AbstractVector{<:Real}, +) + out = Vector{Int}(undef, glm.out_dim) + rand!(rng, glm, out; control_seq=control_seq) + return out +end + +""" + rand!(rng, glm::MvBernoulliGLM, out; control_seq) + +In-place sample. `out` must be a length-`out_dim` integer vector. +Zero allocation. +""" +function Random.rand!( + rng::AbstractRNG, + glm::MvBernoulliGLM, + out::AbstractVector; + control_seq::AbstractVector{<:Real}, +) + length(out) == glm.out_dim || + throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) + length(control_seq) == glm.in_dim || + throw(DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", + )) + + T = eltype(glm.B) + for j in 1:glm.out_dim + η = zero(T) + for r in 1:glm.in_dim + η += glm.B[r, j] * control_seq[r] + end + out[j] = rand(rng) < logistic(η) ? 1 : 0 + end + return out +end - td = TwiceDifferentiable(neglogpost, neglogpost_grad!, neglogpost_hess!, glm.β) - result = optimize(td, glm.β, Newton(), optim_opts) - copyto!(glm.β, Optim.minimizer(result)) +""" + fit!(glm::MvBernoulliGLM, obs_seq, weight_seq; + control_seq, max_iter=50, gtol=1e-8, max_backtrack=20) +Fit each output dimension independently via weighted Newton. Each observation +`obs_seq[i]` must be a length-`k` vector of 0/1 values. Newton workspace +(`g`, `H`, `Δ`) is allocated once and shared across columns. +""" +function StatsAPI.fit!( + glm::MvBernoulliGLM{T}, + obs_seq::AbstractVector{<:AbstractVector}, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, + max_iter::Int=50, + gtol::Real=1e-8, + max_backtrack::Int=20, +) where {T<:Real} + n, p = size(control_seq) + length(obs_seq) == n || + throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || + throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + p == glm.in_dim || + throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) + + k = glm.out_dim + for i in 1:n + length(obs_seq[i]) == k || + throw(DimensionMismatch( + "obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k", + )) + end + + β_buf = Vector{T}(undef, p) + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) + + for j in 1:k + yview = _ColumnElementView(obs_seq, j) + copyto!(β_buf, view(glm.B, :, j)) + _newton_solve!( + β_buf, g, H, Δ, _bernoulli_loss, _bernoulli_gh!, + yview, weight_seq, control_seq, glm.prior; + max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + ) + for r in 1:p + glm.B[r, j] = β_buf[r] + end + end + return glm +end + +""" + MvPoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + +Multivariate Poisson emission as `k` independent log-linear regressions. + +For input x ∈ ℝᵖ the conditional distribution of y ∈ ℤ₊ᵏ factorizes as +`P(y|x) = ∏ⱼ Poisson(exp(B[:,j]ᵀ x))`. Coefficients live in a `p × k` matrix +`B`. `fit!` runs an independent weighted Newton fit per column. + +# Fields +- `B`: coefficient matrix of size `p × k` +- `prior`: regularization prior applied per column (default `NoPrior()`) +""" +mutable struct MvPoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM + B::Matrix{T} + prior::P + in_dim::Int + out_dim::Int +end + +function MvPoissonGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real, P<:AbstractPrior} + p, k = size(B) + p > 0 || throw(ArgumentError("B must have at least one row")) + k > 0 || throw(ArgumentError("B must have at least one column")) + return MvPoissonGLM{T, P}(Matrix{T}(B), prior, p, k) +end + +MvPoissonGLM(B::AbstractMatrix{T}) where {T<:Real} = MvPoissonGLM(B, NoPrior()) + +MvPoissonGLM(B::AbstractMatrix) = MvPoissonGLM(float.(B), NoPrior()) + +MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior) = + MvPoissonGLM(float.(B), prior) + +DensityInterface.DensityKind(::MvPoissonGLM) = DensityInterface.HasDensity() + +function DensityInterface.logdensityof( + glm::MvPoissonGLM, + y::AbstractVector; + control_seq::AbstractVector{<:Real}, +) + length(y) == glm.out_dim || + throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) + length(control_seq) == glm.in_dim || + throw(DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", + )) + + Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) + lp = zero(Tη) + for j in 1:glm.out_dim + yj = y[j] + yj >= 0 || return Tη(-Inf) + η = zero(Tη) + for r in 1:glm.in_dim + η += glm.B[r, j] * control_seq[r] + end + lp += yj * η - exp(η) - logfactorial(yj) + end + return lp +end + +function Random.rand( + rng::AbstractRNG, + glm::MvPoissonGLM; + control_seq::AbstractVector{<:Real}, +) + out = Vector{Int}(undef, glm.out_dim) + rand!(rng, glm, out; control_seq=control_seq) + return out +end + +""" + rand!(rng, glm::MvPoissonGLM, out; control_seq) + +In-place sample. `out` must be a length-`out_dim` integer vector. +""" +function Random.rand!( + rng::AbstractRNG, + glm::MvPoissonGLM, + out::AbstractVector; + control_seq::AbstractVector{<:Real}, +) + length(out) == glm.out_dim || + throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) + length(control_seq) == glm.in_dim || + throw(DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", + )) + + T = eltype(glm.B) + for j in 1:glm.out_dim + η = zero(T) + for r in 1:glm.in_dim + η += glm.B[r, j] * control_seq[r] + end + out[j] = rand(rng, Poisson(exp(η))) + end + return out +end + +""" + fit!(glm::MvPoissonGLM, obs_seq, weight_seq; + control_seq, max_iter=50, gtol=1e-8, max_backtrack=20) + +Fit each output dimension independently via weighted Newton. Each observation +`obs_seq[i]` must be a length-`k` vector of non-negative integers. Newton +workspace (`g`, `H`, `Δ`) is allocated once and shared across columns. +""" +function StatsAPI.fit!( + glm::MvPoissonGLM{T}, + obs_seq::AbstractVector{<:AbstractVector}, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, + max_iter::Int=50, + gtol::Real=1e-8, + max_backtrack::Int=20, +) where {T<:Real} + n, p = size(control_seq) + length(obs_seq) == n || + throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || + throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + p == glm.in_dim || + throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) + + k = glm.out_dim + for i in 1:n + length(obs_seq[i]) == k || + throw(DimensionMismatch( + "obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k", + )) + end + + β_buf = Vector{T}(undef, p) + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) + + for j in 1:k + yview = _ColumnElementView(obs_seq, j) + copyto!(β_buf, view(glm.B, :, j)) + _newton_solve!( + β_buf, g, H, Δ, _poisson_loss, _poisson_gh!, + yview, weight_seq, control_seq, glm.prior; + max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + ) + for r in 1:p + glm.B[r, j] = β_buf[r] + end + end return glm end diff --git a/src/multivariate/t.jl b/src/multivariate/t.jl index 3acdd7b..7a4f069 100644 --- a/src/multivariate/t.jl +++ b/src/multivariate/t.jl @@ -46,7 +46,7 @@ mutable struct MultivariateT{T<:Real} throw(DimensionMismatch("Σ must be $(dim)×$(dim), got $(size(Σ))")) Σ_chol = try - cholesky(Σ) + cholesky(Symmetric(Σ, :L)) catch throw(ArgumentError("Σ must be positive definite")) end @@ -147,11 +147,16 @@ function DensityInterface.logdensityof(dist::MultivariateT, x::AbstractVector) d = dist.dim ν = dist.ν + diff = dist._diff - # 1 allocation. ldiv! solves L\(x-μ) in-place, avoiding a second allocation. - diff = x - dist.μ + #= Zero allocation: reuse the struct scratch buffer for the residual. + Sequential use only — not safe for concurrent calls on the same dist. =# + @. diff = x - dist.μ ldiv!(dist.Σ_chol.L, diff) - mahal² = sum(abs2, diff) + mahal² = zero(eltype(diff)) + for i in eachindex(diff) + mahal² += diff[i] * diff[i] + end log_norm = loggamma((ν + d) / 2) - loggamma(ν / 2) - (d / 2) * log(ν * π) - dist.logdetΣ / 2 @@ -316,11 +321,11 @@ function StatsAPI.fit!( # Update Σ and Cholesky, regularizing if needed Σ_chol_new = try - cholesky(Symmetric(Σ_acc)) + cholesky(Symmetric(Σ_acc, :L)) catch min_eig = minimum(eigvals(Hermitian(Σ_acc))) Σ_acc .+= (abs(min_eig) + 1e-6) * I - cholesky(Symmetric(Σ_acc)) + cholesky(Symmetric(Σ_acc, :L)) end dist.Σ .= Σ_acc dist.Σ_chol = Σ_chol_new diff --git a/test/allocations.jl b/test/allocations.jl new file mode 100644 index 0000000..6db8fd4 --- /dev/null +++ b/test/allocations.jl @@ -0,0 +1,192 @@ +using EmissionModels +using Test +using Random +using LinearAlgebra +using DensityInterface +using StatsAPI + +#= + Allocation regression tests. The pattern + is: warm up by running the operation once, then call it many times in a + type-stable function and divide @allocated by the rep count to get + per-call bytes. + + Steady-state targets (per call): + - All logdensityof: 0 bytes + - All rand!: 0 bytes (rand allocates only the return vector) + - fit! is bounded by O(p² + k²), independent of n + + See `bench_*` helpers below for why a top-level `@allocated foo()` would + report misleading kwarg-lowering noise that does NOT exist in real loops. +=# + +bench_logd(d, y, x, n) = (s = 0.0; for _ in 1:n; s += logdensityof(d, y; control_seq=x); end; s) +bench_logd_unctrl(d, y, n) = (s = 0.0; for _ in 1:n; s += logdensityof(d, y); end; s) +bench_rand_scalar(rng, d, x, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d; control_seq=x); end; s) +bench_rand_int(rng, d, x, n) = (s = 0; for _ in 1:n; s += rand(rng, d; control_seq=x); end; s) +bench_rand!_v(rng, d, out, x, n) = (s = 0.0; for _ in 1:n; rand!(rng, d, out; control_seq=x); s += out[1]; end; s) +bench_rand!_i(rng, d, out, x, n) = (s = 0; for _ in 1:n; rand!(rng, d, out; control_seq=x); s += out[1]; end; s) +bench_rand_unctrl_scalar(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d); end; s) +bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; end; s) + +@testset "Allocations (steady state)" begin + rng = Random.MersenneTwister(0) + REPS = 1000 + + @testset "logdensityof — zero alloc per call" begin + x = [1.0, 2.0] + + g = GaussianGLM([0.5, -1.0], 1.0) + b = BernoulliGLM([0.5, -1.0]) + p = PoissonGLM([0.5, -1.0]) + mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) + mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) + mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) + yv = [0.1, 0.2]; yi = [0, 1] + + # Warm + bench_logd(g, 0.5, x, 1); bench_logd(b, 1, x, 1); bench_logd(p, 2, x, 1) + bench_logd(mg, yv, x, 1); bench_logd(mb, yi, x, 1); bench_logd(mp, yi, x, 1) + + @test (@allocated bench_logd(g, 0.5, x, REPS)) == 0 + @test (@allocated bench_logd(b, 1, x, REPS)) == 0 + @test (@allocated bench_logd(p, 2, x, REPS)) == 0 + @test (@allocated bench_logd(mg, yv, x, REPS)) == 0 + @test (@allocated bench_logd(mb, yi, x, REPS)) == 0 + @test (@allocated bench_logd(mp, yi, x, REPS)) == 0 + end + + @testset "rand!/rand — zero alloc beyond return" begin + x = [1.0, 2.0] + + g = GaussianGLM([0.5, -1.0], 1.0) + b = BernoulliGLM([0.5, -1.0]) + p = PoissonGLM([0.5, -1.0]) + mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) + mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) + mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) + + # Univariate rand returns a scalar — zero alloc + bench_rand_scalar(rng, g, x, 1) + bench_rand_int(rng, b, x, 1) + bench_rand_int(rng, p, x, 1) + @test (@allocated bench_rand_scalar(rng, g, x, REPS)) == 0 + @test (@allocated bench_rand_int(rng, b, x, REPS)) == 0 + @test (@allocated bench_rand_int(rng, p, x, REPS)) == 0 + + # Multivariate rand! into pre-allocated buffer — zero alloc + out_f = zeros(2) + out_i = zeros(Int, 2) + bench_rand!_v(rng, mg, out_f, x, 1) + bench_rand!_i(rng, mb, out_i, x, 1) + bench_rand!_i(rng, mp, out_i, x, 1) + @test (@allocated bench_rand!_v(rng, mg, out_f, x, REPS)) == 0 + @test (@allocated bench_rand!_i(rng, mb, out_i, x, REPS)) == 0 + @test (@allocated bench_rand!_i(rng, mp, out_i, x, REPS)) == 0 + end + + @testset "GLM fit! — bounded, independent of n" begin + n = 500 + X = hcat(ones(n), randn(rng, n)) + w = ones(n) + + # GaussianGLM: closed-form WLS. Workspace: XWX (p²) + XWy (p). + yg = randn(rng, n) + gg = GaussianGLM([0.0, 0.0], 1.0); fit!(gg, yg, w; control_seq=X) + gg = GaussianGLM([0.0, 0.0], 1.0) + @test (@allocated fit!(gg, yg, w; control_seq=X)) ≤ 1_000 + + # MvGaussianGLM: closed-form weighted MvN regression. + # Workspace: XWX (p²) + XWY (p·k) + Σ_new (k²) + r (k). + ymv = [randn(rng, 2) for _ in 1:n] + gmv = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) + fit!(gmv, ymv, w; control_seq=X) + gmv = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) + @test (@allocated fit!(gmv, ymv, w; control_seq=X)) ≤ 2_000 + + # BernoulliGLM/PoissonGLM: hand-rolled Newton. + # Workspace: g, H, Δ — three small alloc, total ~300 bytes for p=2. + yb = Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:n] + gb = BernoulliGLM(zeros(2)); fit!(gb, yb, w; control_seq=X) + gb = BernoulliGLM(zeros(2)) + @test (@allocated fit!(gb, yb, w; control_seq=X)) ≤ 1_000 + + yp = Int[rand(rng, 0:5) for _ in 1:n] + gp = PoissonGLM(zeros(2)); fit!(gp, yp, w; control_seq=X) + gp = PoissonGLM(zeros(2)) + @test (@allocated fit!(gp, yp, w; control_seq=X)) ≤ 1_000 + + # Multivariate Newton: workspace shared across columns. _ColumnElementView + # avoids the n-sized per-column copy that Vector-of-Vectors would force. + ymb = [Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:2] for _ in 1:n] + gmb = MvBernoulliGLM(zeros(2, 2)); fit!(gmb, ymb, w; control_seq=X) + gmb = MvBernoulliGLM(zeros(2, 2)) + @test (@allocated fit!(gmb, ymb, w; control_seq=X)) ≤ 1_000 + + ymp = [Int[rand(rng, 0:3) for _ in 1:2] for _ in 1:n] + gmp = MvPoissonGLM(zeros(2, 2)); fit!(gmp, ymp, w; control_seq=X) + gmp = MvPoissonGLM(zeros(2, 2)) + @test (@allocated fit!(gmp, ymp, w; control_seq=X)) ≤ 1_000 + end + + @testset "PoissonZeroInflated" begin + zip = PoissonZeroInflated(3.0, 0.2) + + bench_logd_unctrl(zip, 0, 1); bench_logd_unctrl(zip, 5, 1) + @test (@allocated bench_logd_unctrl(zip, 0, REPS)) == 0 + @test (@allocated bench_logd_unctrl(zip, 5, REPS)) == 0 + + bench_rand_unctrl_scalar(rng, zip, 1) + @test (@allocated bench_rand_unctrl_scalar(rng, zip, REPS)) == 0 + + n = 500 + y = [rand(rng, zip) for _ in 1:n] + w = ones(n) + zip2 = PoissonZeroInflated(1.0, 0.1); fit!(zip2, y, w) + zip2 = PoissonZeroInflated(1.0, 0.1) + @test (@allocated fit!(zip2, y, w)) ≤ 1_000 + end + + @testset "MultivariateT (full Σ)" begin + d = 2 + mvt = MultivariateT([0.0, 0.0], [1.0 0.3; 0.3 1.0], 5.0) + xv = [0.1, 0.2] + + # Cholesky now stored as :L and residual reuses struct scratch. + bench_logd_unctrl(mvt, xv, 1) + @test (@allocated bench_logd_unctrl(mvt, xv, REPS)) == 0 + + bench_rand_unctrl_vec(rng, mvt, 1) + # rand still allocates the return vector; bound it loosely. + @test (@allocated bench_rand_unctrl_vec(rng, mvt, REPS)) ≤ 600 * REPS + + n = 500 + obs = [randn(rng, d) for _ in 1:n] + w = ones(n) + mvt2 = MultivariateT([0.0, 0.0], Matrix(1.0I, d, d), 5.0) + fit!(mvt2, obs, w; max_iter=5) + mvt2 = MultivariateT([0.0, 0.0], Matrix(1.0I, d, d), 5.0) + # ν step still goes through Optim — loose bound. + @test (@allocated fit!(mvt2, obs, w; max_iter=5)) ≤ 1_000_000 + end + + @testset "MultivariateTDiag" begin + d = 2 + mvtd = MultivariateTDiag([0.0, 0.0], [1.0, 1.0], 5.0) + xv = [0.1, 0.2] + + bench_logd_unctrl(mvtd, xv, 1) + @test (@allocated bench_logd_unctrl(mvtd, xv, REPS)) == 0 + + bench_rand_unctrl_vec(rng, mvtd, 1) + @test (@allocated bench_rand_unctrl_vec(rng, mvtd, REPS)) ≤ 400 * REPS + + n = 500 + obs = [randn(rng, d) for _ in 1:n] + w = ones(n) + mvtd2 = MultivariateTDiag([0.0, 0.0], [1.0, 1.0], 5.0) + fit!(mvtd2, obs, w; max_iter=5) + mvtd2 = MultivariateTDiag([0.0, 0.0], [1.0, 1.0], 5.0) + @test (@allocated fit!(mvtd2, obs, w; max_iter=5)) ≤ 500_000 + end +end diff --git a/test/glm/gaussian.jl b/test/glm/gaussian.jl index f0fb42d..9fd6499 100644 --- a/test/glm/gaussian.jl +++ b/test/glm/gaussian.jl @@ -4,11 +4,10 @@ using Random using LinearAlgebra using Distributions using DensityInterface -using StableRNGs using StatsAPI using EmissionModels -rng = StableRNG(1234) +rng = Random.MersenneTwister(1234) function _synthetic_gaussian_glm(rng, n::Int, p::Int; β::Vector{Float64}, σ2::Float64, weights=:uniform) @assert length(β) == p @@ -132,4 +131,148 @@ _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) @test all(isfinite, glm_ridge.β) end +end + +function _synthetic_mvgaussian_glm(rng, n::Int, p::Int, k::Int; + B::Matrix{Float64}, Σ::Matrix{Float64}, + weights=:uniform) + @assert size(B) == (p, k) + @assert size(Σ) == (k, k) + X = hcat(ones(n), randn(rng, n, p-1)) + L = cholesky(Σ).L + obs_seq = Vector{Vector{Float64}}(undef, n) + for i in 1:n + μ_i = vec(B' * X[i, :]) + obs_seq[i] = μ_i .+ L * randn(rng, k) + end + w = weights === :uniform ? ones(n) : + weights === :random ? (rand(rng, n) .+ 0.5) : + error("weights must be :uniform or :random") + return X, obs_seq, w +end + +@testset "MvGaussianGLM" begin + rng = Random.MersenneTwister(20260507) + + @testset "Constructor" begin + B = [0.5 -0.3; 1.0 0.2] # p=2, k=2 + Σ = [1.0 0.2; 0.2 0.8] + glm = MvGaussianGLM(B, Σ) + @test glm.B == B + @test glm.Σ == Σ + @test glm.in_dim == 2 + @test glm.out_dim == 2 + @test glm.prior isa NoPrior + + glm_float = MvGaussianGLM([1.0 0.0; 0.0 1.0], [1.0 0.0; 0.0 1.0]) + @test glm_float.B == [1.0 0.0; 0.0 1.0] + @test eltype(glm_float.B) === Float64 + + glm_ridge = MvGaussianGLM(B, Σ, RidgePrior(2.0)) + @test glm_ridge.prior isa RidgePrior + + @test_throws DimensionMismatch MvGaussianGLM(B, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]) + @test_throws ArgumentError MvGaussianGLM(B, [1.0 0.0; 0.0 -1.0]) + end + + @testset "DensityInterface" begin + B = [0.5 -1.0; 1.0 0.5] + Σ = [1.0 0.3; 0.3 1.5] + glm = MvGaussianGLM(B, Σ) + + @test DensityKind(glm) == HasDensity() + + x = [1.0, 2.0] + μ = vec(B' * x) + y = μ .+ [0.1, -0.2] + + # Compare against MvNormal closed form + diff = y - μ + expected = -log(2π) - 0.5 * logdet(Σ) - 0.5 * dot(diff, Σ \ diff) + @test logdensityof(glm, y; control_seq=x) ≈ expected rtol=1e-10 + + # Higher density at the mean than far from it + @test logdensityof(glm, μ; control_seq=x) > + logdensityof(glm, μ .+ 10.0; control_seq=x) + + @test_throws DimensionMismatch logdensityof(glm, [1.0]; control_seq=x) + @test_throws DimensionMismatch logdensityof(glm, y; control_seq=[1.0]) + end + + @testset "Random sampling" begin + B = [0.5 -1.0; 1.0 0.5] + Σ = [1.0 0.3; 0.3 1.5] + glm = MvGaussianGLM(B, Σ) + x = [1.0, 0.5] + + n = 20_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + Y = reduce(hcat, samples)' # n × k + + μ_true = vec(B' * x) + m = vec(mean(Y; dims=1)) + @test m ≈ μ_true atol=0.05 + + S = (Y .- m')' * (Y .- m') / (n - 1) + @test S ≈ Σ atol=0.1 + end + + @testset "fit! recovers B and Σ" begin + B_true = [0.5 -1.0; 1.0 0.5; -0.3 0.8] # p=3, k=2 + Σ_true = [1.0 0.3; 0.3 1.5] + X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 4000, 3, 2; + B=B_true, Σ=Σ_true) + + glm = MvGaussianGLM(zeros(3, 2), Matrix(1.0I, 2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + @test glm.B ≈ B_true atol=0.08 + @test glm.Σ ≈ Σ_true atol=0.10 + @test all(isfinite, glm.B) + @test all(isfinite, glm.Σ) + @test isposdef(glm.Σ) + end + + @testset "fit! random weights" begin + B_true = [1.0 0.0; -0.5 1.5] + Σ_true = [0.5 -0.1; -0.1 0.7] + X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 3000, 2, 2; + B=B_true, Σ=Σ_true, + weights=:random) + + glm = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + @test glm.B ≈ B_true atol=0.10 + @test glm.Σ ≈ Σ_true atol=0.12 + end + + @testset "fit! with RidgePrior shrinks toward zero" begin + B_true = [3.0 -3.0; 3.0 3.0] + Σ_true = [1.0 0.0; 0.0 1.0] + X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 200, 2, 2; + B=B_true, Σ=Σ_true) + + glm_noprior = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) + fit!(glm_noprior, obs_seq, w; control_seq=X) + + glm_ridge = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2), RidgePrior(20.0)) + fit!(glm_ridge, obs_seq, w; control_seq=X) + + @test norm(glm_ridge.B) < norm(glm_noprior.B) + @test all(isfinite, glm_ridge.B) + end + + @testset "fit! DimensionMismatch" begin + glm = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) + X = ones(5, 2) + bad_obs = [zeros(2) for _ in 1:4] + good_obs = [zeros(2) for _ in 1:5] + @test_throws DimensionMismatch fit!(glm, bad_obs, ones(5); control_seq=X) + @test_throws DimensionMismatch fit!(glm, good_obs, ones(4); control_seq=X) + @test_throws DimensionMismatch fit!(glm, good_obs, ones(5); control_seq=ones(5, 3)) + + wrong_dim_obs = [zeros(3) for _ in 1:5] + @test_throws DimensionMismatch fit!(glm, wrong_dim_obs, ones(5); control_seq=X) + end end \ No newline at end of file diff --git a/test/glm/test_bernoulli_poisson.jl b/test/glm/test_bernoulli_poisson.jl index 734a8ee..b30429f 100644 --- a/test/glm/test_bernoulli_poisson.jl +++ b/test/glm/test_bernoulli_poisson.jl @@ -271,3 +271,262 @@ end @test all(isfinite, glm.β) end end + +function _synthetic_mvbernoulli(rng, n, B_true; weights=:uniform) + p, k = size(B_true) + X = hcat(ones(n), randn(rng, n, p - 1)) + obs_seq = Vector{Vector{Int}}(undef, n) + for i in 1:n + obs_seq[i] = Int[rand(rng) < _sigmoid(dot(B_true[:, j], X[i, :])) ? 1 : 0 + for j in 1:k] + end + w = weights === :uniform ? ones(n) : rand(rng, n) .+ 0.5 + return X, obs_seq, w +end + +function _synthetic_mvpoisson(rng, n, B_true; weights=:uniform) + p, k = size(B_true) + X = hcat(ones(n), randn(rng, n, p - 1)) + obs_seq = Vector{Vector{Int}}(undef, n) + for i in 1:n + obs_seq[i] = Int[rand(rng, Poisson(exp(dot(B_true[:, j], X[i, :])))) for j in 1:k] + end + w = weights === :uniform ? ones(n) : rand(rng, n) .+ 0.5 + return X, obs_seq, w +end + +@testset "MvBernoulliGLM" begin + rng = Random.MersenneTwister(123) + + @testset "Constructor" begin + B = [0.5 -1.0; 1.0 0.5] + glm = MvBernoulliGLM(B) + @test glm.B == B + @test glm.in_dim == 2 + @test glm.out_dim == 2 + @test glm.prior isa NoPrior + + glm_float = MvBernoulliGLM([1.0 0.0; 0.0 1.0]) + @test eltype(glm_float.B) === Float64 + + glm_ridge = MvBernoulliGLM(B, RidgePrior(1.0)) + @test glm_ridge.prior isa RidgePrior + end + + @testset "DensityInterface" begin + B = [0.0 1.0; 1.0 -0.5] + glm = MvBernoulliGLM(B) + @test DensityKind(glm) == HasDensity() + + x = [1.0, 0.5] + η1 = dot(B[:, 1], x) + η2 = dot(B[:, 2], x) + μ1, μ2 = _sigmoid(η1), _sigmoid(η2) + + @test logdensityof(glm, [1, 1]; control_seq=x) ≈ log(μ1) + log(μ2) rtol=1e-10 + @test logdensityof(glm, [1, 0]; control_seq=x) ≈ log(μ1) + log(1 - μ2) rtol=1e-10 + @test logdensityof(glm, [0, 1]; control_seq=x) ≈ log(1 - μ1) + log(μ2) rtol=1e-10 + @test logdensityof(glm, [0, 0]; control_seq=x) ≈ log(1 - μ1) + log(1 - μ2) rtol=1e-10 + @test logdensityof(glm, [2, 0]; control_seq=x) == -Inf + end + + @testset "logdensityof large η (numerical stability)" begin + glm = MvBernoulliGLM(reshape([50.0, -50.0], 1, 2)) + x = [1.0] + @test isfinite(logdensityof(glm, [1, 1]; control_seq=x)) + @test isfinite(logdensityof(glm, [0, 0]; control_seq=x)) + end + + @testset "Random sampling" begin + B = [0.0 0.0; 2.0 -2.0] + glm = MvBernoulliGLM(B) + x = [1.0, 0.5] + + n = 5_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + Y = reduce(hcat, samples)' # n × k + + @test all(s -> s == 0 || s == 1, vec(Y)) + @test mean(Y[:, 1]) ≈ _sigmoid(dot(B[:, 1], x)) atol=0.05 + @test mean(Y[:, 2]) ≈ _sigmoid(dot(B[:, 2], x)) atol=0.05 + end + + @testset "fit! recovers B" begin + B_true = [0.3 -1.0; -1.2 0.8] + X, obs_seq, w = _synthetic_mvbernoulli(rng, 5000, B_true) + + glm = MvBernoulliGLM(zeros(2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + @test glm.B ≈ B_true atol=0.20 + @test all(isfinite, glm.B) + end + + @testset "fit! random weights" begin + B_true = [1.0 -0.5; 0.5 1.0] + X, obs_seq, w = _synthetic_mvbernoulli(rng, 5000, B_true; weights=:random) + + glm = MvBernoulliGLM(zeros(2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + @test glm.B ≈ B_true atol=0.25 + end + + @testset "fit! with RidgePrior shrinks toward zero" begin + B_true = [2.0 -2.0; -2.0 2.0] + X, obs_seq, w = _synthetic_mvbernoulli(rng, 500, B_true) + + glm_noprior = MvBernoulliGLM(zeros(2, 2)) + fit!(glm_noprior, obs_seq, w; control_seq=X) + + glm_ridge = MvBernoulliGLM(zeros(2, 2), RidgePrior(10.0)) + fit!(glm_ridge, obs_seq, w; control_seq=X) + + @test norm(glm_ridge.B) < norm(glm_noprior.B) + @test all(isfinite, glm_ridge.B) + end + + @testset "matches independent BernoulliGLM fits" begin + B_true = [0.4 -0.7; -0.9 0.3] + X, obs_seq, w = _synthetic_mvbernoulli(rng, 1000, B_true) + + glm = MvBernoulliGLM(zeros(2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + for j in 1:2 + y_j = [obs_seq[i][j] for i in eachindex(obs_seq)] + glm_j = BernoulliGLM(zeros(2)) + fit!(glm_j, y_j, w; control_seq=X) + @test glm.B[:, j] ≈ glm_j.β rtol=1e-6 + end + end + + @testset "fit! DimensionMismatch" begin + glm = MvBernoulliGLM(zeros(2, 2)) + X = ones(5, 2) + good_obs = [zeros(Int, 2) for _ in 1:5] + @test_throws DimensionMismatch fit!( + glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X, + ) + @test_throws DimensionMismatch fit!(glm, good_obs, ones(4); control_seq=X) + @test_throws DimensionMismatch fit!(glm, good_obs, ones(5); control_seq=ones(5, 3)) + + wrong_dim_obs = [zeros(Int, 3) for _ in 1:5] + @test_throws DimensionMismatch fit!(glm, wrong_dim_obs, ones(5); control_seq=X) + end +end + +@testset "MvPoissonGLM" begin + rng = Random.MersenneTwister(456) + + @testset "Constructor" begin + B = [1.0 0.5; -0.3 0.2] + glm = MvPoissonGLM(B) + @test glm.B == B + @test glm.in_dim == 2 + @test glm.out_dim == 2 + @test glm.prior isa NoPrior + + glm_float = MvPoissonGLM([1.0 0.0; 0.0 1.0]) + @test eltype(glm_float.B) === Float64 + + glm_ridge = MvPoissonGLM(B, RidgePrior(0.5)) + @test glm_ridge.prior isa RidgePrior + end + + @testset "DensityInterface" begin + B = [0.5 0.0; 0.2 -0.5] + glm = MvPoissonGLM(B) + @test DensityKind(glm) == HasDensity() + + x = [1.0, 1.0] + μ1 = exp(dot(B[:, 1], x)) + μ2 = exp(dot(B[:, 2], x)) + + for k1 in 0:3, k2 in 0:3 + expected = logpdf(Poisson(μ1), k1) + logpdf(Poisson(μ2), k2) + @test logdensityof(glm, [k1, k2]; control_seq=x) ≈ expected rtol=1e-10 + end + + @test logdensityof(glm, [-1, 0]; control_seq=x) == -Inf + @test logdensityof(glm, [0, -1]; control_seq=x) == -Inf + end + + @testset "Random sampling" begin + B = reshape([1.5, 0.5], 1, 2) + glm = MvPoissonGLM(B) + x = [1.0] + + n = 5_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + Y = reduce(hcat, samples)' + + @test all(s -> s >= 0, vec(Y)) + @test mean(Y[:, 1]) ≈ exp(B[1, 1]) atol=0.2 + @test mean(Y[:, 2]) ≈ exp(B[1, 2]) atol=0.2 + end + + @testset "fit! recovers B" begin + B_true = [1.0 0.5; 0.4 -0.3] + X, obs_seq, w = _synthetic_mvpoisson(rng, 3000, B_true) + + glm = MvPoissonGLM(zeros(2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + @test glm.B ≈ B_true atol=0.15 + @test all(isfinite, glm.B) + end + + @testset "fit! random weights" begin + B_true = [0.5 0.0; -0.3 0.4] + X, obs_seq, w = _synthetic_mvpoisson(rng, 3000, B_true; weights=:random) + + glm = MvPoissonGLM(zeros(2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + @test glm.B ≈ B_true atol=0.20 + end + + @testset "fit! with RidgePrior shrinks toward zero" begin + B_true = [2.0 -1.5; -1.0 1.5] + X, obs_seq, w = _synthetic_mvpoisson(rng, 500, B_true) + + glm_noprior = MvPoissonGLM(zeros(2, 2)) + fit!(glm_noprior, obs_seq, w; control_seq=X) + + glm_ridge = MvPoissonGLM(zeros(2, 2), RidgePrior(10.0)) + fit!(glm_ridge, obs_seq, w; control_seq=X) + + @test norm(glm_ridge.B) < norm(glm_noprior.B) + @test all(isfinite, glm_ridge.B) + end + + @testset "matches independent PoissonGLM fits" begin + B_true = [0.5 -0.3; 0.4 0.2] + X, obs_seq, w = _synthetic_mvpoisson(rng, 1000, B_true) + + glm = MvPoissonGLM(zeros(2, 2)) + fit!(glm, obs_seq, w; control_seq=X) + + for j in 1:2 + y_j = [obs_seq[i][j] for i in eachindex(obs_seq)] + glm_j = PoissonGLM(zeros(2)) + fit!(glm_j, y_j, w; control_seq=X) + @test glm.B[:, j] ≈ glm_j.β rtol=1e-6 + end + end + + @testset "fit! DimensionMismatch" begin + glm = MvPoissonGLM(zeros(2, 2)) + X = ones(5, 2) + good_obs = [zeros(Int, 2) for _ in 1:5] + @test_throws DimensionMismatch fit!( + glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X, + ) + @test_throws DimensionMismatch fit!(glm, good_obs, ones(4); control_seq=X) + @test_throws DimensionMismatch fit!(glm, good_obs, ones(5); control_seq=ones(5, 3)) + + wrong_dim_obs = [zeros(Int, 3) for _ in 1:5] + @test_throws DimensionMismatch fit!(glm, wrong_dim_obs, ones(5); control_seq=X) + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 55c5e4e..a66ea46 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -26,6 +26,10 @@ using StatsAPI include("multivariate/test_t.jl") end + @testset "Allocations" begin + include("allocations.jl") + end + @testset "Code quality (Aqua.jl)" begin Aqua.test_all(EmissionModels) end From 2ae43af3aa247ae59498ab38f912a2b2dcd55bc7 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sun, 10 May 2026 13:21:00 -0400 Subject: [PATCH 06/20] Fix race condition in MvGaussianGLM and MultivariateT logdensityof --- src/glms/glm.jl | 32 ++++++++++++++------------------ src/multivariate/t.jl | 9 +++++---- test/allocations.jl | 35 +++++++++++++++++++++-------------- 3 files changed, 40 insertions(+), 36 deletions(-) diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 4e04498..4fa8b5d 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -566,12 +566,6 @@ mutable struct MvGaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM logdetΣ::T in_dim::Int out_dim::Int - - #= Scratch for hot paths. Sequential use only — not safe for concurrent - calls on the same instance. HMM E-steps run forward/backward serially - per state, so each emission has its own scratch and this is fine. =# - _diff::Vector{T} - _z::Vector{T} end function MvGaussianGLM( @@ -594,7 +588,6 @@ function MvGaussianGLM( return MvGaussianGLM{T, P}( Matrix{T}(B), Matrix{T}(Σ), prior, Σ_chol, logdet(Σ_chol), p, k, - Vector{T}(undef, k), Vector{T}(undef, k), ) end @@ -618,8 +611,8 @@ DensityInterface.DensityKind(::MvGaussianGLM) = DensityInterface.HasDensity() """ logdensityof(glm::MvGaussianGLM, y::AbstractVector; control_seq) -Log density of `y ∈ ℝᵏ` under the conditional MvNormal model. Zero allocation: -the residual buffer `_diff` is pre-allocated in the struct. +Log density of `y ∈ ℝᵏ` under the conditional MvNormal model. Allocates one +length-`k` residual vector per call — thread-safe (matches `Distributions.MvNormal`). """ function DensityInterface.logdensityof( glm::MvGaussianGLM, @@ -636,7 +629,7 @@ function DensityInterface.logdensityof( k = glm.out_dim p = glm.in_dim T = eltype(glm.B) - diff = glm._diff + diff = Vector{T}(undef, k) #= μ = Bᵀ x and diff = y − μ, fused as a single loop to avoid the Adjoint*Vector temporary that BLAS would otherwise allocate. =# @@ -669,8 +662,12 @@ end """ rand!(rng, glm::MvGaussianGLM, out; control_seq) -In-place sample. `out` must be a length-`out_dim` `AbstractVector{T}`. -Zero allocation. +In-place sample into `out` (length `out_dim`). Zero allocation, thread-safe: +the trick is to draw `z ~ N(0, I)` directly into `out`, multiply by `L` +in-place (`lmul!`), then add `μ = Bᵀ x` element-wise. Lower-triangular +multiply is well-defined in place because each output `xᵢ = Σⱼ≤ᵢ Lᵢⱼ zⱼ` +only reads `z[1..i]`, so iterating `i = k, k-1, …, 1` never reads an +already-overwritten entry. """ function Random.rand!( rng::AbstractRNG, @@ -687,19 +684,18 @@ function Random.rand!( k = glm.out_dim p = glm.in_dim - z = glm._z - randn!(rng, z) - #= out = Bᵀ x (manual loop avoids the Adjoint*Vector temporary) =# + randn!(rng, out) + lmul!(glm.Σ_chol.L, out) # out = L * z, in place + + # out += μ = Bᵀ x (Bᵀx is computed inline; no aliasing with out) for j in 1:k sj = zero(T) for r in 1:p sj += glm.B[r, j] * control_seq[r] end - out[j] = sj + out[j] += sj end - #= out += L z via in-place 5-arg mul! =# - mul!(out, glm.Σ_chol.L, z, true, true) return out end diff --git a/src/multivariate/t.jl b/src/multivariate/t.jl index 7a4f069..6a2e311 100644 --- a/src/multivariate/t.jl +++ b/src/multivariate/t.jl @@ -147,11 +147,12 @@ function DensityInterface.logdensityof(dist::MultivariateT, x::AbstractVector) d = dist.dim ν = dist.ν - diff = dist._diff - #= Zero allocation: reuse the struct scratch buffer for the residual. - Sequential use only — not safe for concurrent calls on the same dist. =# - @. diff = x - dist.μ + #= Allocate the residual locally (one length-d vector). Thread-safe: + parallel forward/backward calls on the same dist don't share state. + The struct scratch (`_diff`, `_z`) is reserved for `fit!`, which + runs sequentially per state inside HMM EM. =# + diff = x - dist.μ ldiv!(dist.Σ_chol.L, diff) mahal² = zero(eltype(diff)) for i in eachindex(diff) diff --git a/test/allocations.jl b/test/allocations.jl index 6db8fd4..04212d2 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -6,18 +6,20 @@ using DensityInterface using StatsAPI #= - Allocation regression tests. The pattern - is: warm up by running the operation once, then call it many times in a - type-stable function and divide @allocated by the rep count to get - per-call bytes. + Allocation regression tests. Warm up first, then call the operation many + times in a type-stable function and divide @allocated by the rep count. Steady-state targets (per call): - - All logdensityof: 0 bytes - - All rand!: 0 bytes (rand allocates only the return vector) - - fit! is bounded by O(p² + k²), independent of n - - See `bench_*` helpers below for why a top-level `@allocated foo()` would - report misleading kwarg-lowering noise that does NOT exist in real loops. + - Univariate logdensityof, MvBernoulli/MvPoisson logdensityof, + MultivariateTDiag logdensityof, PoissonZeroInflated logdensityof: 0 B + - MvGaussianGLM logdensityof, MultivariateT logdensityof: 1 length-k + vector per call. Thread-safe by design — the struct scratch was removed + because two threads reading the same dist would race. Bound to 256 B/call. + - All Mv rand!: 0 bytes (zero-alloc, thread-safe via in-place lmul!). + - fit! is bounded by O(p² + k²), independent of n. + + Top-level `@allocated foo()` reports kwarg-lowering noise (~48 B) that + does NOT exist in real loops — measure inside a function. =# bench_logd(d, y, x, n) = (s = 0.0; for _ in 1:n; s += logdensityof(d, y; control_seq=x); end; s) @@ -33,7 +35,7 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; rng = Random.MersenneTwister(0) REPS = 1000 - @testset "logdensityof — zero alloc per call" begin + @testset "logdensityof — bounded per call" begin x = [1.0, 2.0] g = GaussianGLM([0.5, -1.0], 1.0) @@ -48,12 +50,16 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; bench_logd(g, 0.5, x, 1); bench_logd(b, 1, x, 1); bench_logd(p, 2, x, 1) bench_logd(mg, yv, x, 1); bench_logd(mb, yi, x, 1); bench_logd(mp, yi, x, 1) + # Truly zero-alloc (no scratch needed) @test (@allocated bench_logd(g, 0.5, x, REPS)) == 0 @test (@allocated bench_logd(b, 1, x, REPS)) == 0 @test (@allocated bench_logd(p, 2, x, REPS)) == 0 - @test (@allocated bench_logd(mg, yv, x, REPS)) == 0 @test (@allocated bench_logd(mb, yi, x, REPS)) == 0 @test (@allocated bench_logd(mp, yi, x, REPS)) == 0 + + # MvGaussianGLM: one length-k residual per call (thread-safe). + # Bound to ~256 B/call → 256 * REPS for the loop. + @test (@allocated bench_logd(mg, yv, x, REPS)) ≤ 256 * REPS end @testset "rand!/rand — zero alloc beyond return" begin @@ -152,9 +158,10 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; mvt = MultivariateT([0.0, 0.0], [1.0 0.3; 0.3 1.0], 5.0) xv = [0.1, 0.2] - # Cholesky now stored as :L and residual reuses struct scratch. + # Cholesky stored as :L (no .L copy). Residual is allocated locally + # per call so concurrent calls on the same dist are race-free. bench_logd_unctrl(mvt, xv, 1) - @test (@allocated bench_logd_unctrl(mvt, xv, REPS)) == 0 + @test (@allocated bench_logd_unctrl(mvt, xv, REPS)) ≤ 256 * REPS bench_rand_unctrl_vec(rng, mvt, 1) # rand still allocates the return vector; bound it loosely. From 09348a01a24bdec940118a144d2f1764feb7cba4 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sun, 10 May 2026 21:40:26 -0400 Subject: [PATCH 07/20] Format and fix JET issue --- src/glms/glm.jl | 353 +++++++++++++++++------------ src/multivariate/t.jl | 2 +- test/Project.toml | 5 +- test/allocations.jl | 128 ++++++++--- test/glm/gaussian.jl | 216 +++++++++--------- test/glm/test_bernoulli_poisson.jl | 19 +- test/runtests.jl | 4 +- 7 files changed, 431 insertions(+), 296 deletions(-) diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 4fa8b5d..e1e8f36 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -82,33 +82,39 @@ exact MAP solution for a Gaussian prior β ~ N(0, (1/λ)I). - `σ2`: noise variance - `prior`: regularization prior (default `NoPrior()`) """ -mutable struct GaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct GaussianGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} σ2::T prior::P end -GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:Real} = - GaussianGLM{T, NoPrior}(Vector{T}(β), σ2, NoPrior()) +function GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:Real} + return GaussianGLM{T,NoPrior}(Vector{T}(β), σ2, NoPrior()) +end -GaussianGLM(β::AbstractVector{T}, σ2::T, prior::P) where {T<:Real, P<:AbstractPrior} = - GaussianGLM{T, P}(Vector{T}(β), σ2, prior) +function GaussianGLM(β::AbstractVector{T}, σ2::T, prior::P) where {T<:Real,P<:AbstractPrior} + return GaussianGLM{T,P}(Vector{T}(β), σ2, prior) +end # Convenience: promote β eltype and σ2 together -GaussianGLM(β::AbstractVector, σ2::Real) = - GaussianGLM(float.(β), float(σ2)) +GaussianGLM(β::AbstractVector, σ2::Real) = GaussianGLM(float.(β), float(σ2)) -GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior) = - GaussianGLM(float.(β), float(σ2), prior) +function GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior) + return GaussianGLM(float.(β), float(σ2), prior) +end DensityInterface.DensityKind(::GaussianGLM) = DensityInterface.HasDensity() -function DensityInterface.logdensityof(reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real}) +function DensityInterface.logdensityof( + reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real} +) μ = dot(reg.β, control_seq) return -0.5 * log(2π * reg.σ2) - 0.5 * ((y - μ)^2 / reg.σ2) end -function Random.rand(rng::AbstractRNG, reg::GaussianGLM; control_seq::AbstractVector{<:Real}) +function Random.rand( + rng::AbstractRNG, reg::GaussianGLM; control_seq::AbstractVector{<:Real} +) return rand(rng, Normal(dot(reg.β, control_seq), sqrt(reg.σ2))) end @@ -131,21 +137,20 @@ function StatsAPI.fit!( wsum = zero(T) for i in 1:n - w = T(weights[i]) + w = T(weights[i]) wsum += w x_i = view(control_seq, i, :) y_i = T(obs_seq[i]) for a in 1:p - xa = x_i[a] + xa = x_i[a] wxa = w * xa for b in 1:p XWX[a, b] += wxa * x_i[b] end XWy[a] += wxa * y_i end - end + end#= RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. =# - #= RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. =# neglogprior_hess!(reg.prior, XWX, reg.β) F = cholesky!(Symmetric(XWX)) @@ -175,8 +180,9 @@ struct _ColumnElementView{T,V<:AbstractVector{<:AbstractVector}} <: AbstractVect seq::V j::Int end -_ColumnElementView(seq::V, j::Int) where {V<:AbstractVector{<:AbstractVector}} = - _ColumnElementView{eltype(eltype(V)), V}(seq, j) +function _ColumnElementView(seq::V, j::Int) where {V<:AbstractVector{<:AbstractVector}} + return _ColumnElementView{eltype(eltype(V)),V}(seq, j) +end Base.size(c::_ColumnElementView) = (length(c.seq),) Base.IndexStyle(::Type{<:_ColumnElementView}) = IndexLinear() Base.@propagate_inbounds Base.getindex(c::_ColumnElementView, i::Integer) = c.seq[i][c.j] @@ -212,7 +218,7 @@ function _bernoulli_gh!( fill!(H, zero(T)) for i in 1:n x_i = view(X, i, :) - wi = T(w[i]) + wi = T(w[i]) μ_i = logistic(dot(β, x_i)) r_i = wi * (μ_i - T(y[i])) W_i = wi * μ_i * (one(T) - μ_i) @@ -261,9 +267,9 @@ function _poisson_gh!( fill!(H, zero(T)) for i in 1:n x_i = view(X, i, :) - wi = T(w[i]) + wi = T(w[i]) η_i = clamp(dot(β, x_i), T(-500), T(500)) - eη = exp(η_i) + eη = exp(η_i) r_i = wi * (eη - T(y[i])) W_i = wi * eη for a in 1:p @@ -298,7 +304,7 @@ function _newton_solve!( gtol::Real=1e-8, max_backtrack::Int=20, ridge::Real=1e-12, -) where {T<:Real, F1, F2} +) where {T<:Real,F1,F2} p = length(β) f_curr = loss(β, y, w, X, prior) @@ -361,8 +367,9 @@ function _fit_bernoulli_glm!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) length(β) == p || throw(DimensionMismatch("β length $(length(β)) ≠ control_seq columns $p")) @@ -370,9 +377,19 @@ function _fit_bernoulli_glm!( H = Matrix{T}(undef, p, p) Δ = Vector{T}(undef, p) return _newton_solve!( - β, g, H, Δ, _bernoulli_loss, _bernoulli_gh!, - obs_seq, weight_seq, control_seq, prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β, + g, + H, + Δ, + _bernoulli_loss, + _bernoulli_gh!, + obs_seq, + weight_seq, + control_seq, + prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) end @@ -389,8 +406,9 @@ function _fit_poisson_glm!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) length(β) == p || throw(DimensionMismatch("β length $(length(β)) ≠ control_seq columns $p")) @@ -398,9 +416,19 @@ function _fit_poisson_glm!( H = Matrix{T}(undef, p, p) Δ = Vector{T}(undef, p) return _newton_solve!( - β, g, H, Δ, _poisson_loss, _poisson_gh!, - obs_seq, weight_seq, control_seq, prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β, + g, + H, + Δ, + _poisson_loss, + _poisson_gh!, + obs_seq, + weight_seq, + control_seq, + prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) end @@ -418,30 +446,32 @@ to the solver. - `β`: coefficient vector (length p) - `prior`: regularization prior (default `NoPrior()`) """ -mutable struct BernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct BernoulliGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} prior::P end -BernoulliGLM(β::AbstractVector{T}) where {T<:Real} = - BernoulliGLM{T, NoPrior}(Vector{T}(β), NoPrior()) +function BernoulliGLM(β::AbstractVector{T}) where {T<:Real} + return BernoulliGLM{T,NoPrior}(Vector{T}(β), NoPrior()) +end -BernoulliGLM(β::AbstractVector{T}, prior::P) where {T<:Real, P<:AbstractPrior} = - BernoulliGLM{T, P}(Vector{T}(β), prior) +function BernoulliGLM(β::AbstractVector{T}, prior::P) where {T<:Real,P<:AbstractPrior} + return BernoulliGLM{T,P}(Vector{T}(β), prior) +end DensityInterface.DensityKind(::BernoulliGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::BernoulliGLM, - y::Integer; - control_seq::AbstractVector{<:Real}, + glm::BernoulliGLM, y::Integer; control_seq::AbstractVector{<:Real} ) (y == 0 || y == 1) || return oftype(dot(glm.β, control_seq), -Inf) η = dot(glm.β, control_seq) return y == 1 ? -log1pexp(-η) : -log1pexp(η) end -function Random.rand(rng::AbstractRNG, glm::BernoulliGLM; control_seq::AbstractVector{<:Real}) +function Random.rand( + rng::AbstractRNG, glm::BernoulliGLM; control_seq::AbstractVector{<:Real} +) return rand(rng, Bernoulli(logistic(dot(glm.β, control_seq)))) end @@ -464,8 +494,14 @@ function StatsAPI.fit!( max_backtrack::Int=20, ) _fit_bernoulli_glm!( - glm.β, obs_seq, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + glm.β, + obs_seq, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) return glm end @@ -483,23 +519,23 @@ composes without changes to the solver. - `β`: coefficient vector (length p) - `prior`: regularization prior (default `NoPrior()`) """ -mutable struct PoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct PoissonGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} prior::P end -PoissonGLM(β::AbstractVector{T}) where {T<:Real} = - PoissonGLM{T, NoPrior}(Vector{T}(β), NoPrior()) +function PoissonGLM(β::AbstractVector{T}) where {T<:Real} + return PoissonGLM{T,NoPrior}(Vector{T}(β), NoPrior()) +end -PoissonGLM(β::AbstractVector{T}, prior::P) where {T<:Real, P<:AbstractPrior} = - PoissonGLM{T, P}(Vector{T}(β), prior) +function PoissonGLM(β::AbstractVector{T}, prior::P) where {T<:Real,P<:AbstractPrior} + return PoissonGLM{T,P}(Vector{T}(β), prior) +end DensityInterface.DensityKind(::PoissonGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::PoissonGLM, - y::Integer; - control_seq::AbstractVector{<:Real}, + glm::PoissonGLM, y::Integer; control_seq::AbstractVector{<:Real} ) y >= 0 || return oftype(dot(glm.β, control_seq), -Inf) η = dot(glm.β, control_seq) @@ -529,8 +565,14 @@ function StatsAPI.fit!( max_backtrack::Int=20, ) _fit_poisson_glm!( - glm.β, obs_seq, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + glm.β, + obs_seq, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) return glm end @@ -557,27 +599,26 @@ prior on `B` (Frobenius-norm penalty). only returns a wrapper without copying when uplo == 'L'; with the upper-stored default, every `.L` access does `copy(factors')`, costing ~k² floats per call. =# -mutable struct MvGaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct MvGaussianGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM B::Matrix{T} Σ::Matrix{T} prior::P - Σ_chol::Cholesky{T, Matrix{T}} + Σ_chol::Cholesky{T,Matrix{T}} logdetΣ::T in_dim::Int out_dim::Int end function MvGaussianGLM( - B::AbstractMatrix{T}, - Σ::AbstractMatrix{T}, - prior::P, -) where {T<:Real, P<:AbstractPrior} + B::AbstractMatrix{T}, Σ::AbstractMatrix{T}, prior::P +) where {T<:Real,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) - size(Σ) == (k, k) || - throw(DimensionMismatch("Σ must be $(k)×$(k) for B with $(k) columns, got $(size(Σ))")) + size(Σ) == (k, k) || throw( + DimensionMismatch("Σ must be $(k)×$(k) for B with $(k) columns, got $(size(Σ))") + ) Σ_chol = try cholesky(Symmetric(Σ, :L)) @@ -585,14 +626,14 @@ function MvGaussianGLM( throw(ArgumentError("Σ must be positive definite")) end - return MvGaussianGLM{T, P}( - Matrix{T}(B), Matrix{T}(Σ), prior, - Σ_chol, logdet(Σ_chol), p, k, + return MvGaussianGLM{T,P}( + Matrix{T}(B), Matrix{T}(Σ), prior, Σ_chol, logdet(Σ_chol), p, k ) end -MvGaussianGLM(B::AbstractMatrix{T}, Σ::AbstractMatrix{T}) where {T<:Real} = - MvGaussianGLM(B, Σ, NoPrior()) +function MvGaussianGLM(B::AbstractMatrix{T}, Σ::AbstractMatrix{T}) where {T<:Real} + return MvGaussianGLM(B, Σ, NoPrior()) +end function MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix) T = promote_type(eltype(B), eltype(Σ)) @@ -615,16 +656,15 @@ Log density of `y ∈ ℝᵏ` under the conditional MvNormal model. Allocates on length-`k` residual vector per call — thread-safe (matches `Distributions.MvNormal`). """ function DensityInterface.logdensityof( - glm::MvGaussianGLM, - y::AbstractVector; - control_seq::AbstractVector{<:Real}, + glm::MvGaussianGLM, y::AbstractVector; control_seq::AbstractVector{<:Real} ) length(y) == glm.out_dim || throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) k = glm.out_dim p = glm.in_dim @@ -650,9 +690,7 @@ function DensityInterface.logdensityof( end function Random.rand( - rng::AbstractRNG, - glm::MvGaussianGLM{T}; - control_seq::AbstractVector{<:Real}, + rng::AbstractRNG, glm::MvGaussianGLM{T}; control_seq::AbstractVector{<:Real} ) where {T<:Real} out = Vector{T}(undef, glm.out_dim) rand!(rng, glm, out; control_seq=control_seq) @@ -677,10 +715,11 @@ function Random.rand!( ) where {T<:Real} length(out) == glm.out_dim || throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) k = glm.out_dim p = glm.in_dim @@ -715,8 +754,9 @@ function StatsAPI.fit!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) p == glm.in_dim || throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) @@ -724,20 +764,17 @@ function StatsAPI.fit!( XWX = zeros(T, p, p) XWY = zeros(T, p, k) - wsum = zero(T) + wsum = zero(T)#= Build XᵀWX and XᵀWY in a single pass — no Y matrix, no temporaries. =# - #= 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 || - throw(DimensionMismatch( - "obs_seq[$i] length $(length(obs_i)) ≠ out_dim $k", - )) - w = T(weight_seq[i]) + throw(DimensionMismatch("obs_seq[$i] length $(length(obs_i)) ≠ out_dim $k")) + w = T(weight_seq[i]) wsum += w x_i = view(control_seq, i, :) for a in 1:p - xa = x_i[a] + xa = x_i[a] wxa = w * xa for b in 1:p XWX[a, b] += wxa * x_i[b] @@ -770,7 +807,7 @@ function StatsAPI.fit!( Σ_new = zeros(T, k, k) for i in 1:n obs_i = obs_seq[i] - w = T(weight_seq[i]) + w = T(weight_seq[i]) x_i = view(control_seq, i, :) for j in 1:k rj = T(obs_i[j]) @@ -820,40 +857,38 @@ For input x ∈ ℝᵖ the conditional distribution of y ∈ {0,1}ᵏ factorizes - `B`: coefficient matrix of size `p × k` - `prior`: regularization prior applied per column (default `NoPrior()`) """ -mutable struct MvBernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct MvBernoulliGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM B::Matrix{T} prior::P in_dim::Int out_dim::Int end -function MvBernoulliGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real, P<:AbstractPrior} +function MvBernoulliGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) - return MvBernoulliGLM{T, P}(Matrix{T}(B), prior, p, k) + return MvBernoulliGLM{T,P}(Matrix{T}(B), prior, p, k) end MvBernoulliGLM(B::AbstractMatrix{T}) where {T<:Real} = MvBernoulliGLM(B, NoPrior()) MvBernoulliGLM(B::AbstractMatrix) = MvBernoulliGLM(float.(B), NoPrior()) -MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior) = - MvBernoulliGLM(float.(B), prior) +MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior) = MvBernoulliGLM(float.(B), prior) DensityInterface.DensityKind(::MvBernoulliGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::MvBernoulliGLM, - y::AbstractVector; - control_seq::AbstractVector{<:Real}, + glm::MvBernoulliGLM, y::AbstractVector; control_seq::AbstractVector{<:Real} ) length(y) == glm.out_dim || throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) lp = zero(Tη) @@ -870,9 +905,7 @@ function DensityInterface.logdensityof( end function Random.rand( - rng::AbstractRNG, - glm::MvBernoulliGLM; - control_seq::AbstractVector{<:Real}, + rng::AbstractRNG, glm::MvBernoulliGLM; control_seq::AbstractVector{<:Real} ) out = Vector{Int}(undef, glm.out_dim) rand!(rng, glm, out; control_seq=control_seq) @@ -893,10 +926,11 @@ function Random.rand!( ) length(out) == glm.out_dim || throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) T = eltype(glm.B) for j in 1:glm.out_dim @@ -929,31 +963,41 @@ function StatsAPI.fit!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) p == glm.in_dim || throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) k = glm.out_dim for i in 1:n - length(obs_seq[i]) == k || - throw(DimensionMismatch( - "obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k", - )) + length(obs_seq[i]) == k || throw( + DimensionMismatch("obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k") + ) end β_buf = Vector{T}(undef, p) - g = Vector{T}(undef, p) - H = Matrix{T}(undef, p, p) - Δ = Vector{T}(undef, p) + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) for j in 1:k yview = _ColumnElementView(obs_seq, j) copyto!(β_buf, view(glm.B, :, j)) _newton_solve!( - β_buf, g, H, Δ, _bernoulli_loss, _bernoulli_gh!, - yview, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β_buf, + g, + H, + Δ, + _bernoulli_loss, + _bernoulli_gh!, + yview, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) for r in 1:p glm.B[r, j] = β_buf[r] @@ -975,40 +1019,38 @@ For input x ∈ ℝᵖ the conditional distribution of y ∈ ℤ₊ᵏ factorize - `B`: coefficient matrix of size `p × k` - `prior`: regularization prior applied per column (default `NoPrior()`) """ -mutable struct MvPoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct MvPoissonGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM B::Matrix{T} prior::P in_dim::Int out_dim::Int end -function MvPoissonGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real, P<:AbstractPrior} +function MvPoissonGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) - return MvPoissonGLM{T, P}(Matrix{T}(B), prior, p, k) + return MvPoissonGLM{T,P}(Matrix{T}(B), prior, p, k) end MvPoissonGLM(B::AbstractMatrix{T}) where {T<:Real} = MvPoissonGLM(B, NoPrior()) MvPoissonGLM(B::AbstractMatrix) = MvPoissonGLM(float.(B), NoPrior()) -MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior) = - MvPoissonGLM(float.(B), prior) +MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior) = MvPoissonGLM(float.(B), prior) DensityInterface.DensityKind(::MvPoissonGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::MvPoissonGLM, - y::AbstractVector; - control_seq::AbstractVector{<:Real}, + glm::MvPoissonGLM, y::AbstractVector; control_seq::AbstractVector{<:Real} ) length(y) == glm.out_dim || throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) lp = zero(Tη) @@ -1025,9 +1067,7 @@ function DensityInterface.logdensityof( end function Random.rand( - rng::AbstractRNG, - glm::MvPoissonGLM; - control_seq::AbstractVector{<:Real}, + rng::AbstractRNG, glm::MvPoissonGLM; control_seq::AbstractVector{<:Real} ) out = Vector{Int}(undef, glm.out_dim) rand!(rng, glm, out; control_seq=control_seq) @@ -1047,10 +1087,11 @@ function Random.rand!( ) length(out) == glm.out_dim || throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) T = eltype(glm.B) for j in 1:glm.out_dim @@ -1083,31 +1124,41 @@ function StatsAPI.fit!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) p == glm.in_dim || throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) k = glm.out_dim for i in 1:n - length(obs_seq[i]) == k || - throw(DimensionMismatch( - "obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k", - )) + length(obs_seq[i]) == k || throw( + DimensionMismatch("obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k") + ) end β_buf = Vector{T}(undef, p) - g = Vector{T}(undef, p) - H = Matrix{T}(undef, p, p) - Δ = Vector{T}(undef, p) + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) for j in 1:k yview = _ColumnElementView(obs_seq, j) copyto!(β_buf, view(glm.B, :, j)) _newton_solve!( - β_buf, g, H, Δ, _poisson_loss, _poisson_gh!, - yview, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β_buf, + g, + H, + Δ, + _poisson_loss, + _poisson_gh!, + yview, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) for r in 1:p glm.B[r, j] = β_buf[r] diff --git a/src/multivariate/t.jl b/src/multivariate/t.jl index 6a2e311..e181c69 100644 --- a/src/multivariate/t.jl +++ b/src/multivariate/t.jl @@ -314,7 +314,7 @@ function StatsAPI.fit!( end # Symmetrize numerical noise - for j in 1:d, k in 1:j-1 + for j in 1:d, k in 1:(j - 1) s = (Σ_acc[j, k] + Σ_acc[k, j]) / 2 Σ_acc[j, k] = s Σ_acc[k, j] = s diff --git a/test/Project.toml b/test/Project.toml index b89abb9..880ac11 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,14 +3,11 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" 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" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[compat] -JET = "0.11.2" diff --git a/test/allocations.jl b/test/allocations.jl index 04212d2..07e1d5b 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -22,14 +22,72 @@ using StatsAPI does NOT exist in real loops — measure inside a function. =# -bench_logd(d, y, x, n) = (s = 0.0; for _ in 1:n; s += logdensityof(d, y; control_seq=x); end; s) -bench_logd_unctrl(d, y, n) = (s = 0.0; for _ in 1:n; s += logdensityof(d, y); end; s) -bench_rand_scalar(rng, d, x, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d; control_seq=x); end; s) -bench_rand_int(rng, d, x, n) = (s = 0; for _ in 1:n; s += rand(rng, d; control_seq=x); end; s) -bench_rand!_v(rng, d, out, x, n) = (s = 0.0; for _ in 1:n; rand!(rng, d, out; control_seq=x); s += out[1]; end; s) -bench_rand!_i(rng, d, out, x, n) = (s = 0; for _ in 1:n; rand!(rng, d, out; control_seq=x); s += out[1]; end; s) -bench_rand_unctrl_scalar(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d); end; s) -bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; end; s) +bench_logd(d, y, x, n) = ( + s=0.0; + for _ in 1:n + ; + s += logdensityof(d, y; control_seq=x); + end; + s +) +bench_logd_unctrl(d, y, n) = ( + s=0.0; + for _ in 1:n + ; + s += logdensityof(d, y); + end; + s +) +bench_rand_scalar(rng, d, x, n) = ( + s=0.0; + for _ in 1:n + ; + s += rand(rng, d; control_seq=x); + end; + s +) +bench_rand_int(rng, d, x, n) = ( + s=0; + for _ in 1:n + ; + s += rand(rng, d; control_seq=x); + end; + s +) +bench_rand!_v(rng, d, out, x, n) = ( + s=0.0; + for _ in 1:n + ; + rand!(rng, d, out; control_seq=x); + s += out[1]; + end; + s +) +bench_rand!_i(rng, d, out, x, n) = ( + s=0; + for _ in 1:n + ; + rand!(rng, d, out; control_seq=x); + s += out[1]; + end; + s +) +bench_rand_unctrl_scalar(rng, d, n) = ( + s=0.0; + for _ in 1:n + ; + s += rand(rng, d); + end; + s +) +bench_rand_unctrl_vec(rng, d, n) = ( + s=0.0; + for _ in 1:n + ; + s += rand(rng, d)[1]; + end; + s +) @testset "Allocations (steady state)" begin rng = Random.MersenneTwister(0) @@ -38,24 +96,29 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; @testset "logdensityof — bounded per call" begin x = [1.0, 2.0] - g = GaussianGLM([0.5, -1.0], 1.0) - b = BernoulliGLM([0.5, -1.0]) - p = PoissonGLM([0.5, -1.0]) + g = GaussianGLM([0.5, -1.0], 1.0) + b = BernoulliGLM([0.5, -1.0]) + p = PoissonGLM([0.5, -1.0]) mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) - yv = [0.1, 0.2]; yi = [0, 1] + yv = [0.1, 0.2]; + yi = [0, 1] # Warm - bench_logd(g, 0.5, x, 1); bench_logd(b, 1, x, 1); bench_logd(p, 2, x, 1) - bench_logd(mg, yv, x, 1); bench_logd(mb, yi, x, 1); bench_logd(mp, yi, x, 1) + bench_logd(g, 0.5, x, 1); + bench_logd(b, 1, x, 1); + bench_logd(p, 2, x, 1) + bench_logd(mg, yv, x, 1); + bench_logd(mb, yi, x, 1); + bench_logd(mp, yi, x, 1) # Truly zero-alloc (no scratch needed) - @test (@allocated bench_logd(g, 0.5, x, REPS)) == 0 - @test (@allocated bench_logd(b, 1, x, REPS)) == 0 - @test (@allocated bench_logd(p, 2, x, REPS)) == 0 - @test (@allocated bench_logd(mb, yi, x, REPS)) == 0 - @test (@allocated bench_logd(mp, yi, x, REPS)) == 0 + @test (@allocated bench_logd(g, 0.5, x, REPS)) == 0 + @test (@allocated bench_logd(b, 1, x, REPS)) == 0 + @test (@allocated bench_logd(p, 2, x, REPS)) == 0 + @test (@allocated bench_logd(mb, yi, x, REPS)) == 0 + @test (@allocated bench_logd(mp, yi, x, REPS)) == 0 # MvGaussianGLM: one length-k residual per call (thread-safe). # Bound to ~256 B/call → 256 * REPS for the loop. @@ -65,9 +128,9 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; @testset "rand!/rand — zero alloc beyond return" begin x = [1.0, 2.0] - g = GaussianGLM([0.5, -1.0], 1.0) - b = BernoulliGLM([0.5, -1.0]) - p = PoissonGLM([0.5, -1.0]) + g = GaussianGLM([0.5, -1.0], 1.0) + b = BernoulliGLM([0.5, -1.0]) + p = PoissonGLM([0.5, -1.0]) mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) @@ -98,7 +161,8 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; # GaussianGLM: closed-form WLS. Workspace: XWX (p²) + XWy (p). yg = randn(rng, n) - gg = GaussianGLM([0.0, 0.0], 1.0); fit!(gg, yg, w; control_seq=X) + gg = GaussianGLM([0.0, 0.0], 1.0); + fit!(gg, yg, w; control_seq=X) gg = GaussianGLM([0.0, 0.0], 1.0) @test (@allocated fit!(gg, yg, w; control_seq=X)) ≤ 1_000 @@ -113,24 +177,28 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; # BernoulliGLM/PoissonGLM: hand-rolled Newton. # Workspace: g, H, Δ — three small alloc, total ~300 bytes for p=2. yb = Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:n] - gb = BernoulliGLM(zeros(2)); fit!(gb, yb, w; control_seq=X) + gb = BernoulliGLM(zeros(2)); + fit!(gb, yb, w; control_seq=X) gb = BernoulliGLM(zeros(2)) @test (@allocated fit!(gb, yb, w; control_seq=X)) ≤ 1_000 yp = Int[rand(rng, 0:5) for _ in 1:n] - gp = PoissonGLM(zeros(2)); fit!(gp, yp, w; control_seq=X) + gp = PoissonGLM(zeros(2)); + fit!(gp, yp, w; control_seq=X) gp = PoissonGLM(zeros(2)) @test (@allocated fit!(gp, yp, w; control_seq=X)) ≤ 1_000 # Multivariate Newton: workspace shared across columns. _ColumnElementView # avoids the n-sized per-column copy that Vector-of-Vectors would force. ymb = [Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:2] for _ in 1:n] - gmb = MvBernoulliGLM(zeros(2, 2)); fit!(gmb, ymb, w; control_seq=X) + gmb = MvBernoulliGLM(zeros(2, 2)); + fit!(gmb, ymb, w; control_seq=X) gmb = MvBernoulliGLM(zeros(2, 2)) @test (@allocated fit!(gmb, ymb, w; control_seq=X)) ≤ 1_000 ymp = [Int[rand(rng, 0:3) for _ in 1:2] for _ in 1:n] - gmp = MvPoissonGLM(zeros(2, 2)); fit!(gmp, ymp, w; control_seq=X) + gmp = MvPoissonGLM(zeros(2, 2)); + fit!(gmp, ymp, w; control_seq=X) gmp = MvPoissonGLM(zeros(2, 2)) @test (@allocated fit!(gmp, ymp, w; control_seq=X)) ≤ 1_000 end @@ -138,7 +206,8 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; @testset "PoissonZeroInflated" begin zip = PoissonZeroInflated(3.0, 0.2) - bench_logd_unctrl(zip, 0, 1); bench_logd_unctrl(zip, 5, 1) + bench_logd_unctrl(zip, 0, 1); + bench_logd_unctrl(zip, 5, 1) @test (@allocated bench_logd_unctrl(zip, 0, REPS)) == 0 @test (@allocated bench_logd_unctrl(zip, 5, REPS)) == 0 @@ -148,7 +217,8 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; n = 500 y = [rand(rng, zip) for _ in 1:n] w = ones(n) - zip2 = PoissonZeroInflated(1.0, 0.1); fit!(zip2, y, w) + zip2 = PoissonZeroInflated(1.0, 0.1); + fit!(zip2, y, w) zip2 = PoissonZeroInflated(1.0, 0.1) @test (@allocated fit!(zip2, y, w)) ≤ 1_000 end diff --git a/test/glm/gaussian.jl b/test/glm/gaussian.jl index 9fd6499..d58f1c1 100644 --- a/test/glm/gaussian.jl +++ b/test/glm/gaussian.jl @@ -9,133 +9,143 @@ using EmissionModels rng = Random.MersenneTwister(1234) -function _synthetic_gaussian_glm(rng, n::Int, p::Int; β::Vector{Float64}, σ2::Float64, weights=:uniform) +function _synthetic_gaussian_glm( + rng, n::Int, p::Int; β::Vector{Float64}, σ2::Float64, weights=:uniform +) @assert length(β) == p X = hcat(ones(n), randn(rng, n, p-1)) # intercept + random features μ = X * β y = rand.(Ref(rng), Normal.(μ, sqrt(σ2))) - w = weights === :uniform ? ones(n) : - weights === :random ? (rand(rng, n) .+ 0.5) : + w = if weights === :uniform + ones(n) + elseif weights === :random + (rand(rng, n) .+ 0.5) + else error("weights must be :uniform or :random") + end return X, y, w end _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) @testset "GaussianGLM Basic Functionality" begin - @testset "Constructor" begin - glm = GaussianGLM([0.5, -1.0], 1.0) - @test glm.β == [0.5, -1.0] - @test glm.σ2 == 1.0 - - glm2 = GaussianGLM([1, 2], 3) - @test glm2.β == [1, 2] - @test glm2.σ2 == 3 - end + @testset "Constructor" begin + glm = GaussianGLM([0.5, -1.0], 1.0) + @test glm.β == [0.5, -1.0] + @test glm.σ2 == 1.0 - @testset "DensityInterface" begin - glm = GaussianGLM([0.5, -1.0], 1.0) + glm2 = GaussianGLM([1, 2], 3) + @test glm2.β == [1, 2] + @test glm2.σ2 == 3 + end - # Trait - @test DensityKind(glm) == HasDensity() + @testset "DensityInterface" begin + glm = GaussianGLM([0.5, -1.0], 1.0) - X = [1.0 2.0; 1.0 3.0; 1.0 4.0] - μ = X * glm.β - y = rand.(Ref(rng), Normal.(μ, sqrt(glm.σ2))) + # Trait + @test DensityKind(glm) == HasDensity() - # logdensityof matches closed form - for i in eachindex(y) - x_i = vec(X[i, :]) # ensure Vector - logp = logdensityof(glm, y[i]; control_seq=x_i) - @test isfinite(logp) + X = [1.0 2.0; 1.0 3.0; 1.0 4.0] + μ = X * glm.β + y = rand.(Ref(rng), Normal.(μ, sqrt(glm.σ2))) - expected = _expected_logpdf(y[i], dot(glm.β, x_i), glm.σ2) - @test logp ≈ expected rtol=1e-12 atol=0.0 - end + # logdensityof matches closed form + for i in eachindex(y) + x_i = vec(X[i, :]) # ensure Vector + logp = logdensityof(glm, y[i]; control_seq=x_i) + @test isfinite(logp) - # Higher density near the mean than far away - x0 = vec(X[1, :]) - y_at_mean = dot(glm.β, x0) - y_far = y_at_mean + 10.0 - @test logdensityof(glm, y_at_mean; control_seq=x0) > logdensityof(glm, y_far; control_seq=x0) + expected = _expected_logpdf(y[i], dot(glm.β, x_i), glm.σ2) + @test logp ≈ expected rtol=1e-12 atol=0.0 end - @testset "Random sampling" begin - glm = GaussianGLM([0.5, -1.0], 2.0) # variance 2.0 - x = [1.0, 3.0] - n = 10_000 - samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + # Higher density near the mean than far away + x0 = vec(X[1, :]) + y_at_mean = dot(glm.β, x0) + y_far = y_at_mean + 10.0 + @test logdensityof(glm, y_at_mean; control_seq=x0) > + logdensityof(glm, y_far; control_seq=x0) + end - @test all(s -> s isa Real, samples) + @testset "Random sampling" begin + glm = GaussianGLM([0.5, -1.0], 2.0) # variance 2.0 + x = [1.0, 3.0] + n = 10_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] - # Empirical mean should be close to dot(β,x) - μ = dot(glm.β, x) - m = mean(samples) - @test m ≈ μ atol=0.15 + @test all(s -> s isa Real, samples) - # Empirical variance should be close to σ2 - v = var(samples) - @test v ≈ glm.σ2 atol=0.25 - end + # Empirical mean should be close to dot(β,x) + μ = dot(glm.β, x) + m = mean(samples) + @test m ≈ μ atol=0.15 - @testset "fit! with uniform weights" begin - β_true = [0.3, -1.2] # p=2 - σ2_true = 0.7 - X, y, w = _synthetic_gaussian_glm(rng, 2000, 2; β=β_true, σ2=σ2_true, weights=:uniform) + # Empirical variance should be close to σ2 + v = var(samples) + @test v ≈ glm.σ2 atol=0.25 + end - glm = GaussianGLM([0.0, 0.0], 1.0) - fit!(glm, y, w; control_seq=X) + @testset "fit! with uniform weights" begin + β_true = [0.3, -1.2] # p=2 + σ2_true = 0.7 + X, y, w = _synthetic_gaussian_glm( + rng, 2000, 2; β=β_true, σ2=σ2_true, weights=:uniform + ) - # Coefficients should be close with enough data - @test glm.β ≈ β_true atol=0.08 - @test glm.σ2 ≈ σ2_true atol=0.10 + glm = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm, y, w; control_seq=X) - @test all(isfinite, glm.β) - @test isfinite(glm.σ2) - @test glm.σ2 > 0 - end + # Coefficients should be close with enough data + @test glm.β ≈ β_true atol=0.08 + @test glm.σ2 ≈ σ2_true atol=0.10 - @testset "fit! with weighted observations" begin - β_true = [1.0, 0.6] - σ2_true = 1.5 - X, y, w = _synthetic_gaussian_glm(rng, 2500, 2; β=β_true, σ2=σ2_true, weights=:random) + @test all(isfinite, glm.β) + @test isfinite(glm.σ2) + @test glm.σ2 > 0 + end - glm = GaussianGLM([0.0, 0.0], 1.0) - fit!(glm, y, w; control_seq=X) + @testset "fit! with weighted observations" begin + β_true = [1.0, 0.6] + σ2_true = 1.5 + X, y, w = _synthetic_gaussian_glm( + rng, 2500, 2; β=β_true, σ2=σ2_true, weights=:random + ) - @test glm.β ≈ β_true atol=0.10 - @test glm.σ2 ≈ σ2_true atol=0.15 - end + glm = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm, y, w; control_seq=X) - @testset "Constructor with prior" begin - glm = GaussianGLM([0.0, 0.0], 1.0) - @test glm.prior isa NoPrior + @test glm.β ≈ β_true atol=0.10 + @test glm.σ2 ≈ σ2_true atol=0.15 + end - glm2 = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(2.0)) - @test glm2.prior isa RidgePrior - @test glm2.prior.λ == 2.0 - end + @testset "Constructor with prior" begin + glm = GaussianGLM([0.0, 0.0], 1.0) + @test glm.prior isa NoPrior - @testset "fit! with RidgePrior shrinks toward zero" begin - β_true = [3.0, -3.0] - σ2_true = 1.0 - X, y, w = _synthetic_gaussian_glm(rng, 300, 2; β=β_true, σ2=σ2_true) + glm2 = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(2.0)) + @test glm2.prior isa RidgePrior + @test glm2.prior.λ == 2.0 + end - glm_noprior = GaussianGLM([0.0, 0.0], 1.0) - fit!(glm_noprior, y, w; control_seq=X) + @testset "fit! with RidgePrior shrinks toward zero" begin + β_true = [3.0, -3.0] + σ2_true = 1.0 + X, y, w = _synthetic_gaussian_glm(rng, 300, 2; β=β_true, σ2=σ2_true) - glm_ridge = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(10.0)) - fit!(glm_ridge, y, w; control_seq=X) + glm_noprior = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm_noprior, y, w; control_seq=X) - @test norm(glm_ridge.β) < norm(glm_noprior.β) - @test all(isfinite, glm_ridge.β) - end + glm_ridge = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(10.0)) + fit!(glm_ridge, y, w; control_seq=X) + @test norm(glm_ridge.β) < norm(glm_noprior.β) + @test all(isfinite, glm_ridge.β) + end end -function _synthetic_mvgaussian_glm(rng, n::Int, p::Int, k::Int; - B::Matrix{Float64}, Σ::Matrix{Float64}, - weights=:uniform) +function _synthetic_mvgaussian_glm( + rng, n::Int, p::Int, k::Int; B::Matrix{Float64}, Σ::Matrix{Float64}, weights=:uniform +) @assert size(B) == (p, k) @assert size(Σ) == (k, k) X = hcat(ones(n), randn(rng, n, p-1)) @@ -145,9 +155,13 @@ function _synthetic_mvgaussian_glm(rng, n::Int, p::Int, k::Int; μ_i = vec(B' * X[i, :]) obs_seq[i] = μ_i .+ L * randn(rng, k) end - w = weights === :uniform ? ones(n) : - weights === :random ? (rand(rng, n) .+ 0.5) : + w = if weights === :uniform + ones(n) + elseif weights === :random + (rand(rng, n) .+ 0.5) + else error("weights must be :uniform or :random") + end return X, obs_seq, w end @@ -171,7 +185,9 @@ end glm_ridge = MvGaussianGLM(B, Σ, RidgePrior(2.0)) @test glm_ridge.prior isa RidgePrior - @test_throws DimensionMismatch MvGaussianGLM(B, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]) + @test_throws DimensionMismatch MvGaussianGLM( + B, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0] + ) @test_throws ArgumentError MvGaussianGLM(B, [1.0 0.0; 0.0 -1.0]) end @@ -193,7 +209,7 @@ end # Higher density at the mean than far from it @test logdensityof(glm, μ; control_seq=x) > - logdensityof(glm, μ .+ 10.0; control_seq=x) + logdensityof(glm, μ .+ 10.0; control_seq=x) @test_throws DimensionMismatch logdensityof(glm, [1.0]; control_seq=x) @test_throws DimensionMismatch logdensityof(glm, y; control_seq=[1.0]) @@ -220,8 +236,7 @@ end @testset "fit! recovers B and Σ" begin B_true = [0.5 -1.0; 1.0 0.5; -0.3 0.8] # p=3, k=2 Σ_true = [1.0 0.3; 0.3 1.5] - X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 4000, 3, 2; - B=B_true, Σ=Σ_true) + X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 4000, 3, 2; B=B_true, Σ=Σ_true) glm = MvGaussianGLM(zeros(3, 2), Matrix(1.0I, 2, 2)) fit!(glm, obs_seq, w; control_seq=X) @@ -236,9 +251,9 @@ end @testset "fit! random weights" begin B_true = [1.0 0.0; -0.5 1.5] Σ_true = [0.5 -0.1; -0.1 0.7] - X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 3000, 2, 2; - B=B_true, Σ=Σ_true, - weights=:random) + X, obs_seq, w = _synthetic_mvgaussian_glm( + rng, 3000, 2, 2; B=B_true, Σ=Σ_true, weights=:random + ) glm = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) fit!(glm, obs_seq, w; control_seq=X) @@ -250,8 +265,7 @@ end @testset "fit! with RidgePrior shrinks toward zero" begin B_true = [3.0 -3.0; 3.0 3.0] Σ_true = [1.0 0.0; 0.0 1.0] - X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 200, 2, 2; - B=B_true, Σ=Σ_true) + X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 200, 2, 2; B=B_true, Σ=Σ_true) glm_noprior = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) fit!(glm_noprior, obs_seq, w; control_seq=X) @@ -275,4 +289,4 @@ end wrong_dim_obs = [zeros(3) for _ in 1:5] @test_throws DimensionMismatch fit!(glm, wrong_dim_obs, ones(5); control_seq=X) end -end \ No newline at end of file +end diff --git a/test/glm/test_bernoulli_poisson.jl b/test/glm/test_bernoulli_poisson.jl index b30429f..f481438 100644 --- a/test/glm/test_bernoulli_poisson.jl +++ b/test/glm/test_bernoulli_poisson.jl @@ -6,7 +6,6 @@ using StatsAPI using LinearAlgebra using Distributions: Bernoulli, Poisson, logpdf - function _sigmoid(η) return one(η) / (one(η) + exp(-η)) end @@ -160,7 +159,9 @@ end X = ones(5, 2) @test_throws DimensionMismatch fit!(glm, ones(Int, 4), ones(5); control_seq=X) @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(4); control_seq=X) - @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(5); control_seq=ones(5, 3)) + @test_throws DimensionMismatch fit!( + glm, ones(Int, 5), ones(5); control_seq=ones(5, 3) + ) end @testset "fit! all-zero observations" begin @@ -177,7 +178,6 @@ end end end - @testset "PoissonGLM" begin rng = Random.MersenneTwister(99) @@ -255,7 +255,9 @@ end X = ones(5, 2) @test_throws DimensionMismatch fit!(glm, ones(Int, 4), ones(5); control_seq=X) @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(4); control_seq=X) - @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(5); control_seq=ones(5, 3)) + @test_throws DimensionMismatch fit!( + glm, ones(Int, 5), ones(5); control_seq=ones(5, 3) + ) end @testset "fit! all-zero counts" begin @@ -277,8 +279,9 @@ function _synthetic_mvbernoulli(rng, n, B_true; weights=:uniform) X = hcat(ones(n), randn(rng, n, p - 1)) obs_seq = Vector{Vector{Int}}(undef, n) for i in 1:n - obs_seq[i] = Int[rand(rng) < _sigmoid(dot(B_true[:, j], X[i, :])) ? 1 : 0 - for j in 1:k] + obs_seq[i] = Int[ + rand(rng) < _sigmoid(dot(B_true[:, j], X[i, :])) ? 1 : 0 for j in 1:k + ] end w = weights === :uniform ? ones(n) : rand(rng, n) .+ 0.5 return X, obs_seq, w @@ -406,7 +409,7 @@ end X = ones(5, 2) good_obs = [zeros(Int, 2) for _ in 1:5] @test_throws DimensionMismatch fit!( - glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X, + glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X ) @test_throws DimensionMismatch fit!(glm, good_obs, ones(4); control_seq=X) @test_throws DimensionMismatch fit!(glm, good_obs, ones(5); control_seq=ones(5, 3)) @@ -521,7 +524,7 @@ end X = ones(5, 2) good_obs = [zeros(Int, 2) for _ in 1:5] @test_throws DimensionMismatch fit!( - glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X, + glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X ) @test_throws DimensionMismatch fit!(glm, good_obs, ones(4); control_seq=X) @test_throws DimensionMismatch fit!(glm, good_obs, ones(5); control_seq=ones(5, 3)) diff --git a/test/runtests.jl b/test/runtests.jl index a66ea46..a738d27 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,13 +1,13 @@ using EmissionModels using Test using Aqua -using JET -using JuliaFormatter using StatsAPI @testset "EmissionModels.jl" begin @testset "Code linting" begin if VERSION >= v"1.10" && isempty(VERSION.prerelease) + using Pkg + Pkg.add("JET") JET.test_package(EmissionModels; target_modules=(EmissionModels,)) else @info "Skipping JET on Julia $VERSION (requires >=1.10 and a release build)" From 90c6b8fd902f3784bfafe06209c423093f6c18c6 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 11 May 2026 11:53:34 -0400 Subject: [PATCH 08/20] docs fixes --- docs/Project.toml | 2 ++ docs/make.jl | 3 +++ docs/src/distributions.md | 17 +++++++++++++ docs/src/glm.md | 50 ++++++++++++++++++++++++++++++++++++++- docs/src/priors.md | 8 +++++++ src/glms/glm.jl | 7 ++---- test/runtests.jl | 1 + 7 files changed, 82 insertions(+), 6 deletions(-) diff --git a/docs/Project.toml b/docs/Project.toml index c56c224..be32084 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,3 +1,5 @@ [deps] +DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" EmissionModels = "1e2dd27c-41a5-43b2-863c-3eddd0c72c67" +StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" diff --git a/docs/make.jl b/docs/make.jl index d6cac3d..f5bcbd8 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,5 +1,8 @@ using EmissionModels using Documenter +using DensityInterface +using Random +using StatsAPI DocMeta.setdocmeta!(EmissionModels, :DocTestSetup, :(using EmissionModels); recursive=true) diff --git a/docs/src/distributions.md b/docs/src/distributions.md index 6bbca5c..b6d1782 100644 --- a/docs/src/distributions.md +++ b/docs/src/distributions.md @@ -81,3 +81,20 @@ fit!(dist, obs_seq, weight_seq; max_iter=100, tol=1e-6, fix_nu=false) - The M-step updates ``μ`` and ``Σ`` (and optionally ``ν``). - Pass `fix_nu=true` to keep the degrees of freedom fixed during fitting. - `μ` and ``Σ`` are re-regularized on-the-fly if the Cholesky factorization fails. + +## API Reference + +```@docs +PoissonZeroInflated +MultivariateT +MultivariateTDiag +DensityInterface.logdensityof(::PoissonZeroInflated, ::Integer) +DensityInterface.logdensityof(::MultivariateT, ::AbstractVector) +DensityInterface.logdensityof(::MultivariateTDiag, ::AbstractVector) +Base.rand(::Random.AbstractRNG, ::PoissonZeroInflated) +Base.rand(::Random.AbstractRNG, ::MultivariateT) +Base.rand(::Random.AbstractRNG, ::MultivariateTDiag) +StatsAPI.fit!(::PoissonZeroInflated, ::AbstractVector, ::AbstractVector) +StatsAPI.fit!(::MultivariateT, ::Any, ::Any) +StatsAPI.fit!(::MultivariateTDiag, ::Any, ::Any) +``` diff --git a/docs/src/glm.md b/docs/src/glm.md index 9ad53c4..1f74600 100644 --- a/docs/src/glm.md +++ b/docs/src/glm.md @@ -42,6 +42,28 @@ Log-linear regression for count observations ``Y ∈ \{0, 1, \ldots\}``. glm = PoissonGLM(zeros(3)) ``` +## Multivariate GLM types + +Each univariate GLM has a multivariate counterpart that emits a length-``k`` observation vector for a single input ``x``. Coefficients are stored as a ``p × k`` matrix ``B``; column ``j`` is the coefficient vector for output dimension ``j``. + +### `MvGaussianGLM(B, Σ)` + +Multivariate linear regression with shared full covariance ``Σ``. + +``Y \mid x \sim \mathcal{N}(Bᵀx, \ Σ)`` + +### `MvBernoulliGLM(B)` + +``k`` independent logistic regressions sharing the same input ``x``. + +``P(Y \mid x) = \prod_{j=1}^{k} \text{Bernoulli}(σ(B[:,j]ᵀx))`` + +### `MvPoissonGLM(B)` + +``k`` independent Poisson log-linear regressions sharing the same input ``x``. + +``P(Y \mid x) = \prod_{j=1}^{k} \text{Poisson}(\exp(B[:,j]ᵀx))`` + ## Fitting GLM emissions GLM emissions are fit by minimizing the negative log-posterior (or weighted negative log-likelihood when no prior is used): @@ -49,5 +71,31 @@ GLM emissions are fit by minimizing the negative log-posterior (or weighted nega ``\ell(β) = -\sum_{i=1}^{n} w_i \, \log p(y_i \mid x_i, β) + \text{neglogprior}(β)`` ```julia -fit!(glm, y_seq, w_seq; control_seq=X, optim_opts=Optim.Options()) +# Gaussian is closed-form weighted least squares. +fit!(glm, y_seq, w_seq; control_seq=X) + +# Bernoulli / Poisson use a hand-rolled Newton with backtracking line search. +fit!(glm, y_seq, w_seq; control_seq=X, + max_iter=50, gtol=1e-8, max_backtrack=20) +``` + +## API Reference + +```@docs +EmissionModels.AbstractGLM +GaussianGLM +BernoulliGLM +PoissonGLM +MvGaussianGLM +MvBernoulliGLM +MvPoissonGLM +DensityInterface.logdensityof(::MvGaussianGLM, ::AbstractVector) +StatsAPI.fit!(::BernoulliGLM, ::AbstractVector, ::AbstractVector{<:Real}) +StatsAPI.fit!(::PoissonGLM, ::AbstractVector, ::AbstractVector{<:Real}) +StatsAPI.fit!(::MvGaussianGLM{T}, ::AbstractVector{<:AbstractVector}, ::AbstractVector{<:Real}) where {T<:Real} +StatsAPI.fit!(::MvBernoulliGLM{T}, ::AbstractVector{<:AbstractVector}, ::AbstractVector{<:Real}) where {T<:Real} +StatsAPI.fit!(::MvPoissonGLM{T}, ::AbstractVector{<:AbstractVector}, ::AbstractVector{<:Real}) where {T<:Real} +Random.rand!(::Random.AbstractRNG, ::MvGaussianGLM{T}, ::AbstractVector) where {T<:Real} +Random.rand!(::Random.AbstractRNG, ::MvBernoulliGLM, ::AbstractVector) +Random.rand!(::Random.AbstractRNG, ::MvPoissonGLM, ::AbstractVector) ``` diff --git a/docs/src/priors.md b/docs/src/priors.md index 458916d..4b2316e 100644 --- a/docs/src/priors.md +++ b/docs/src/priors.md @@ -71,3 +71,11 @@ glm = GaussianGLM(zeros(5), 1.0, prior) ``` The composed prior simply sums the individual penalties, gradients, and Hessians. + +## API Reference + +```@docs +AbstractPrior +NoPrior +RidgePrior +``` diff --git a/src/glms/glm.jl b/src/glms/glm.jl index e1e8f36..237eeea 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -149,8 +149,9 @@ function StatsAPI.fit!( end XWy[a] += wxa * y_i end - end#= RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. =# + end + #= RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. =# neglogprior_hess!(reg.prior, XWX, reg.β) F = cholesky!(Symmetric(XWX)) @@ -595,10 +596,6 @@ prior on `B` (Frobenius-norm penalty). - `Σ`: residual covariance of size `k × k` - `prior`: regularization prior on `B` (default `NoPrior()`) """ -#= Cholesky uplo: we always store the LOWER triangle. `getproperty(::Cholesky, :L)` - only returns a wrapper without copying when uplo == 'L'; with the upper-stored - default, every `.L` access does `copy(factors')`, costing ~k² floats per call. =# - mutable struct MvGaussianGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM B::Matrix{T} Σ::Matrix{T} diff --git a/test/runtests.jl b/test/runtests.jl index a738d27..2593bf3 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -8,6 +8,7 @@ using StatsAPI if VERSION >= v"1.10" && isempty(VERSION.prerelease) using Pkg Pkg.add("JET") + using JET JET.test_package(EmissionModels; target_modules=(EmissionModels,)) else @info "Skipping JET on Julia $VERSION (requires >=1.10 and a release build)" From 12eae700066b37095a7c4e4728f1bdd46a154f85 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 11 May 2026 17:00:56 -0400 Subject: [PATCH 09/20] Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/allocations.jl | 132 ++++++++++++----------------- test/glm/gaussian.jl | 32 +++---- test/glm/test_bernoulli_poisson.jl | 49 +++++------ 3 files changed, 95 insertions(+), 118 deletions(-) diff --git a/test/allocations.jl b/test/allocations.jl index 07e1d5b..7c5a885 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -22,72 +22,48 @@ using StatsAPI does NOT exist in real loops — measure inside a function. =# -bench_logd(d, y, x, n) = ( - s=0.0; - for _ in 1:n - ; - s += logdensityof(d, y; control_seq=x); - end; - s -) -bench_logd_unctrl(d, y, n) = ( - s=0.0; - for _ in 1:n - ; - s += logdensityof(d, y); - end; - s -) -bench_rand_scalar(rng, d, x, n) = ( - s=0.0; - for _ in 1:n - ; - s += rand(rng, d; control_seq=x); - end; - s -) -bench_rand_int(rng, d, x, n) = ( - s=0; - for _ in 1:n - ; - s += rand(rng, d; control_seq=x); - end; - s -) -bench_rand!_v(rng, d, out, x, n) = ( - s=0.0; - for _ in 1:n - ; - rand!(rng, d, out; control_seq=x); - s += out[1]; - end; - s -) -bench_rand!_i(rng, d, out, x, n) = ( - s=0; - for _ in 1:n - ; - rand!(rng, d, out; control_seq=x); - s += out[1]; - end; - s -) -bench_rand_unctrl_scalar(rng, d, n) = ( - s=0.0; - for _ in 1:n - ; - s += rand(rng, d); - end; - s -) -bench_rand_unctrl_vec(rng, d, n) = ( - s=0.0; - for _ in 1:n - ; - s += rand(rng, d)[1]; - end; - s -) +bench_logd(d, y, x, n) = (s = 0.0; +for _ in 1:n + s += logdensityof(d, y; control_seq=x) +end; +s) +bench_logd_unctrl(d, y, n) = (s = 0.0; +for _ in 1:n + s += logdensityof(d, y) +end; +s) +bench_rand_scalar(rng, d, x, n) = (s = 0.0; +for _ in 1:n + s += rand(rng, d; control_seq=x) +end; +s) +bench_rand_int(rng, d, x, n) = (s = 0; +for _ in 1:n + s += rand(rng, d; control_seq=x) +end; +s) +bench_rand!_v(rng, d, out, x, n) = (s = 0.0; +for _ in 1:n + rand!(rng, d, out; control_seq=x) + s += out[1] +end; +s) +bench_rand!_i(rng, d, out, x, n) = (s = 0; +for _ in 1:n + rand!(rng, d, out; control_seq=x) + s += out[1] +end; +s) +bench_rand_unctrl_scalar(rng, d, n) = (s = 0.0; +for _ in 1:n + s += rand(rng, d) +end; +s) +bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; +for _ in 1:n + s += rand(rng, d)[1] +end; +s) @testset "Allocations (steady state)" begin rng = Random.MersenneTwister(0) @@ -102,15 +78,15 @@ bench_rand_unctrl_vec(rng, d, n) = ( mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) - yv = [0.1, 0.2]; + yv = [0.1, 0.2] yi = [0, 1] # Warm - bench_logd(g, 0.5, x, 1); - bench_logd(b, 1, x, 1); + bench_logd(g, 0.5, x, 1) + bench_logd(b, 1, x, 1) bench_logd(p, 2, x, 1) - bench_logd(mg, yv, x, 1); - bench_logd(mb, yi, x, 1); + bench_logd(mg, yv, x, 1) + bench_logd(mb, yi, x, 1) bench_logd(mp, yi, x, 1) # Truly zero-alloc (no scratch needed) @@ -161,7 +137,7 @@ bench_rand_unctrl_vec(rng, d, n) = ( # GaussianGLM: closed-form WLS. Workspace: XWX (p²) + XWy (p). yg = randn(rng, n) - gg = GaussianGLM([0.0, 0.0], 1.0); + gg = GaussianGLM([0.0, 0.0], 1.0) fit!(gg, yg, w; control_seq=X) gg = GaussianGLM([0.0, 0.0], 1.0) @test (@allocated fit!(gg, yg, w; control_seq=X)) ≤ 1_000 @@ -177,13 +153,13 @@ bench_rand_unctrl_vec(rng, d, n) = ( # BernoulliGLM/PoissonGLM: hand-rolled Newton. # Workspace: g, H, Δ — three small alloc, total ~300 bytes for p=2. yb = Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:n] - gb = BernoulliGLM(zeros(2)); + gb = BernoulliGLM(zeros(2)) fit!(gb, yb, w; control_seq=X) gb = BernoulliGLM(zeros(2)) @test (@allocated fit!(gb, yb, w; control_seq=X)) ≤ 1_000 yp = Int[rand(rng, 0:5) for _ in 1:n] - gp = PoissonGLM(zeros(2)); + gp = PoissonGLM(zeros(2)) fit!(gp, yp, w; control_seq=X) gp = PoissonGLM(zeros(2)) @test (@allocated fit!(gp, yp, w; control_seq=X)) ≤ 1_000 @@ -191,13 +167,13 @@ bench_rand_unctrl_vec(rng, d, n) = ( # Multivariate Newton: workspace shared across columns. _ColumnElementView # avoids the n-sized per-column copy that Vector-of-Vectors would force. ymb = [Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:2] for _ in 1:n] - gmb = MvBernoulliGLM(zeros(2, 2)); + gmb = MvBernoulliGLM(zeros(2, 2)) fit!(gmb, ymb, w; control_seq=X) gmb = MvBernoulliGLM(zeros(2, 2)) @test (@allocated fit!(gmb, ymb, w; control_seq=X)) ≤ 1_000 ymp = [Int[rand(rng, 0:3) for _ in 1:2] for _ in 1:n] - gmp = MvPoissonGLM(zeros(2, 2)); + gmp = MvPoissonGLM(zeros(2, 2)) fit!(gmp, ymp, w; control_seq=X) gmp = MvPoissonGLM(zeros(2, 2)) @test (@allocated fit!(gmp, ymp, w; control_seq=X)) ≤ 1_000 @@ -206,7 +182,7 @@ bench_rand_unctrl_vec(rng, d, n) = ( @testset "PoissonZeroInflated" begin zip = PoissonZeroInflated(3.0, 0.2) - bench_logd_unctrl(zip, 0, 1); + bench_logd_unctrl(zip, 0, 1) bench_logd_unctrl(zip, 5, 1) @test (@allocated bench_logd_unctrl(zip, 0, REPS)) == 0 @test (@allocated bench_logd_unctrl(zip, 5, REPS)) == 0 @@ -217,7 +193,7 @@ bench_rand_unctrl_vec(rng, d, n) = ( n = 500 y = [rand(rng, zip) for _ in 1:n] w = ones(n) - zip2 = PoissonZeroInflated(1.0, 0.1); + zip2 = PoissonZeroInflated(1.0, 0.1) fit!(zip2, y, w) zip2 = PoissonZeroInflated(1.0, 0.1) @test (@allocated fit!(zip2, y, w)) ≤ 1_000 diff --git a/test/glm/gaussian.jl b/test/glm/gaussian.jl index d58f1c1..0564bf0 100644 --- a/test/glm/gaussian.jl +++ b/test/glm/gaussian.jl @@ -13,7 +13,7 @@ function _synthetic_gaussian_glm( rng, n::Int, p::Int; β::Vector{Float64}, σ2::Float64, weights=:uniform ) @assert length(β) == p - X = hcat(ones(n), randn(rng, n, p-1)) # intercept + random features + X = hcat(ones(n), randn(rng, n, p - 1)) # intercept + random features μ = X * β y = rand.(Ref(rng), Normal.(μ, sqrt(σ2))) w = if weights === :uniform @@ -56,7 +56,7 @@ _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) @test isfinite(logp) expected = _expected_logpdf(y[i], dot(glm.β, x_i), glm.σ2) - @test logp ≈ expected rtol=1e-12 atol=0.0 + @test logp ≈ expected rtol = 1e-12 atol = 0.0 end # Higher density near the mean than far away @@ -78,11 +78,11 @@ _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) # Empirical mean should be close to dot(β,x) μ = dot(glm.β, x) m = mean(samples) - @test m ≈ μ atol=0.15 + @test m ≈ μ atol = 0.15 # Empirical variance should be close to σ2 v = var(samples) - @test v ≈ glm.σ2 atol=0.25 + @test v ≈ glm.σ2 atol = 0.25 end @testset "fit! with uniform weights" begin @@ -96,8 +96,8 @@ _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) fit!(glm, y, w; control_seq=X) # Coefficients should be close with enough data - @test glm.β ≈ β_true atol=0.08 - @test glm.σ2 ≈ σ2_true atol=0.10 + @test glm.β ≈ β_true atol = 0.08 + @test glm.σ2 ≈ σ2_true atol = 0.10 @test all(isfinite, glm.β) @test isfinite(glm.σ2) @@ -114,8 +114,8 @@ _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) glm = GaussianGLM([0.0, 0.0], 1.0) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.10 - @test glm.σ2 ≈ σ2_true atol=0.15 + @test glm.β ≈ β_true atol = 0.10 + @test glm.σ2 ≈ σ2_true atol = 0.15 end @testset "Constructor with prior" begin @@ -148,7 +148,7 @@ function _synthetic_mvgaussian_glm( ) @assert size(B) == (p, k) @assert size(Σ) == (k, k) - X = hcat(ones(n), randn(rng, n, p-1)) + X = hcat(ones(n), randn(rng, n, p - 1)) L = cholesky(Σ).L obs_seq = Vector{Vector{Float64}}(undef, n) for i in 1:n @@ -205,7 +205,7 @@ end # Compare against MvNormal closed form diff = y - μ expected = -log(2π) - 0.5 * logdet(Σ) - 0.5 * dot(diff, Σ \ diff) - @test logdensityof(glm, y; control_seq=x) ≈ expected rtol=1e-10 + @test logdensityof(glm, y; control_seq=x) ≈ expected rtol = 1e-10 # Higher density at the mean than far from it @test logdensityof(glm, μ; control_seq=x) > @@ -227,10 +227,10 @@ end μ_true = vec(B' * x) m = vec(mean(Y; dims=1)) - @test m ≈ μ_true atol=0.05 + @test m ≈ μ_true atol = 0.05 S = (Y .- m')' * (Y .- m') / (n - 1) - @test S ≈ Σ atol=0.1 + @test S ≈ Σ atol = 0.1 end @testset "fit! recovers B and Σ" begin @@ -241,8 +241,8 @@ end glm = MvGaussianGLM(zeros(3, 2), Matrix(1.0I, 2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.08 - @test glm.Σ ≈ Σ_true atol=0.10 + @test glm.B ≈ B_true atol = 0.08 + @test glm.Σ ≈ Σ_true atol = 0.10 @test all(isfinite, glm.B) @test all(isfinite, glm.Σ) @test isposdef(glm.Σ) @@ -258,8 +258,8 @@ end glm = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.10 - @test glm.Σ ≈ Σ_true atol=0.12 + @test glm.B ≈ B_true atol = 0.10 + @test glm.Σ ≈ Σ_true atol = 0.12 end @testset "fit! with RidgePrior shrinks toward zero" begin diff --git a/test/glm/test_bernoulli_poisson.jl b/test/glm/test_bernoulli_poisson.jl index f481438..e3722e5 100644 --- a/test/glm/test_bernoulli_poisson.jl +++ b/test/glm/test_bernoulli_poisson.jl @@ -91,8 +91,8 @@ end η = dot(glm.β, x) μ = _sigmoid(η) - @test logdensityof(glm, 1; control_seq=x) ≈ log(μ) rtol=1e-10 - @test logdensityof(glm, 0; control_seq=x) ≈ log(1 - μ) rtol=1e-10 + @test logdensityof(glm, 1; control_seq=x) ≈ log(μ) rtol = 1e-10 + @test logdensityof(glm, 0; control_seq=x) ≈ log(1 - μ) rtol = 1e-10 @test logdensityof(glm, 2; control_seq=x) == -Inf @test logdensityof(glm, -1; control_seq=x) == -Inf end @@ -115,7 +115,7 @@ end samples = [rand(rng, glm; control_seq=x) for _ in 1:n] @test all(s -> s == 0 || s == 1, samples) - @test mean(samples) ≈ _sigmoid(dot(glm.β, x)) atol=0.05 + @test mean(samples) ≈ _sigmoid(dot(glm.β, x)) atol = 0.05 end @testset "fit! uniform weights" begin @@ -125,7 +125,7 @@ end glm = BernoulliGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.15 + @test glm.β ≈ β_true atol = 0.15 @test all(isfinite, glm.β) end @@ -136,7 +136,7 @@ end glm = BernoulliGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.20 + @test glm.β ≈ β_true atol = 0.20 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -199,7 +199,7 @@ end μ = exp(η) for k in 0:5 - @test logdensityof(glm, k; control_seq=x) ≈ logpdf(Poisson(μ), k) rtol=1e-10 + @test logdensityof(glm, k; control_seq=x) ≈ logpdf(Poisson(μ), k) rtol = 1e-10 end @test logdensityof(glm, -1; control_seq=x) == -Inf @@ -212,7 +212,7 @@ end samples = [rand(rng, glm; control_seq=x) for _ in 1:n] @test all(s -> s >= 0, samples) - @test mean(samples) ≈ exp(glm.β[1]) atol=0.2 + @test mean(samples) ≈ exp(glm.β[1]) atol = 0.2 end @testset "fit! uniform weights" begin @@ -222,7 +222,7 @@ end glm = PoissonGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.15 + @test glm.β ≈ β_true atol = 0.15 @test all(isfinite, glm.β) end @@ -233,7 +233,7 @@ end glm = PoissonGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.15 + @test glm.β ≈ β_true atol = 0.15 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -326,10 +326,11 @@ end η2 = dot(B[:, 2], x) μ1, μ2 = _sigmoid(η1), _sigmoid(η2) - @test logdensityof(glm, [1, 1]; control_seq=x) ≈ log(μ1) + log(μ2) rtol=1e-10 - @test logdensityof(glm, [1, 0]; control_seq=x) ≈ log(μ1) + log(1 - μ2) rtol=1e-10 - @test logdensityof(glm, [0, 1]; control_seq=x) ≈ log(1 - μ1) + log(μ2) rtol=1e-10 - @test logdensityof(glm, [0, 0]; control_seq=x) ≈ log(1 - μ1) + log(1 - μ2) rtol=1e-10 + @test logdensityof(glm, [1, 1]; control_seq=x) ≈ log(μ1) + log(μ2) rtol = 1e-10 + @test logdensityof(glm, [1, 0]; control_seq=x) ≈ log(μ1) + log(1 - μ2) rtol = 1e-10 + @test logdensityof(glm, [0, 1]; control_seq=x) ≈ log(1 - μ1) + log(μ2) rtol = 1e-10 + @test logdensityof(glm, [0, 0]; control_seq=x) ≈ log(1 - μ1) + log(1 - μ2) rtol = + 1e-10 @test logdensityof(glm, [2, 0]; control_seq=x) == -Inf end @@ -350,8 +351,8 @@ end Y = reduce(hcat, samples)' # n × k @test all(s -> s == 0 || s == 1, vec(Y)) - @test mean(Y[:, 1]) ≈ _sigmoid(dot(B[:, 1], x)) atol=0.05 - @test mean(Y[:, 2]) ≈ _sigmoid(dot(B[:, 2], x)) atol=0.05 + @test mean(Y[:, 1]) ≈ _sigmoid(dot(B[:, 1], x)) atol = 0.05 + @test mean(Y[:, 2]) ≈ _sigmoid(dot(B[:, 2], x)) atol = 0.05 end @testset "fit! recovers B" begin @@ -361,7 +362,7 @@ end glm = MvBernoulliGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.20 + @test glm.B ≈ B_true atol = 0.20 @test all(isfinite, glm.B) end @@ -372,7 +373,7 @@ end glm = MvBernoulliGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.25 + @test glm.B ≈ B_true atol = 0.25 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -400,7 +401,7 @@ end y_j = [obs_seq[i][j] for i in eachindex(obs_seq)] glm_j = BernoulliGLM(zeros(2)) fit!(glm_j, y_j, w; control_seq=X) - @test glm.B[:, j] ≈ glm_j.β rtol=1e-6 + @test glm.B[:, j] ≈ glm_j.β rtol = 1e-6 end end @@ -448,7 +449,7 @@ end for k1 in 0:3, k2 in 0:3 expected = logpdf(Poisson(μ1), k1) + logpdf(Poisson(μ2), k2) - @test logdensityof(glm, [k1, k2]; control_seq=x) ≈ expected rtol=1e-10 + @test logdensityof(glm, [k1, k2]; control_seq=x) ≈ expected rtol = 1e-10 end @test logdensityof(glm, [-1, 0]; control_seq=x) == -Inf @@ -465,8 +466,8 @@ end Y = reduce(hcat, samples)' @test all(s -> s >= 0, vec(Y)) - @test mean(Y[:, 1]) ≈ exp(B[1, 1]) atol=0.2 - @test mean(Y[:, 2]) ≈ exp(B[1, 2]) atol=0.2 + @test mean(Y[:, 1]) ≈ exp(B[1, 1]) atol = 0.2 + @test mean(Y[:, 2]) ≈ exp(B[1, 2]) atol = 0.2 end @testset "fit! recovers B" begin @@ -476,7 +477,7 @@ end glm = MvPoissonGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.15 + @test glm.B ≈ B_true atol = 0.15 @test all(isfinite, glm.B) end @@ -487,7 +488,7 @@ end glm = MvPoissonGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.20 + @test glm.B ≈ B_true atol = 0.20 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -515,7 +516,7 @@ end y_j = [obs_seq[i][j] for i in eachindex(obs_seq)] glm_j = PoissonGLM(zeros(2)) fit!(glm_j, y_j, w; control_seq=X) - @test glm.B[:, j] ≈ glm_j.β rtol=1e-6 + @test glm.B[:, j] ≈ glm_j.β rtol = 1e-6 end end From 4515a4da196dbe7ed79c36eaedd3fc8761d2f87d Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 11 May 2026 17:11:33 -0400 Subject: [PATCH 10/20] Apply suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/glms/glm.jl | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 237eeea..4bc4111 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -761,7 +761,7 @@ 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] @@ -889,11 +889,11 @@ function DensityInterface.logdensityof( Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) lp = zero(Tη) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) yj = y[j] (yj == 0 || yj == 1) || return Tη(-Inf) η = zero(Tη) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end lp += yj == 1 ? -log1pexp(-η) : -log1pexp(η) @@ -932,7 +932,7 @@ function Random.rand!( T = eltype(glm.B) for j in 1:glm.out_dim η = zero(T) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end out[j] = rand(rng) < logistic(η) ? 1 : 0 @@ -1051,11 +1051,11 @@ function DensityInterface.logdensityof( Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) lp = zero(Tη) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) yj = y[j] yj >= 0 || return Tη(-Inf) η = zero(Tη) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end lp += yj * η - exp(η) - logfactorial(yj) @@ -1091,9 +1091,9 @@ function Random.rand!( ) T = eltype(glm.B) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) η = zero(T) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end out[j] = rand(rng, Poisson(exp(η))) From e227109da3095f6e6e4b19b699276dbac30783d0 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 11 May 2026 17:12:05 -0400 Subject: [PATCH 11/20] Update src/glms/glm.jl Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- src/glms/glm.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 4bc4111..952ab22 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -930,7 +930,7 @@ function Random.rand!( ) T = eltype(glm.B) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) η = zero(T) for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] From 863f31d5197b513eb3e97f6685db69a585355cc8 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 11 May 2026 17:14:22 -0400 Subject: [PATCH 12/20] Jet --- src/glms/glm.jl | 2 +- test/runtests.jl | 8 ++++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 237eeea..76bde0c 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -151,7 +151,7 @@ function StatsAPI.fit!( end end - #= RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. =# + # RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. neglogprior_hess!(reg.prior, XWX, reg.β) F = cholesky!(Symmetric(XWX)) diff --git a/test/runtests.jl b/test/runtests.jl index 2593bf3..4e8fd29 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,13 +5,17 @@ using StatsAPI @testset "EmissionModels.jl" begin @testset "Code linting" begin - if VERSION >= v"1.10" && isempty(VERSION.prerelease) + # JET ≤ 0.10.x crashes inside Base._which on Julia ≥ 1.12 (reflection + # API change). Until JET ships a 1.12-compatible release we skip the + # linting test on 1.12+. + jet_ok = VERSION >= v"1.10" && VERSION < v"1.12" && isempty(VERSION.prerelease) + if jet_ok using Pkg Pkg.add("JET") using JET JET.test_package(EmissionModels; target_modules=(EmissionModels,)) else - @info "Skipping JET on Julia $VERSION (requires >=1.10 and a release build)" + @info "Skipping JET on Julia $VERSION (requires 1.10 ≤ v < 1.12)" end end From ada77bf6405bdaef50e37b57d57329452f8cba01 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Mon, 11 May 2026 17:14:23 -0400 Subject: [PATCH 13/20] c --- test/runtests.jl | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/test/runtests.jl b/test/runtests.jl index 4e8fd29..badb822 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,11 +5,7 @@ using StatsAPI @testset "EmissionModels.jl" begin @testset "Code linting" begin - # JET ≤ 0.10.x crashes inside Base._which on Julia ≥ 1.12 (reflection - # API change). Until JET ships a 1.12-compatible release we skip the - # linting test on 1.12+. - jet_ok = VERSION >= v"1.10" && VERSION < v"1.12" && isempty(VERSION.prerelease) - if jet_ok + if isempty(VERSION.prerelease) using Pkg Pkg.add("JET") using JET From 1d9bbc033e40220666476db6a90c041958648a4a Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 14 May 2026 10:48:19 -0400 Subject: [PATCH 14/20] attempt fix of juliaformatter action --- .github/workflows/Format-PR.yml | 8 --- .github/workflows/Format.yml | 12 +++- src/glms/glm.jl | 2 +- test/allocations.jl | 100 ++++++++++++++++++-------------- 4 files changed, 68 insertions(+), 54 deletions(-) delete mode 100644 .github/workflows/Format-PR.yml diff --git a/.github/workflows/Format-PR.yml b/.github/workflows/Format-PR.yml deleted file mode 100644 index 233c677..0000000 --- a/.github/workflows/Format-PR.yml +++ /dev/null @@ -1,8 +0,0 @@ -name: Format Suggestions -on: - pull_request: -jobs: - code-style: - runs-on: ubuntu-latest - steps: - - uses: julia-actions/julia-format@v4 \ No newline at end of file diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml index 9a95aa1..6e4f8bf 100644 --- a/.github/workflows/Format.yml +++ b/.github/workflows/Format.yml @@ -2,10 +2,16 @@ name: Formatter on: schedule: - cron: '0 0 * * *' + push: branches: - - 'main' - tags: ['*'] + - main + tags: + - '*' + + pull_request: + branches: + - main permissions: contents: write @@ -29,7 +35,7 @@ jobs: shell: julia --color=yes {0} run: | using Pkg - Pkg.add("JuliaFormatter") + Pkg.add(Pkg.PackageSpec(name = "JuliaFormatter", version = "2")) using JuliaFormatter format(".") diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 1f755b6..cef2df0 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -761,7 +761,7 @@ 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] diff --git a/test/allocations.jl b/test/allocations.jl index 7c5a885..2636de1 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -22,48 +22,64 @@ using StatsAPI does NOT exist in real loops — measure inside a function. =# -bench_logd(d, y, x, n) = (s = 0.0; -for _ in 1:n - s += logdensityof(d, y; control_seq=x) -end; -s) -bench_logd_unctrl(d, y, n) = (s = 0.0; -for _ in 1:n - s += logdensityof(d, y) -end; -s) -bench_rand_scalar(rng, d, x, n) = (s = 0.0; -for _ in 1:n - s += rand(rng, d; control_seq=x) -end; -s) -bench_rand_int(rng, d, x, n) = (s = 0; -for _ in 1:n - s += rand(rng, d; control_seq=x) -end; -s) -bench_rand!_v(rng, d, out, x, n) = (s = 0.0; -for _ in 1:n - rand!(rng, d, out; control_seq=x) - s += out[1] -end; -s) -bench_rand!_i(rng, d, out, x, n) = (s = 0; -for _ in 1:n - rand!(rng, d, out; control_seq=x) - s += out[1] -end; -s) -bench_rand_unctrl_scalar(rng, d, n) = (s = 0.0; -for _ in 1:n - s += rand(rng, d) -end; -s) -bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; -for _ in 1:n - s += rand(rng, d)[1] -end; -s) +bench_logd(d, y, x, n) = ( + s=0.0; + for _ in 1:n + s += logdensityof(d, y; control_seq=x) + end; + s +) +bench_logd_unctrl(d, y, n) = ( + s=0.0; + for _ in 1:n + s += logdensityof(d, y) + end; + s +) +bench_rand_scalar(rng, d, x, n) = ( + s=0.0; + for _ in 1:n + s += rand(rng, d; control_seq=x) + end; + s +) +bench_rand_int(rng, d, x, n) = ( + s=0; + for _ in 1:n + s += rand(rng, d; control_seq=x) + end; + s +) +bench_rand!_v(rng, d, out, x, n) = ( + s=0.0; + for _ in 1:n + rand!(rng, d, out; control_seq=x) + s += out[1] + end; + s +) +bench_rand!_i(rng, d, out, x, n) = ( + s=0; + for _ in 1:n + rand!(rng, d, out; control_seq=x) + s += out[1] + end; + s +) +bench_rand_unctrl_scalar(rng, d, n) = ( + s=0.0; + for _ in 1:n + s += rand(rng, d) + end; + s +) +bench_rand_unctrl_vec(rng, d, n) = ( + s=0.0; + for _ in 1:n + s += rand(rng, d)[1] + end; + s +) @testset "Allocations (steady state)" begin rng = Random.MersenneTwister(0) From 9ede16cd307b704ea9edeacff5af367f66f8771b Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 14 May 2026 10:54:07 -0400 Subject: [PATCH 15/20] try 2 --- .github/workflows/Format.yml | 23 ++++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml index 6e4f8bf..1a5fe20 100644 --- a/.github/workflows/Format.yml +++ b/.github/workflows/Format.yml @@ -1,4 +1,5 @@ name: Formatter + on: schedule: - cron: '0 0 * * *' @@ -6,13 +7,13 @@ on: push: branches: - main - tags: - - '*' pull_request: branches: - main + workflow_dispatch: + permissions: contents: write pull-requests: write @@ -20,8 +21,11 @@ permissions: jobs: build: runs-on: ubuntu-latest + steps: - uses: actions/checkout@v6 + with: + ref: ${{ github.event_name == 'pull_request' && github.head_ref || github.ref_name }} - uses: julia-actions/setup-julia@v3 with: @@ -32,6 +36,7 @@ jobs: - uses: julia-actions/cache@v3 - name: Install JuliaFormatter and format + if: github.event_name != 'pull_request' shell: julia --color=yes {0} run: | using Pkg @@ -39,18 +44,30 @@ jobs: using JuliaFormatter format(".") + - name: Check formatting on PR + if: github.event_name == 'pull_request' + shell: julia --color=yes {0} + run: | + using Pkg + Pkg.add(Pkg.PackageSpec(name = "JuliaFormatter", version = "2")) + using JuliaFormatter + @assert format(".", overwrite = false) + - name: Create Pull Request id: cpr + if: github.event_name != 'pull_request' uses: peter-evans/create-pull-request@v8 with: token: ${{ github.token }} commit-message: Format .jl files title: 'Automatic JuliaFormatter.jl run' branch: auto-juliaformatter-pr + base: main delete-branch: true labels: formatting, automated pr, no changelog - name: Check outputs + if: github.event_name != 'pull_request' run: | echo "Pull Request Number - ${{ steps.cpr.outputs.pull-request-number }}" - echo "Pull Request URL - ${{ steps.cpr.outputs.pull-request-url }}" + echo "Pull Request URL - ${{ steps.cpr.outputs.pull-request-url }}" \ No newline at end of file From 1ede960c7888a7cf8cc4ab4ab0883d72db324562 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 14 May 2026 11:00:55 -0400 Subject: [PATCH 16/20] jet fix --- test/Project.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/Project.toml b/test/Project.toml index 880ac11..2ebe9b5 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -11,3 +11,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +JET = "c3a54625-4413-5d7c-8a35-3738d8913d89" + +[compat] +JET = "0.9, 0.10, 0.11" \ No newline at end of file From 4cd68babc48dc8ab9355f9b814b6a48168e061fc Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 14 May 2026 11:03:55 -0400 Subject: [PATCH 17/20] fix --- test/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Project.toml b/test/Project.toml index 2ebe9b5..92dcd25 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -11,7 +11,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -JET = "c3a54625-4413-5d7c-8a35-3738d8913d89" +JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" [compat] JET = "0.9, 0.10, 0.11" \ No newline at end of file From 76f726844ad3d065a9b4cb2cfc20d48023cde788 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 14 May 2026 11:10:03 -0400 Subject: [PATCH 18/20] fix --- test/Project.toml | 7 +------ test/runtests.jl | 2 -- 2 files changed, 1 insertion(+), 8 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index 92dcd25..42352bc 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,15 +3,10 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" HiddenMarkovModels = "84ca31d5-effc-45e0-bfda-5a68cd981f47" -JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" -JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" - -[compat] -JET = "0.9, 0.10, 0.11" \ No newline at end of file +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" \ No newline at end of file diff --git a/test/runtests.jl b/test/runtests.jl index badb822..eec4621 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,8 +10,6 @@ using StatsAPI Pkg.add("JET") using JET JET.test_package(EmissionModels; target_modules=(EmissionModels,)) - else - @info "Skipping JET on Julia $VERSION (requires 1.10 ≤ v < 1.12)" end end From 9d9b97781c72b4a37a931a5d919547a43283e7b2 Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Sun, 10 May 2026 13:21:00 -0400 Subject: [PATCH 19/20] Fix race condition in MvGaussianGLM and MultivariateT logdensityof Gate JET linting on supported Julia versions Avoid resolving JET on prerelease Julia versions where no compatible registered release is available. Install JET only during supported lint runs so normal package tests can run on newer Julia prereleases. --- .github/workflows/Format-PR.yml | 8 - .github/workflows/Format.yml | 31 ++- docs/Project.toml | 2 + docs/make.jl | 3 + docs/src/distributions.md | 17 ++ docs/src/glm.md | 50 +++- docs/src/priors.md | 8 + src/glms/glm.jl | 404 ++++++++++++++++------------- src/multivariate/t.jl | 11 +- test/Project.toml | 8 +- test/allocations.jl | 157 +++++++---- test/glm/gaussian.jl | 234 +++++++++-------- test/glm/test_bernoulli_poisson.jl | 68 ++--- test/runtests.jl | 9 +- 14 files changed, 615 insertions(+), 395 deletions(-) delete mode 100644 .github/workflows/Format-PR.yml diff --git a/.github/workflows/Format-PR.yml b/.github/workflows/Format-PR.yml deleted file mode 100644 index 233c677..0000000 --- a/.github/workflows/Format-PR.yml +++ /dev/null @@ -1,8 +0,0 @@ -name: Format Suggestions -on: - pull_request: -jobs: - code-style: - runs-on: ubuntu-latest - steps: - - uses: julia-actions/julia-format@v4 \ No newline at end of file diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml index 9a95aa1..1a5fe20 100644 --- a/.github/workflows/Format.yml +++ b/.github/workflows/Format.yml @@ -1,11 +1,18 @@ name: Formatter + on: schedule: - cron: '0 0 * * *' + push: branches: - - 'main' - tags: ['*'] + - main + + pull_request: + branches: + - main + + workflow_dispatch: permissions: contents: write @@ -14,8 +21,11 @@ permissions: jobs: build: runs-on: ubuntu-latest + steps: - uses: actions/checkout@v6 + with: + ref: ${{ github.event_name == 'pull_request' && github.head_ref || github.ref_name }} - uses: julia-actions/setup-julia@v3 with: @@ -26,25 +36,38 @@ jobs: - uses: julia-actions/cache@v3 - name: Install JuliaFormatter and format + if: github.event_name != 'pull_request' shell: julia --color=yes {0} run: | using Pkg - Pkg.add("JuliaFormatter") + Pkg.add(Pkg.PackageSpec(name = "JuliaFormatter", version = "2")) using JuliaFormatter format(".") + - name: Check formatting on PR + if: github.event_name == 'pull_request' + shell: julia --color=yes {0} + run: | + using Pkg + Pkg.add(Pkg.PackageSpec(name = "JuliaFormatter", version = "2")) + using JuliaFormatter + @assert format(".", overwrite = false) + - name: Create Pull Request id: cpr + if: github.event_name != 'pull_request' uses: peter-evans/create-pull-request@v8 with: token: ${{ github.token }} commit-message: Format .jl files title: 'Automatic JuliaFormatter.jl run' branch: auto-juliaformatter-pr + base: main delete-branch: true labels: formatting, automated pr, no changelog - name: Check outputs + if: github.event_name != 'pull_request' run: | echo "Pull Request Number - ${{ steps.cpr.outputs.pull-request-number }}" - echo "Pull Request URL - ${{ steps.cpr.outputs.pull-request-url }}" + echo "Pull Request URL - ${{ steps.cpr.outputs.pull-request-url }}" \ No newline at end of file diff --git a/docs/Project.toml b/docs/Project.toml index c56c224..be32084 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,3 +1,5 @@ [deps] +DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" EmissionModels = "1e2dd27c-41a5-43b2-863c-3eddd0c72c67" +StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" diff --git a/docs/make.jl b/docs/make.jl index d6cac3d..f5bcbd8 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,5 +1,8 @@ using EmissionModels using Documenter +using DensityInterface +using Random +using StatsAPI DocMeta.setdocmeta!(EmissionModels, :DocTestSetup, :(using EmissionModels); recursive=true) diff --git a/docs/src/distributions.md b/docs/src/distributions.md index 6bbca5c..b6d1782 100644 --- a/docs/src/distributions.md +++ b/docs/src/distributions.md @@ -81,3 +81,20 @@ fit!(dist, obs_seq, weight_seq; max_iter=100, tol=1e-6, fix_nu=false) - The M-step updates ``μ`` and ``Σ`` (and optionally ``ν``). - Pass `fix_nu=true` to keep the degrees of freedom fixed during fitting. - `μ` and ``Σ`` are re-regularized on-the-fly if the Cholesky factorization fails. + +## API Reference + +```@docs +PoissonZeroInflated +MultivariateT +MultivariateTDiag +DensityInterface.logdensityof(::PoissonZeroInflated, ::Integer) +DensityInterface.logdensityof(::MultivariateT, ::AbstractVector) +DensityInterface.logdensityof(::MultivariateTDiag, ::AbstractVector) +Base.rand(::Random.AbstractRNG, ::PoissonZeroInflated) +Base.rand(::Random.AbstractRNG, ::MultivariateT) +Base.rand(::Random.AbstractRNG, ::MultivariateTDiag) +StatsAPI.fit!(::PoissonZeroInflated, ::AbstractVector, ::AbstractVector) +StatsAPI.fit!(::MultivariateT, ::Any, ::Any) +StatsAPI.fit!(::MultivariateTDiag, ::Any, ::Any) +``` diff --git a/docs/src/glm.md b/docs/src/glm.md index 9ad53c4..1f74600 100644 --- a/docs/src/glm.md +++ b/docs/src/glm.md @@ -42,6 +42,28 @@ Log-linear regression for count observations ``Y ∈ \{0, 1, \ldots\}``. glm = PoissonGLM(zeros(3)) ``` +## Multivariate GLM types + +Each univariate GLM has a multivariate counterpart that emits a length-``k`` observation vector for a single input ``x``. Coefficients are stored as a ``p × k`` matrix ``B``; column ``j`` is the coefficient vector for output dimension ``j``. + +### `MvGaussianGLM(B, Σ)` + +Multivariate linear regression with shared full covariance ``Σ``. + +``Y \mid x \sim \mathcal{N}(Bᵀx, \ Σ)`` + +### `MvBernoulliGLM(B)` + +``k`` independent logistic regressions sharing the same input ``x``. + +``P(Y \mid x) = \prod_{j=1}^{k} \text{Bernoulli}(σ(B[:,j]ᵀx))`` + +### `MvPoissonGLM(B)` + +``k`` independent Poisson log-linear regressions sharing the same input ``x``. + +``P(Y \mid x) = \prod_{j=1}^{k} \text{Poisson}(\exp(B[:,j]ᵀx))`` + ## Fitting GLM emissions GLM emissions are fit by minimizing the negative log-posterior (or weighted negative log-likelihood when no prior is used): @@ -49,5 +71,31 @@ GLM emissions are fit by minimizing the negative log-posterior (or weighted nega ``\ell(β) = -\sum_{i=1}^{n} w_i \, \log p(y_i \mid x_i, β) + \text{neglogprior}(β)`` ```julia -fit!(glm, y_seq, w_seq; control_seq=X, optim_opts=Optim.Options()) +# Gaussian is closed-form weighted least squares. +fit!(glm, y_seq, w_seq; control_seq=X) + +# Bernoulli / Poisson use a hand-rolled Newton with backtracking line search. +fit!(glm, y_seq, w_seq; control_seq=X, + max_iter=50, gtol=1e-8, max_backtrack=20) +``` + +## API Reference + +```@docs +EmissionModels.AbstractGLM +GaussianGLM +BernoulliGLM +PoissonGLM +MvGaussianGLM +MvBernoulliGLM +MvPoissonGLM +DensityInterface.logdensityof(::MvGaussianGLM, ::AbstractVector) +StatsAPI.fit!(::BernoulliGLM, ::AbstractVector, ::AbstractVector{<:Real}) +StatsAPI.fit!(::PoissonGLM, ::AbstractVector, ::AbstractVector{<:Real}) +StatsAPI.fit!(::MvGaussianGLM{T}, ::AbstractVector{<:AbstractVector}, ::AbstractVector{<:Real}) where {T<:Real} +StatsAPI.fit!(::MvBernoulliGLM{T}, ::AbstractVector{<:AbstractVector}, ::AbstractVector{<:Real}) where {T<:Real} +StatsAPI.fit!(::MvPoissonGLM{T}, ::AbstractVector{<:AbstractVector}, ::AbstractVector{<:Real}) where {T<:Real} +Random.rand!(::Random.AbstractRNG, ::MvGaussianGLM{T}, ::AbstractVector) where {T<:Real} +Random.rand!(::Random.AbstractRNG, ::MvBernoulliGLM, ::AbstractVector) +Random.rand!(::Random.AbstractRNG, ::MvPoissonGLM, ::AbstractVector) ``` diff --git a/docs/src/priors.md b/docs/src/priors.md index 458916d..4b2316e 100644 --- a/docs/src/priors.md +++ b/docs/src/priors.md @@ -71,3 +71,11 @@ glm = GaussianGLM(zeros(5), 1.0, prior) ``` The composed prior simply sums the individual penalties, gradients, and Hessians. + +## API Reference + +```@docs +AbstractPrior +NoPrior +RidgePrior +``` diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 4e04498..cef2df0 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -82,33 +82,39 @@ exact MAP solution for a Gaussian prior β ~ N(0, (1/λ)I). - `σ2`: noise variance - `prior`: regularization prior (default `NoPrior()`) """ -mutable struct GaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct GaussianGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} σ2::T prior::P end -GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:Real} = - GaussianGLM{T, NoPrior}(Vector{T}(β), σ2, NoPrior()) +function GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:Real} + return GaussianGLM{T,NoPrior}(Vector{T}(β), σ2, NoPrior()) +end -GaussianGLM(β::AbstractVector{T}, σ2::T, prior::P) where {T<:Real, P<:AbstractPrior} = - GaussianGLM{T, P}(Vector{T}(β), σ2, prior) +function GaussianGLM(β::AbstractVector{T}, σ2::T, prior::P) where {T<:Real,P<:AbstractPrior} + return GaussianGLM{T,P}(Vector{T}(β), σ2, prior) +end # Convenience: promote β eltype and σ2 together -GaussianGLM(β::AbstractVector, σ2::Real) = - GaussianGLM(float.(β), float(σ2)) +GaussianGLM(β::AbstractVector, σ2::Real) = GaussianGLM(float.(β), float(σ2)) -GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior) = - GaussianGLM(float.(β), float(σ2), prior) +function GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior) + return GaussianGLM(float.(β), float(σ2), prior) +end DensityInterface.DensityKind(::GaussianGLM) = DensityInterface.HasDensity() -function DensityInterface.logdensityof(reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real}) +function DensityInterface.logdensityof( + reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real} +) μ = dot(reg.β, control_seq) return -0.5 * log(2π * reg.σ2) - 0.5 * ((y - μ)^2 / reg.σ2) end -function Random.rand(rng::AbstractRNG, reg::GaussianGLM; control_seq::AbstractVector{<:Real}) +function Random.rand( + rng::AbstractRNG, reg::GaussianGLM; control_seq::AbstractVector{<:Real} +) return rand(rng, Normal(dot(reg.β, control_seq), sqrt(reg.σ2))) end @@ -131,12 +137,12 @@ function StatsAPI.fit!( wsum = zero(T) for i in 1:n - w = T(weights[i]) + w = T(weights[i]) wsum += w x_i = view(control_seq, i, :) y_i = T(obs_seq[i]) for a in 1:p - xa = x_i[a] + xa = x_i[a] wxa = w * xa for b in 1:p XWX[a, b] += wxa * x_i[b] @@ -145,7 +151,7 @@ function StatsAPI.fit!( end end - #= RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. =# + # RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. neglogprior_hess!(reg.prior, XWX, reg.β) F = cholesky!(Symmetric(XWX)) @@ -175,8 +181,9 @@ struct _ColumnElementView{T,V<:AbstractVector{<:AbstractVector}} <: AbstractVect seq::V j::Int end -_ColumnElementView(seq::V, j::Int) where {V<:AbstractVector{<:AbstractVector}} = - _ColumnElementView{eltype(eltype(V)), V}(seq, j) +function _ColumnElementView(seq::V, j::Int) where {V<:AbstractVector{<:AbstractVector}} + return _ColumnElementView{eltype(eltype(V)),V}(seq, j) +end Base.size(c::_ColumnElementView) = (length(c.seq),) Base.IndexStyle(::Type{<:_ColumnElementView}) = IndexLinear() Base.@propagate_inbounds Base.getindex(c::_ColumnElementView, i::Integer) = c.seq[i][c.j] @@ -212,7 +219,7 @@ function _bernoulli_gh!( fill!(H, zero(T)) for i in 1:n x_i = view(X, i, :) - wi = T(w[i]) + wi = T(w[i]) μ_i = logistic(dot(β, x_i)) r_i = wi * (μ_i - T(y[i])) W_i = wi * μ_i * (one(T) - μ_i) @@ -261,9 +268,9 @@ function _poisson_gh!( fill!(H, zero(T)) for i in 1:n x_i = view(X, i, :) - wi = T(w[i]) + wi = T(w[i]) η_i = clamp(dot(β, x_i), T(-500), T(500)) - eη = exp(η_i) + eη = exp(η_i) r_i = wi * (eη - T(y[i])) W_i = wi * eη for a in 1:p @@ -298,7 +305,7 @@ function _newton_solve!( gtol::Real=1e-8, max_backtrack::Int=20, ridge::Real=1e-12, -) where {T<:Real, F1, F2} +) where {T<:Real,F1,F2} p = length(β) f_curr = loss(β, y, w, X, prior) @@ -361,8 +368,9 @@ function _fit_bernoulli_glm!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) length(β) == p || throw(DimensionMismatch("β length $(length(β)) ≠ control_seq columns $p")) @@ -370,9 +378,19 @@ function _fit_bernoulli_glm!( H = Matrix{T}(undef, p, p) Δ = Vector{T}(undef, p) return _newton_solve!( - β, g, H, Δ, _bernoulli_loss, _bernoulli_gh!, - obs_seq, weight_seq, control_seq, prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β, + g, + H, + Δ, + _bernoulli_loss, + _bernoulli_gh!, + obs_seq, + weight_seq, + control_seq, + prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) end @@ -389,8 +407,9 @@ function _fit_poisson_glm!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) length(β) == p || throw(DimensionMismatch("β length $(length(β)) ≠ control_seq columns $p")) @@ -398,9 +417,19 @@ function _fit_poisson_glm!( H = Matrix{T}(undef, p, p) Δ = Vector{T}(undef, p) return _newton_solve!( - β, g, H, Δ, _poisson_loss, _poisson_gh!, - obs_seq, weight_seq, control_seq, prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β, + g, + H, + Δ, + _poisson_loss, + _poisson_gh!, + obs_seq, + weight_seq, + control_seq, + prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) end @@ -418,30 +447,32 @@ to the solver. - `β`: coefficient vector (length p) - `prior`: regularization prior (default `NoPrior()`) """ -mutable struct BernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct BernoulliGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} prior::P end -BernoulliGLM(β::AbstractVector{T}) where {T<:Real} = - BernoulliGLM{T, NoPrior}(Vector{T}(β), NoPrior()) +function BernoulliGLM(β::AbstractVector{T}) where {T<:Real} + return BernoulliGLM{T,NoPrior}(Vector{T}(β), NoPrior()) +end -BernoulliGLM(β::AbstractVector{T}, prior::P) where {T<:Real, P<:AbstractPrior} = - BernoulliGLM{T, P}(Vector{T}(β), prior) +function BernoulliGLM(β::AbstractVector{T}, prior::P) where {T<:Real,P<:AbstractPrior} + return BernoulliGLM{T,P}(Vector{T}(β), prior) +end DensityInterface.DensityKind(::BernoulliGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::BernoulliGLM, - y::Integer; - control_seq::AbstractVector{<:Real}, + glm::BernoulliGLM, y::Integer; control_seq::AbstractVector{<:Real} ) (y == 0 || y == 1) || return oftype(dot(glm.β, control_seq), -Inf) η = dot(glm.β, control_seq) return y == 1 ? -log1pexp(-η) : -log1pexp(η) end -function Random.rand(rng::AbstractRNG, glm::BernoulliGLM; control_seq::AbstractVector{<:Real}) +function Random.rand( + rng::AbstractRNG, glm::BernoulliGLM; control_seq::AbstractVector{<:Real} +) return rand(rng, Bernoulli(logistic(dot(glm.β, control_seq)))) end @@ -464,8 +495,14 @@ function StatsAPI.fit!( max_backtrack::Int=20, ) _fit_bernoulli_glm!( - glm.β, obs_seq, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + glm.β, + obs_seq, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) return glm end @@ -483,23 +520,23 @@ composes without changes to the solver. - `β`: coefficient vector (length p) - `prior`: regularization prior (default `NoPrior()`) """ -mutable struct PoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct PoissonGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} prior::P end -PoissonGLM(β::AbstractVector{T}) where {T<:Real} = - PoissonGLM{T, NoPrior}(Vector{T}(β), NoPrior()) +function PoissonGLM(β::AbstractVector{T}) where {T<:Real} + return PoissonGLM{T,NoPrior}(Vector{T}(β), NoPrior()) +end -PoissonGLM(β::AbstractVector{T}, prior::P) where {T<:Real, P<:AbstractPrior} = - PoissonGLM{T, P}(Vector{T}(β), prior) +function PoissonGLM(β::AbstractVector{T}, prior::P) where {T<:Real,P<:AbstractPrior} + return PoissonGLM{T,P}(Vector{T}(β), prior) +end DensityInterface.DensityKind(::PoissonGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::PoissonGLM, - y::Integer; - control_seq::AbstractVector{<:Real}, + glm::PoissonGLM, y::Integer; control_seq::AbstractVector{<:Real} ) y >= 0 || return oftype(dot(glm.β, control_seq), -Inf) η = dot(glm.β, control_seq) @@ -529,8 +566,14 @@ function StatsAPI.fit!( max_backtrack::Int=20, ) _fit_poisson_glm!( - glm.β, obs_seq, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + glm.β, + obs_seq, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) return glm end @@ -553,37 +596,26 @@ prior on `B` (Frobenius-norm penalty). - `Σ`: residual covariance of size `k × k` - `prior`: regularization prior on `B` (default `NoPrior()`) """ -#= Cholesky uplo: we always store the LOWER triangle. `getproperty(::Cholesky, :L)` - only returns a wrapper without copying when uplo == 'L'; with the upper-stored - default, every `.L` access does `copy(factors')`, costing ~k² floats per call. =# - -mutable struct MvGaussianGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct MvGaussianGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM B::Matrix{T} Σ::Matrix{T} prior::P - Σ_chol::Cholesky{T, Matrix{T}} + Σ_chol::Cholesky{T,Matrix{T}} logdetΣ::T in_dim::Int out_dim::Int - - #= Scratch for hot paths. Sequential use only — not safe for concurrent - calls on the same instance. HMM E-steps run forward/backward serially - per state, so each emission has its own scratch and this is fine. =# - _diff::Vector{T} - _z::Vector{T} end function MvGaussianGLM( - B::AbstractMatrix{T}, - Σ::AbstractMatrix{T}, - prior::P, -) where {T<:Real, P<:AbstractPrior} + B::AbstractMatrix{T}, Σ::AbstractMatrix{T}, prior::P +) where {T<:Real,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) - size(Σ) == (k, k) || - throw(DimensionMismatch("Σ must be $(k)×$(k) for B with $(k) columns, got $(size(Σ))")) + size(Σ) == (k, k) || throw( + DimensionMismatch("Σ must be $(k)×$(k) for B with $(k) columns, got $(size(Σ))") + ) Σ_chol = try cholesky(Symmetric(Σ, :L)) @@ -591,15 +623,14 @@ function MvGaussianGLM( throw(ArgumentError("Σ must be positive definite")) end - return MvGaussianGLM{T, P}( - Matrix{T}(B), Matrix{T}(Σ), prior, - Σ_chol, logdet(Σ_chol), p, k, - Vector{T}(undef, k), Vector{T}(undef, k), + return MvGaussianGLM{T,P}( + Matrix{T}(B), Matrix{T}(Σ), prior, Σ_chol, logdet(Σ_chol), p, k ) end -MvGaussianGLM(B::AbstractMatrix{T}, Σ::AbstractMatrix{T}) where {T<:Real} = - MvGaussianGLM(B, Σ, NoPrior()) +function MvGaussianGLM(B::AbstractMatrix{T}, Σ::AbstractMatrix{T}) where {T<:Real} + return MvGaussianGLM(B, Σ, NoPrior()) +end function MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix) T = promote_type(eltype(B), eltype(Σ)) @@ -618,25 +649,24 @@ DensityInterface.DensityKind(::MvGaussianGLM) = DensityInterface.HasDensity() """ logdensityof(glm::MvGaussianGLM, y::AbstractVector; control_seq) -Log density of `y ∈ ℝᵏ` under the conditional MvNormal model. Zero allocation: -the residual buffer `_diff` is pre-allocated in the struct. +Log density of `y ∈ ℝᵏ` under the conditional MvNormal model. Allocates one +length-`k` residual vector per call — thread-safe (matches `Distributions.MvNormal`). """ function DensityInterface.logdensityof( - glm::MvGaussianGLM, - y::AbstractVector; - control_seq::AbstractVector{<:Real}, + glm::MvGaussianGLM, y::AbstractVector; control_seq::AbstractVector{<:Real} ) length(y) == glm.out_dim || throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) k = glm.out_dim p = glm.in_dim T = eltype(glm.B) - diff = glm._diff + diff = Vector{T}(undef, k) #= μ = Bᵀ x and diff = y − μ, fused as a single loop to avoid the Adjoint*Vector temporary that BLAS would otherwise allocate. =# @@ -657,9 +687,7 @@ function DensityInterface.logdensityof( end function Random.rand( - rng::AbstractRNG, - glm::MvGaussianGLM{T}; - control_seq::AbstractVector{<:Real}, + rng::AbstractRNG, glm::MvGaussianGLM{T}; control_seq::AbstractVector{<:Real} ) where {T<:Real} out = Vector{T}(undef, glm.out_dim) rand!(rng, glm, out; control_seq=control_seq) @@ -669,8 +697,12 @@ end """ rand!(rng, glm::MvGaussianGLM, out; control_seq) -In-place sample. `out` must be a length-`out_dim` `AbstractVector{T}`. -Zero allocation. +In-place sample into `out` (length `out_dim`). Zero allocation, thread-safe: +the trick is to draw `z ~ N(0, I)` directly into `out`, multiply by `L` +in-place (`lmul!`), then add `μ = Bᵀ x` element-wise. Lower-triangular +multiply is well-defined in place because each output `xᵢ = Σⱼ≤ᵢ Lᵢⱼ zⱼ` +only reads `z[1..i]`, so iterating `i = k, k-1, …, 1` never reads an +already-overwritten entry. """ function Random.rand!( rng::AbstractRNG, @@ -680,26 +712,26 @@ function Random.rand!( ) where {T<:Real} length(out) == glm.out_dim || throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) k = glm.out_dim p = glm.in_dim - z = glm._z - randn!(rng, z) - #= out = Bᵀ x (manual loop avoids the Adjoint*Vector temporary) =# + randn!(rng, out) + lmul!(glm.Σ_chol.L, out) # out = L * z, in place + + # out += μ = Bᵀ x (Bᵀx is computed inline; no aliasing with out) for j in 1:k sj = zero(T) for r in 1:p sj += glm.B[r, j] * control_seq[r] end - out[j] = sj + out[j] += sj end - #= out += L z via in-place 5-arg mul! =# - mul!(out, glm.Σ_chol.L, z, true, true) return out end @@ -719,8 +751,9 @@ function StatsAPI.fit!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) p == glm.in_dim || throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) @@ -728,20 +761,17 @@ function StatsAPI.fit!( XWX = zeros(T, p, p) XWY = zeros(T, p, k) - wsum = zero(T) + wsum = zero(T)#= Build XᵀWX and XᵀWY in a single pass — no Y matrix, no temporaries. =# - #= 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 || - throw(DimensionMismatch( - "obs_seq[$i] length $(length(obs_i)) ≠ out_dim $k", - )) - w = T(weight_seq[i]) + throw(DimensionMismatch("obs_seq[$i] length $(length(obs_i)) ≠ out_dim $k")) + w = T(weight_seq[i]) wsum += w x_i = view(control_seq, i, :) for a in 1:p - xa = x_i[a] + xa = x_i[a] wxa = w * xa for b in 1:p XWX[a, b] += wxa * x_i[b] @@ -774,7 +804,7 @@ function StatsAPI.fit!( Σ_new = zeros(T, k, k) for i in 1:n obs_i = obs_seq[i] - w = T(weight_seq[i]) + w = T(weight_seq[i]) x_i = view(control_seq, i, :) for j in 1:k rj = T(obs_i[j]) @@ -824,48 +854,46 @@ For input x ∈ ℝᵖ the conditional distribution of y ∈ {0,1}ᵏ factorizes - `B`: coefficient matrix of size `p × k` - `prior`: regularization prior applied per column (default `NoPrior()`) """ -mutable struct MvBernoulliGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct MvBernoulliGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM B::Matrix{T} prior::P in_dim::Int out_dim::Int end -function MvBernoulliGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real, P<:AbstractPrior} +function MvBernoulliGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) - return MvBernoulliGLM{T, P}(Matrix{T}(B), prior, p, k) + return MvBernoulliGLM{T,P}(Matrix{T}(B), prior, p, k) end MvBernoulliGLM(B::AbstractMatrix{T}) where {T<:Real} = MvBernoulliGLM(B, NoPrior()) MvBernoulliGLM(B::AbstractMatrix) = MvBernoulliGLM(float.(B), NoPrior()) -MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior) = - MvBernoulliGLM(float.(B), prior) +MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior) = MvBernoulliGLM(float.(B), prior) DensityInterface.DensityKind(::MvBernoulliGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::MvBernoulliGLM, - y::AbstractVector; - control_seq::AbstractVector{<:Real}, + glm::MvBernoulliGLM, y::AbstractVector; control_seq::AbstractVector{<:Real} ) length(y) == glm.out_dim || throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) lp = zero(Tη) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) yj = y[j] (yj == 0 || yj == 1) || return Tη(-Inf) η = zero(Tη) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end lp += yj == 1 ? -log1pexp(-η) : -log1pexp(η) @@ -874,9 +902,7 @@ function DensityInterface.logdensityof( end function Random.rand( - rng::AbstractRNG, - glm::MvBernoulliGLM; - control_seq::AbstractVector{<:Real}, + rng::AbstractRNG, glm::MvBernoulliGLM; control_seq::AbstractVector{<:Real} ) out = Vector{Int}(undef, glm.out_dim) rand!(rng, glm, out; control_seq=control_seq) @@ -897,15 +923,16 @@ function Random.rand!( ) length(out) == glm.out_dim || throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) T = eltype(glm.B) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) η = zero(T) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end out[j] = rand(rng) < logistic(η) ? 1 : 0 @@ -933,31 +960,41 @@ function StatsAPI.fit!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) p == glm.in_dim || throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) k = glm.out_dim for i in 1:n - length(obs_seq[i]) == k || - throw(DimensionMismatch( - "obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k", - )) + length(obs_seq[i]) == k || throw( + DimensionMismatch("obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k") + ) end β_buf = Vector{T}(undef, p) - g = Vector{T}(undef, p) - H = Matrix{T}(undef, p, p) - Δ = Vector{T}(undef, p) + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) for j in 1:k yview = _ColumnElementView(obs_seq, j) copyto!(β_buf, view(glm.B, :, j)) _newton_solve!( - β_buf, g, H, Δ, _bernoulli_loss, _bernoulli_gh!, - yview, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β_buf, + g, + H, + Δ, + _bernoulli_loss, + _bernoulli_gh!, + yview, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) for r in 1:p glm.B[r, j] = β_buf[r] @@ -979,48 +1016,46 @@ For input x ∈ ℝᵖ the conditional distribution of y ∈ ℤ₊ᵏ factorize - `B`: coefficient matrix of size `p × k` - `prior`: regularization prior applied per column (default `NoPrior()`) """ -mutable struct MvPoissonGLM{T<:Real, P<:AbstractPrior} <: AbstractGLM +mutable struct MvPoissonGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM B::Matrix{T} prior::P in_dim::Int out_dim::Int end -function MvPoissonGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real, P<:AbstractPrior} +function MvPoissonGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) - return MvPoissonGLM{T, P}(Matrix{T}(B), prior, p, k) + return MvPoissonGLM{T,P}(Matrix{T}(B), prior, p, k) end MvPoissonGLM(B::AbstractMatrix{T}) where {T<:Real} = MvPoissonGLM(B, NoPrior()) MvPoissonGLM(B::AbstractMatrix) = MvPoissonGLM(float.(B), NoPrior()) -MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior) = - MvPoissonGLM(float.(B), prior) +MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior) = MvPoissonGLM(float.(B), prior) DensityInterface.DensityKind(::MvPoissonGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( - glm::MvPoissonGLM, - y::AbstractVector; - control_seq::AbstractVector{<:Real}, + glm::MvPoissonGLM, y::AbstractVector; control_seq::AbstractVector{<:Real} ) length(y) == glm.out_dim || throw(DimensionMismatch("y length $(length(y)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) Tη = float(promote_type(eltype(glm.B), eltype(control_seq))) lp = zero(Tη) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) yj = y[j] yj >= 0 || return Tη(-Inf) η = zero(Tη) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end lp += yj * η - exp(η) - logfactorial(yj) @@ -1029,9 +1064,7 @@ function DensityInterface.logdensityof( end function Random.rand( - rng::AbstractRNG, - glm::MvPoissonGLM; - control_seq::AbstractVector{<:Real}, + rng::AbstractRNG, glm::MvPoissonGLM; control_seq::AbstractVector{<:Real} ) out = Vector{Int}(undef, glm.out_dim) rand!(rng, glm, out; control_seq=control_seq) @@ -1051,15 +1084,16 @@ function Random.rand!( ) length(out) == glm.out_dim || throw(DimensionMismatch("out length $(length(out)) ≠ out_dim $(glm.out_dim)")) - length(control_seq) == glm.in_dim || - throw(DimensionMismatch( - "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)", - )) + length(control_seq) == glm.in_dim || throw( + DimensionMismatch( + "control_seq length $(length(control_seq)) ≠ in_dim $(glm.in_dim)" + ), + ) T = eltype(glm.B) - for j in 1:glm.out_dim + for j in 1:(glm.out_dim) η = zero(T) - for r in 1:glm.in_dim + for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end out[j] = rand(rng, Poisson(exp(η))) @@ -1087,31 +1121,41 @@ function StatsAPI.fit!( n, p = size(control_seq) length(obs_seq) == n || throw(DimensionMismatch("obs_seq length $(length(obs_seq)) ≠ control_seq rows $n")) - length(weight_seq) == n || - throw(DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n")) + length(weight_seq) == n || throw( + DimensionMismatch("weight_seq length $(length(weight_seq)) ≠ control_seq rows $n"), + ) p == glm.in_dim || throw(DimensionMismatch("control_seq columns $p ≠ in_dim $(glm.in_dim)")) k = glm.out_dim for i in 1:n - length(obs_seq[i]) == k || - throw(DimensionMismatch( - "obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k", - )) + length(obs_seq[i]) == k || throw( + DimensionMismatch("obs_seq[$i] length $(length(obs_seq[i])) ≠ out_dim $k") + ) end β_buf = Vector{T}(undef, p) - g = Vector{T}(undef, p) - H = Matrix{T}(undef, p, p) - Δ = Vector{T}(undef, p) + g = Vector{T}(undef, p) + H = Matrix{T}(undef, p, p) + Δ = Vector{T}(undef, p) for j in 1:k yview = _ColumnElementView(obs_seq, j) copyto!(β_buf, view(glm.B, :, j)) _newton_solve!( - β_buf, g, H, Δ, _poisson_loss, _poisson_gh!, - yview, weight_seq, control_seq, glm.prior; - max_iter=max_iter, gtol=gtol, max_backtrack=max_backtrack, + β_buf, + g, + H, + Δ, + _poisson_loss, + _poisson_gh!, + yview, + weight_seq, + control_seq, + glm.prior; + max_iter=max_iter, + gtol=gtol, + max_backtrack=max_backtrack, ) for r in 1:p glm.B[r, j] = β_buf[r] diff --git a/src/multivariate/t.jl b/src/multivariate/t.jl index 7a4f069..e181c69 100644 --- a/src/multivariate/t.jl +++ b/src/multivariate/t.jl @@ -147,11 +147,12 @@ function DensityInterface.logdensityof(dist::MultivariateT, x::AbstractVector) d = dist.dim ν = dist.ν - diff = dist._diff - #= Zero allocation: reuse the struct scratch buffer for the residual. - Sequential use only — not safe for concurrent calls on the same dist. =# - @. diff = x - dist.μ + #= Allocate the residual locally (one length-d vector). Thread-safe: + parallel forward/backward calls on the same dist don't share state. + The struct scratch (`_diff`, `_z`) is reserved for `fit!`, which + runs sequentially per state inside HMM EM. =# + diff = x - dist.μ ldiv!(dist.Σ_chol.L, diff) mahal² = zero(eltype(diff)) for i in eachindex(diff) @@ -313,7 +314,7 @@ function StatsAPI.fit!( end # Symmetrize numerical noise - for j in 1:d, k in 1:j-1 + for j in 1:d, k in 1:(j - 1) s = (Σ_acc[j, k] + Σ_acc[k, j]) / 2 Σ_acc[j, k] = s Σ_acc[k, j] = s diff --git a/test/Project.toml b/test/Project.toml index b89abb9..42352bc 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,14 +3,10 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" 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" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsAPI = "82ae8749-77ed-4fe6-ae5f-f523153014b0" -Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[compat] -JET = "0.11.2" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" \ No newline at end of file diff --git a/test/allocations.jl b/test/allocations.jl index 6db8fd4..2636de1 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -6,62 +6,123 @@ using DensityInterface using StatsAPI #= - Allocation regression tests. The pattern - is: warm up by running the operation once, then call it many times in a - type-stable function and divide @allocated by the rep count to get - per-call bytes. + Allocation regression tests. Warm up first, then call the operation many + times in a type-stable function and divide @allocated by the rep count. Steady-state targets (per call): - - All logdensityof: 0 bytes - - All rand!: 0 bytes (rand allocates only the return vector) - - fit! is bounded by O(p² + k²), independent of n - - See `bench_*` helpers below for why a top-level `@allocated foo()` would - report misleading kwarg-lowering noise that does NOT exist in real loops. + - Univariate logdensityof, MvBernoulli/MvPoisson logdensityof, + MultivariateTDiag logdensityof, PoissonZeroInflated logdensityof: 0 B + - MvGaussianGLM logdensityof, MultivariateT logdensityof: 1 length-k + vector per call. Thread-safe by design — the struct scratch was removed + because two threads reading the same dist would race. Bound to 256 B/call. + - All Mv rand!: 0 bytes (zero-alloc, thread-safe via in-place lmul!). + - fit! is bounded by O(p² + k²), independent of n. + + Top-level `@allocated foo()` reports kwarg-lowering noise (~48 B) that + does NOT exist in real loops — measure inside a function. =# -bench_logd(d, y, x, n) = (s = 0.0; for _ in 1:n; s += logdensityof(d, y; control_seq=x); end; s) -bench_logd_unctrl(d, y, n) = (s = 0.0; for _ in 1:n; s += logdensityof(d, y); end; s) -bench_rand_scalar(rng, d, x, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d; control_seq=x); end; s) -bench_rand_int(rng, d, x, n) = (s = 0; for _ in 1:n; s += rand(rng, d; control_seq=x); end; s) -bench_rand!_v(rng, d, out, x, n) = (s = 0.0; for _ in 1:n; rand!(rng, d, out; control_seq=x); s += out[1]; end; s) -bench_rand!_i(rng, d, out, x, n) = (s = 0; for _ in 1:n; rand!(rng, d, out; control_seq=x); s += out[1]; end; s) -bench_rand_unctrl_scalar(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d); end; s) -bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; end; s) +bench_logd(d, y, x, n) = ( + s=0.0; + for _ in 1:n + s += logdensityof(d, y; control_seq=x) + end; + s +) +bench_logd_unctrl(d, y, n) = ( + s=0.0; + for _ in 1:n + s += logdensityof(d, y) + end; + s +) +bench_rand_scalar(rng, d, x, n) = ( + s=0.0; + for _ in 1:n + s += rand(rng, d; control_seq=x) + end; + s +) +bench_rand_int(rng, d, x, n) = ( + s=0; + for _ in 1:n + s += rand(rng, d; control_seq=x) + end; + s +) +bench_rand!_v(rng, d, out, x, n) = ( + s=0.0; + for _ in 1:n + rand!(rng, d, out; control_seq=x) + s += out[1] + end; + s +) +bench_rand!_i(rng, d, out, x, n) = ( + s=0; + for _ in 1:n + rand!(rng, d, out; control_seq=x) + s += out[1] + end; + s +) +bench_rand_unctrl_scalar(rng, d, n) = ( + s=0.0; + for _ in 1:n + s += rand(rng, d) + end; + s +) +bench_rand_unctrl_vec(rng, d, n) = ( + s=0.0; + for _ in 1:n + s += rand(rng, d)[1] + end; + s +) @testset "Allocations (steady state)" begin rng = Random.MersenneTwister(0) REPS = 1000 - @testset "logdensityof — zero alloc per call" begin + @testset "logdensityof — bounded per call" begin x = [1.0, 2.0] - g = GaussianGLM([0.5, -1.0], 1.0) - b = BernoulliGLM([0.5, -1.0]) - p = PoissonGLM([0.5, -1.0]) + g = GaussianGLM([0.5, -1.0], 1.0) + b = BernoulliGLM([0.5, -1.0]) + p = PoissonGLM([0.5, -1.0]) mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) - yv = [0.1, 0.2]; yi = [0, 1] + yv = [0.1, 0.2] + yi = [0, 1] # Warm - bench_logd(g, 0.5, x, 1); bench_logd(b, 1, x, 1); bench_logd(p, 2, x, 1) - bench_logd(mg, yv, x, 1); bench_logd(mb, yi, x, 1); bench_logd(mp, yi, x, 1) - - @test (@allocated bench_logd(g, 0.5, x, REPS)) == 0 - @test (@allocated bench_logd(b, 1, x, REPS)) == 0 - @test (@allocated bench_logd(p, 2, x, REPS)) == 0 - @test (@allocated bench_logd(mg, yv, x, REPS)) == 0 - @test (@allocated bench_logd(mb, yi, x, REPS)) == 0 - @test (@allocated bench_logd(mp, yi, x, REPS)) == 0 + bench_logd(g, 0.5, x, 1) + bench_logd(b, 1, x, 1) + bench_logd(p, 2, x, 1) + bench_logd(mg, yv, x, 1) + bench_logd(mb, yi, x, 1) + bench_logd(mp, yi, x, 1) + + # Truly zero-alloc (no scratch needed) + @test (@allocated bench_logd(g, 0.5, x, REPS)) == 0 + @test (@allocated bench_logd(b, 1, x, REPS)) == 0 + @test (@allocated bench_logd(p, 2, x, REPS)) == 0 + @test (@allocated bench_logd(mb, yi, x, REPS)) == 0 + @test (@allocated bench_logd(mp, yi, x, REPS)) == 0 + + # MvGaussianGLM: one length-k residual per call (thread-safe). + # Bound to ~256 B/call → 256 * REPS for the loop. + @test (@allocated bench_logd(mg, yv, x, REPS)) ≤ 256 * REPS end @testset "rand!/rand — zero alloc beyond return" begin x = [1.0, 2.0] - g = GaussianGLM([0.5, -1.0], 1.0) - b = BernoulliGLM([0.5, -1.0]) - p = PoissonGLM([0.5, -1.0]) + g = GaussianGLM([0.5, -1.0], 1.0) + b = BernoulliGLM([0.5, -1.0]) + p = PoissonGLM([0.5, -1.0]) mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) @@ -92,7 +153,8 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; # GaussianGLM: closed-form WLS. Workspace: XWX (p²) + XWy (p). yg = randn(rng, n) - gg = GaussianGLM([0.0, 0.0], 1.0); fit!(gg, yg, w; control_seq=X) + gg = GaussianGLM([0.0, 0.0], 1.0) + fit!(gg, yg, w; control_seq=X) gg = GaussianGLM([0.0, 0.0], 1.0) @test (@allocated fit!(gg, yg, w; control_seq=X)) ≤ 1_000 @@ -107,24 +169,28 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; # BernoulliGLM/PoissonGLM: hand-rolled Newton. # Workspace: g, H, Δ — three small alloc, total ~300 bytes for p=2. yb = Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:n] - gb = BernoulliGLM(zeros(2)); fit!(gb, yb, w; control_seq=X) + gb = BernoulliGLM(zeros(2)) + fit!(gb, yb, w; control_seq=X) gb = BernoulliGLM(zeros(2)) @test (@allocated fit!(gb, yb, w; control_seq=X)) ≤ 1_000 yp = Int[rand(rng, 0:5) for _ in 1:n] - gp = PoissonGLM(zeros(2)); fit!(gp, yp, w; control_seq=X) + gp = PoissonGLM(zeros(2)) + fit!(gp, yp, w; control_seq=X) gp = PoissonGLM(zeros(2)) @test (@allocated fit!(gp, yp, w; control_seq=X)) ≤ 1_000 # Multivariate Newton: workspace shared across columns. _ColumnElementView # avoids the n-sized per-column copy that Vector-of-Vectors would force. ymb = [Int[rand(rng) < 0.5 ? 1 : 0 for _ in 1:2] for _ in 1:n] - gmb = MvBernoulliGLM(zeros(2, 2)); fit!(gmb, ymb, w; control_seq=X) + gmb = MvBernoulliGLM(zeros(2, 2)) + fit!(gmb, ymb, w; control_seq=X) gmb = MvBernoulliGLM(zeros(2, 2)) @test (@allocated fit!(gmb, ymb, w; control_seq=X)) ≤ 1_000 ymp = [Int[rand(rng, 0:3) for _ in 1:2] for _ in 1:n] - gmp = MvPoissonGLM(zeros(2, 2)); fit!(gmp, ymp, w; control_seq=X) + gmp = MvPoissonGLM(zeros(2, 2)) + fit!(gmp, ymp, w; control_seq=X) gmp = MvPoissonGLM(zeros(2, 2)) @test (@allocated fit!(gmp, ymp, w; control_seq=X)) ≤ 1_000 end @@ -132,7 +198,8 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; @testset "PoissonZeroInflated" begin zip = PoissonZeroInflated(3.0, 0.2) - bench_logd_unctrl(zip, 0, 1); bench_logd_unctrl(zip, 5, 1) + bench_logd_unctrl(zip, 0, 1) + bench_logd_unctrl(zip, 5, 1) @test (@allocated bench_logd_unctrl(zip, 0, REPS)) == 0 @test (@allocated bench_logd_unctrl(zip, 5, REPS)) == 0 @@ -142,7 +209,8 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; n = 500 y = [rand(rng, zip) for _ in 1:n] w = ones(n) - zip2 = PoissonZeroInflated(1.0, 0.1); fit!(zip2, y, w) + zip2 = PoissonZeroInflated(1.0, 0.1) + fit!(zip2, y, w) zip2 = PoissonZeroInflated(1.0, 0.1) @test (@allocated fit!(zip2, y, w)) ≤ 1_000 end @@ -152,9 +220,10 @@ bench_rand_unctrl_vec(rng, d, n) = (s = 0.0; for _ in 1:n; s += rand(rng, d)[1]; mvt = MultivariateT([0.0, 0.0], [1.0 0.3; 0.3 1.0], 5.0) xv = [0.1, 0.2] - # Cholesky now stored as :L and residual reuses struct scratch. + # Cholesky stored as :L (no .L copy). Residual is allocated locally + # per call so concurrent calls on the same dist are race-free. bench_logd_unctrl(mvt, xv, 1) - @test (@allocated bench_logd_unctrl(mvt, xv, REPS)) == 0 + @test (@allocated bench_logd_unctrl(mvt, xv, REPS)) ≤ 256 * REPS bench_rand_unctrl_vec(rng, mvt, 1) # rand still allocates the return vector; bound it loosely. diff --git a/test/glm/gaussian.jl b/test/glm/gaussian.jl index 9fd6499..0564bf0 100644 --- a/test/glm/gaussian.jl +++ b/test/glm/gaussian.jl @@ -9,145 +9,159 @@ using EmissionModels rng = Random.MersenneTwister(1234) -function _synthetic_gaussian_glm(rng, n::Int, p::Int; β::Vector{Float64}, σ2::Float64, weights=:uniform) +function _synthetic_gaussian_glm( + rng, n::Int, p::Int; β::Vector{Float64}, σ2::Float64, weights=:uniform +) @assert length(β) == p - X = hcat(ones(n), randn(rng, n, p-1)) # intercept + random features + X = hcat(ones(n), randn(rng, n, p - 1)) # intercept + random features μ = X * β y = rand.(Ref(rng), Normal.(μ, sqrt(σ2))) - w = weights === :uniform ? ones(n) : - weights === :random ? (rand(rng, n) .+ 0.5) : + w = if weights === :uniform + ones(n) + elseif weights === :random + (rand(rng, n) .+ 0.5) + else error("weights must be :uniform or :random") + end return X, y, w end _expected_logpdf(y, μ, σ2) = -0.5 * log(2π * σ2) - 0.5 * ((y - μ)^2 / σ2) @testset "GaussianGLM Basic Functionality" begin - @testset "Constructor" begin - glm = GaussianGLM([0.5, -1.0], 1.0) - @test glm.β == [0.5, -1.0] - @test glm.σ2 == 1.0 - - glm2 = GaussianGLM([1, 2], 3) - @test glm2.β == [1, 2] - @test glm2.σ2 == 3 - end + @testset "Constructor" begin + glm = GaussianGLM([0.5, -1.0], 1.0) + @test glm.β == [0.5, -1.0] + @test glm.σ2 == 1.0 - @testset "DensityInterface" begin - glm = GaussianGLM([0.5, -1.0], 1.0) + glm2 = GaussianGLM([1, 2], 3) + @test glm2.β == [1, 2] + @test glm2.σ2 == 3 + end - # Trait - @test DensityKind(glm) == HasDensity() + @testset "DensityInterface" begin + glm = GaussianGLM([0.5, -1.0], 1.0) - X = [1.0 2.0; 1.0 3.0; 1.0 4.0] - μ = X * glm.β - y = rand.(Ref(rng), Normal.(μ, sqrt(glm.σ2))) + # Trait + @test DensityKind(glm) == HasDensity() - # logdensityof matches closed form - for i in eachindex(y) - x_i = vec(X[i, :]) # ensure Vector - logp = logdensityof(glm, y[i]; control_seq=x_i) - @test isfinite(logp) + X = [1.0 2.0; 1.0 3.0; 1.0 4.0] + μ = X * glm.β + y = rand.(Ref(rng), Normal.(μ, sqrt(glm.σ2))) - expected = _expected_logpdf(y[i], dot(glm.β, x_i), glm.σ2) - @test logp ≈ expected rtol=1e-12 atol=0.0 - end + # logdensityof matches closed form + for i in eachindex(y) + x_i = vec(X[i, :]) # ensure Vector + logp = logdensityof(glm, y[i]; control_seq=x_i) + @test isfinite(logp) - # Higher density near the mean than far away - x0 = vec(X[1, :]) - y_at_mean = dot(glm.β, x0) - y_far = y_at_mean + 10.0 - @test logdensityof(glm, y_at_mean; control_seq=x0) > logdensityof(glm, y_far; control_seq=x0) + expected = _expected_logpdf(y[i], dot(glm.β, x_i), glm.σ2) + @test logp ≈ expected rtol = 1e-12 atol = 0.0 end - @testset "Random sampling" begin - glm = GaussianGLM([0.5, -1.0], 2.0) # variance 2.0 - x = [1.0, 3.0] - n = 10_000 - samples = [rand(rng, glm; control_seq=x) for _ in 1:n] + # Higher density near the mean than far away + x0 = vec(X[1, :]) + y_at_mean = dot(glm.β, x0) + y_far = y_at_mean + 10.0 + @test logdensityof(glm, y_at_mean; control_seq=x0) > + logdensityof(glm, y_far; control_seq=x0) + end - @test all(s -> s isa Real, samples) + @testset "Random sampling" begin + glm = GaussianGLM([0.5, -1.0], 2.0) # variance 2.0 + x = [1.0, 3.0] + n = 10_000 + samples = [rand(rng, glm; control_seq=x) for _ in 1:n] - # Empirical mean should be close to dot(β,x) - μ = dot(glm.β, x) - m = mean(samples) - @test m ≈ μ atol=0.15 + @test all(s -> s isa Real, samples) - # Empirical variance should be close to σ2 - v = var(samples) - @test v ≈ glm.σ2 atol=0.25 - end + # Empirical mean should be close to dot(β,x) + μ = dot(glm.β, x) + m = mean(samples) + @test m ≈ μ atol = 0.15 - @testset "fit! with uniform weights" begin - β_true = [0.3, -1.2] # p=2 - σ2_true = 0.7 - X, y, w = _synthetic_gaussian_glm(rng, 2000, 2; β=β_true, σ2=σ2_true, weights=:uniform) + # Empirical variance should be close to σ2 + v = var(samples) + @test v ≈ glm.σ2 atol = 0.25 + end - glm = GaussianGLM([0.0, 0.0], 1.0) - fit!(glm, y, w; control_seq=X) + @testset "fit! with uniform weights" begin + β_true = [0.3, -1.2] # p=2 + σ2_true = 0.7 + X, y, w = _synthetic_gaussian_glm( + rng, 2000, 2; β=β_true, σ2=σ2_true, weights=:uniform + ) - # Coefficients should be close with enough data - @test glm.β ≈ β_true atol=0.08 - @test glm.σ2 ≈ σ2_true atol=0.10 + glm = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm, y, w; control_seq=X) - @test all(isfinite, glm.β) - @test isfinite(glm.σ2) - @test glm.σ2 > 0 - end + # Coefficients should be close with enough data + @test glm.β ≈ β_true atol = 0.08 + @test glm.σ2 ≈ σ2_true atol = 0.10 - @testset "fit! with weighted observations" begin - β_true = [1.0, 0.6] - σ2_true = 1.5 - X, y, w = _synthetic_gaussian_glm(rng, 2500, 2; β=β_true, σ2=σ2_true, weights=:random) + @test all(isfinite, glm.β) + @test isfinite(glm.σ2) + @test glm.σ2 > 0 + end - glm = GaussianGLM([0.0, 0.0], 1.0) - fit!(glm, y, w; control_seq=X) + @testset "fit! with weighted observations" begin + β_true = [1.0, 0.6] + σ2_true = 1.5 + X, y, w = _synthetic_gaussian_glm( + rng, 2500, 2; β=β_true, σ2=σ2_true, weights=:random + ) - @test glm.β ≈ β_true atol=0.10 - @test glm.σ2 ≈ σ2_true atol=0.15 - end + glm = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm, y, w; control_seq=X) - @testset "Constructor with prior" begin - glm = GaussianGLM([0.0, 0.0], 1.0) - @test glm.prior isa NoPrior + @test glm.β ≈ β_true atol = 0.10 + @test glm.σ2 ≈ σ2_true atol = 0.15 + end - glm2 = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(2.0)) - @test glm2.prior isa RidgePrior - @test glm2.prior.λ == 2.0 - end + @testset "Constructor with prior" begin + glm = GaussianGLM([0.0, 0.0], 1.0) + @test glm.prior isa NoPrior - @testset "fit! with RidgePrior shrinks toward zero" begin - β_true = [3.0, -3.0] - σ2_true = 1.0 - X, y, w = _synthetic_gaussian_glm(rng, 300, 2; β=β_true, σ2=σ2_true) + glm2 = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(2.0)) + @test glm2.prior isa RidgePrior + @test glm2.prior.λ == 2.0 + end - glm_noprior = GaussianGLM([0.0, 0.0], 1.0) - fit!(glm_noprior, y, w; control_seq=X) + @testset "fit! with RidgePrior shrinks toward zero" begin + β_true = [3.0, -3.0] + σ2_true = 1.0 + X, y, w = _synthetic_gaussian_glm(rng, 300, 2; β=β_true, σ2=σ2_true) - glm_ridge = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(10.0)) - fit!(glm_ridge, y, w; control_seq=X) + glm_noprior = GaussianGLM([0.0, 0.0], 1.0) + fit!(glm_noprior, y, w; control_seq=X) - @test norm(glm_ridge.β) < norm(glm_noprior.β) - @test all(isfinite, glm_ridge.β) - end + glm_ridge = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(10.0)) + fit!(glm_ridge, y, w; control_seq=X) + @test norm(glm_ridge.β) < norm(glm_noprior.β) + @test all(isfinite, glm_ridge.β) + end end -function _synthetic_mvgaussian_glm(rng, n::Int, p::Int, k::Int; - B::Matrix{Float64}, Σ::Matrix{Float64}, - weights=:uniform) +function _synthetic_mvgaussian_glm( + rng, n::Int, p::Int, k::Int; B::Matrix{Float64}, Σ::Matrix{Float64}, weights=:uniform +) @assert size(B) == (p, k) @assert size(Σ) == (k, k) - X = hcat(ones(n), randn(rng, n, p-1)) + X = hcat(ones(n), randn(rng, n, p - 1)) L = cholesky(Σ).L obs_seq = Vector{Vector{Float64}}(undef, n) for i in 1:n μ_i = vec(B' * X[i, :]) obs_seq[i] = μ_i .+ L * randn(rng, k) end - w = weights === :uniform ? ones(n) : - weights === :random ? (rand(rng, n) .+ 0.5) : + w = if weights === :uniform + ones(n) + elseif weights === :random + (rand(rng, n) .+ 0.5) + else error("weights must be :uniform or :random") + end return X, obs_seq, w end @@ -171,7 +185,9 @@ end glm_ridge = MvGaussianGLM(B, Σ, RidgePrior(2.0)) @test glm_ridge.prior isa RidgePrior - @test_throws DimensionMismatch MvGaussianGLM(B, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]) + @test_throws DimensionMismatch MvGaussianGLM( + B, [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0] + ) @test_throws ArgumentError MvGaussianGLM(B, [1.0 0.0; 0.0 -1.0]) end @@ -189,11 +205,11 @@ end # Compare against MvNormal closed form diff = y - μ expected = -log(2π) - 0.5 * logdet(Σ) - 0.5 * dot(diff, Σ \ diff) - @test logdensityof(glm, y; control_seq=x) ≈ expected rtol=1e-10 + @test logdensityof(glm, y; control_seq=x) ≈ expected rtol = 1e-10 # Higher density at the mean than far from it @test logdensityof(glm, μ; control_seq=x) > - logdensityof(glm, μ .+ 10.0; control_seq=x) + logdensityof(glm, μ .+ 10.0; control_seq=x) @test_throws DimensionMismatch logdensityof(glm, [1.0]; control_seq=x) @test_throws DimensionMismatch logdensityof(glm, y; control_seq=[1.0]) @@ -211,23 +227,22 @@ end μ_true = vec(B' * x) m = vec(mean(Y; dims=1)) - @test m ≈ μ_true atol=0.05 + @test m ≈ μ_true atol = 0.05 S = (Y .- m')' * (Y .- m') / (n - 1) - @test S ≈ Σ atol=0.1 + @test S ≈ Σ atol = 0.1 end @testset "fit! recovers B and Σ" begin B_true = [0.5 -1.0; 1.0 0.5; -0.3 0.8] # p=3, k=2 Σ_true = [1.0 0.3; 0.3 1.5] - X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 4000, 3, 2; - B=B_true, Σ=Σ_true) + X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 4000, 3, 2; B=B_true, Σ=Σ_true) glm = MvGaussianGLM(zeros(3, 2), Matrix(1.0I, 2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.08 - @test glm.Σ ≈ Σ_true atol=0.10 + @test glm.B ≈ B_true atol = 0.08 + @test glm.Σ ≈ Σ_true atol = 0.10 @test all(isfinite, glm.B) @test all(isfinite, glm.Σ) @test isposdef(glm.Σ) @@ -236,22 +251,21 @@ end @testset "fit! random weights" begin B_true = [1.0 0.0; -0.5 1.5] Σ_true = [0.5 -0.1; -0.1 0.7] - X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 3000, 2, 2; - B=B_true, Σ=Σ_true, - weights=:random) + X, obs_seq, w = _synthetic_mvgaussian_glm( + rng, 3000, 2, 2; B=B_true, Σ=Σ_true, weights=:random + ) glm = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.10 - @test glm.Σ ≈ Σ_true atol=0.12 + @test glm.B ≈ B_true atol = 0.10 + @test glm.Σ ≈ Σ_true atol = 0.12 end @testset "fit! with RidgePrior shrinks toward zero" begin B_true = [3.0 -3.0; 3.0 3.0] Σ_true = [1.0 0.0; 0.0 1.0] - X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 200, 2, 2; - B=B_true, Σ=Σ_true) + X, obs_seq, w = _synthetic_mvgaussian_glm(rng, 200, 2, 2; B=B_true, Σ=Σ_true) glm_noprior = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) fit!(glm_noprior, obs_seq, w; control_seq=X) @@ -275,4 +289,4 @@ end wrong_dim_obs = [zeros(3) for _ in 1:5] @test_throws DimensionMismatch fit!(glm, wrong_dim_obs, ones(5); control_seq=X) end -end \ No newline at end of file +end diff --git a/test/glm/test_bernoulli_poisson.jl b/test/glm/test_bernoulli_poisson.jl index b30429f..e3722e5 100644 --- a/test/glm/test_bernoulli_poisson.jl +++ b/test/glm/test_bernoulli_poisson.jl @@ -6,7 +6,6 @@ using StatsAPI using LinearAlgebra using Distributions: Bernoulli, Poisson, logpdf - function _sigmoid(η) return one(η) / (one(η) + exp(-η)) end @@ -92,8 +91,8 @@ end η = dot(glm.β, x) μ = _sigmoid(η) - @test logdensityof(glm, 1; control_seq=x) ≈ log(μ) rtol=1e-10 - @test logdensityof(glm, 0; control_seq=x) ≈ log(1 - μ) rtol=1e-10 + @test logdensityof(glm, 1; control_seq=x) ≈ log(μ) rtol = 1e-10 + @test logdensityof(glm, 0; control_seq=x) ≈ log(1 - μ) rtol = 1e-10 @test logdensityof(glm, 2; control_seq=x) == -Inf @test logdensityof(glm, -1; control_seq=x) == -Inf end @@ -116,7 +115,7 @@ end samples = [rand(rng, glm; control_seq=x) for _ in 1:n] @test all(s -> s == 0 || s == 1, samples) - @test mean(samples) ≈ _sigmoid(dot(glm.β, x)) atol=0.05 + @test mean(samples) ≈ _sigmoid(dot(glm.β, x)) atol = 0.05 end @testset "fit! uniform weights" begin @@ -126,7 +125,7 @@ end glm = BernoulliGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.15 + @test glm.β ≈ β_true atol = 0.15 @test all(isfinite, glm.β) end @@ -137,7 +136,7 @@ end glm = BernoulliGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.20 + @test glm.β ≈ β_true atol = 0.20 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -160,7 +159,9 @@ end X = ones(5, 2) @test_throws DimensionMismatch fit!(glm, ones(Int, 4), ones(5); control_seq=X) @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(4); control_seq=X) - @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(5); control_seq=ones(5, 3)) + @test_throws DimensionMismatch fit!( + glm, ones(Int, 5), ones(5); control_seq=ones(5, 3) + ) end @testset "fit! all-zero observations" begin @@ -177,7 +178,6 @@ end end end - @testset "PoissonGLM" begin rng = Random.MersenneTwister(99) @@ -199,7 +199,7 @@ end μ = exp(η) for k in 0:5 - @test logdensityof(glm, k; control_seq=x) ≈ logpdf(Poisson(μ), k) rtol=1e-10 + @test logdensityof(glm, k; control_seq=x) ≈ logpdf(Poisson(μ), k) rtol = 1e-10 end @test logdensityof(glm, -1; control_seq=x) == -Inf @@ -212,7 +212,7 @@ end samples = [rand(rng, glm; control_seq=x) for _ in 1:n] @test all(s -> s >= 0, samples) - @test mean(samples) ≈ exp(glm.β[1]) atol=0.2 + @test mean(samples) ≈ exp(glm.β[1]) atol = 0.2 end @testset "fit! uniform weights" begin @@ -222,7 +222,7 @@ end glm = PoissonGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.15 + @test glm.β ≈ β_true atol = 0.15 @test all(isfinite, glm.β) end @@ -233,7 +233,7 @@ end glm = PoissonGLM(zeros(2)) fit!(glm, y, w; control_seq=X) - @test glm.β ≈ β_true atol=0.15 + @test glm.β ≈ β_true atol = 0.15 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -255,7 +255,9 @@ end X = ones(5, 2) @test_throws DimensionMismatch fit!(glm, ones(Int, 4), ones(5); control_seq=X) @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(4); control_seq=X) - @test_throws DimensionMismatch fit!(glm, ones(Int, 5), ones(5); control_seq=ones(5, 3)) + @test_throws DimensionMismatch fit!( + glm, ones(Int, 5), ones(5); control_seq=ones(5, 3) + ) end @testset "fit! all-zero counts" begin @@ -277,8 +279,9 @@ function _synthetic_mvbernoulli(rng, n, B_true; weights=:uniform) X = hcat(ones(n), randn(rng, n, p - 1)) obs_seq = Vector{Vector{Int}}(undef, n) for i in 1:n - obs_seq[i] = Int[rand(rng) < _sigmoid(dot(B_true[:, j], X[i, :])) ? 1 : 0 - for j in 1:k] + obs_seq[i] = Int[ + rand(rng) < _sigmoid(dot(B_true[:, j], X[i, :])) ? 1 : 0 for j in 1:k + ] end w = weights === :uniform ? ones(n) : rand(rng, n) .+ 0.5 return X, obs_seq, w @@ -323,10 +326,11 @@ end η2 = dot(B[:, 2], x) μ1, μ2 = _sigmoid(η1), _sigmoid(η2) - @test logdensityof(glm, [1, 1]; control_seq=x) ≈ log(μ1) + log(μ2) rtol=1e-10 - @test logdensityof(glm, [1, 0]; control_seq=x) ≈ log(μ1) + log(1 - μ2) rtol=1e-10 - @test logdensityof(glm, [0, 1]; control_seq=x) ≈ log(1 - μ1) + log(μ2) rtol=1e-10 - @test logdensityof(glm, [0, 0]; control_seq=x) ≈ log(1 - μ1) + log(1 - μ2) rtol=1e-10 + @test logdensityof(glm, [1, 1]; control_seq=x) ≈ log(μ1) + log(μ2) rtol = 1e-10 + @test logdensityof(glm, [1, 0]; control_seq=x) ≈ log(μ1) + log(1 - μ2) rtol = 1e-10 + @test logdensityof(glm, [0, 1]; control_seq=x) ≈ log(1 - μ1) + log(μ2) rtol = 1e-10 + @test logdensityof(glm, [0, 0]; control_seq=x) ≈ log(1 - μ1) + log(1 - μ2) rtol = + 1e-10 @test logdensityof(glm, [2, 0]; control_seq=x) == -Inf end @@ -347,8 +351,8 @@ end Y = reduce(hcat, samples)' # n × k @test all(s -> s == 0 || s == 1, vec(Y)) - @test mean(Y[:, 1]) ≈ _sigmoid(dot(B[:, 1], x)) atol=0.05 - @test mean(Y[:, 2]) ≈ _sigmoid(dot(B[:, 2], x)) atol=0.05 + @test mean(Y[:, 1]) ≈ _sigmoid(dot(B[:, 1], x)) atol = 0.05 + @test mean(Y[:, 2]) ≈ _sigmoid(dot(B[:, 2], x)) atol = 0.05 end @testset "fit! recovers B" begin @@ -358,7 +362,7 @@ end glm = MvBernoulliGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.20 + @test glm.B ≈ B_true atol = 0.20 @test all(isfinite, glm.B) end @@ -369,7 +373,7 @@ end glm = MvBernoulliGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.25 + @test glm.B ≈ B_true atol = 0.25 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -397,7 +401,7 @@ end y_j = [obs_seq[i][j] for i in eachindex(obs_seq)] glm_j = BernoulliGLM(zeros(2)) fit!(glm_j, y_j, w; control_seq=X) - @test glm.B[:, j] ≈ glm_j.β rtol=1e-6 + @test glm.B[:, j] ≈ glm_j.β rtol = 1e-6 end end @@ -406,7 +410,7 @@ end X = ones(5, 2) good_obs = [zeros(Int, 2) for _ in 1:5] @test_throws DimensionMismatch fit!( - glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X, + glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X ) @test_throws DimensionMismatch fit!(glm, good_obs, ones(4); control_seq=X) @test_throws DimensionMismatch fit!(glm, good_obs, ones(5); control_seq=ones(5, 3)) @@ -445,7 +449,7 @@ end for k1 in 0:3, k2 in 0:3 expected = logpdf(Poisson(μ1), k1) + logpdf(Poisson(μ2), k2) - @test logdensityof(glm, [k1, k2]; control_seq=x) ≈ expected rtol=1e-10 + @test logdensityof(glm, [k1, k2]; control_seq=x) ≈ expected rtol = 1e-10 end @test logdensityof(glm, [-1, 0]; control_seq=x) == -Inf @@ -462,8 +466,8 @@ end Y = reduce(hcat, samples)' @test all(s -> s >= 0, vec(Y)) - @test mean(Y[:, 1]) ≈ exp(B[1, 1]) atol=0.2 - @test mean(Y[:, 2]) ≈ exp(B[1, 2]) atol=0.2 + @test mean(Y[:, 1]) ≈ exp(B[1, 1]) atol = 0.2 + @test mean(Y[:, 2]) ≈ exp(B[1, 2]) atol = 0.2 end @testset "fit! recovers B" begin @@ -473,7 +477,7 @@ end glm = MvPoissonGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.15 + @test glm.B ≈ B_true atol = 0.15 @test all(isfinite, glm.B) end @@ -484,7 +488,7 @@ end glm = MvPoissonGLM(zeros(2, 2)) fit!(glm, obs_seq, w; control_seq=X) - @test glm.B ≈ B_true atol=0.20 + @test glm.B ≈ B_true atol = 0.20 end @testset "fit! with RidgePrior shrinks toward zero" begin @@ -512,7 +516,7 @@ end y_j = [obs_seq[i][j] for i in eachindex(obs_seq)] glm_j = PoissonGLM(zeros(2)) fit!(glm_j, y_j, w; control_seq=X) - @test glm.B[:, j] ≈ glm_j.β rtol=1e-6 + @test glm.B[:, j] ≈ glm_j.β rtol = 1e-6 end end @@ -521,7 +525,7 @@ end X = ones(5, 2) good_obs = [zeros(Int, 2) for _ in 1:5] @test_throws DimensionMismatch fit!( - glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X, + glm, [zeros(Int, 2) for _ in 1:4], ones(5); control_seq=X ) @test_throws DimensionMismatch fit!(glm, good_obs, ones(4); control_seq=X) @test_throws DimensionMismatch fit!(glm, good_obs, ones(5); control_seq=ones(5, 3)) diff --git a/test/runtests.jl b/test/runtests.jl index a66ea46..eec4621 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,16 +1,15 @@ using EmissionModels using Test using Aqua -using JET -using JuliaFormatter using StatsAPI @testset "EmissionModels.jl" begin @testset "Code linting" begin - if VERSION >= v"1.10" && isempty(VERSION.prerelease) + if isempty(VERSION.prerelease) + using Pkg + Pkg.add("JET") + using JET JET.test_package(EmissionModels; target_modules=(EmissionModels,)) - else - @info "Skipping JET on Julia $VERSION (requires >=1.10 and a release build)" end end From 941714f2556ad7d8e7d88bc4cbf914f0a16e185a Mon Sep 17 00:00:00 2001 From: Ryan Senne <50930199+rsenne@users.noreply.github.com> Date: Thu, 14 May 2026 12:31:30 -0400 Subject: [PATCH 20/20] Add tests for codecov. Also remove numerical hackyenss --- src/glms/glm.jl | 243 +++++++++++++++++++-------- src/multivariate/t.jl | 27 +-- test/glm/test_promotion_and_types.jl | 197 ++++++++++++++++++++++ test/runtests.jl | 1 + 4 files changed, 382 insertions(+), 86 deletions(-) create mode 100644 test/glm/test_promotion_and_types.jl diff --git a/src/glms/glm.jl b/src/glms/glm.jl index cef2df0..3a1d237 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -86,30 +86,41 @@ mutable struct GaussianGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} σ2::T prior::P + function GaussianGLM{T,P}( + β::Vector{T}, σ2::T, prior::P + ) where {T<:Real,P<:AbstractPrior} + return new{T,P}(β, σ2, prior) + end end -function GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:Real} - return GaussianGLM{T,NoPrior}(Vector{T}(β), σ2, NoPrior()) -end - -function GaussianGLM(β::AbstractVector{T}, σ2::T, prior::P) where {T<:Real,P<:AbstractPrior} +function GaussianGLM( + β::AbstractVector{T}, σ2::T, prior::P +) where {T<:AbstractFloat,P<:AbstractPrior} return GaussianGLM{T,P}(Vector{T}(β), σ2, prior) end - -# Convenience: promote β eltype and σ2 together -GaussianGLM(β::AbstractVector, σ2::Real) = GaussianGLM(float.(β), float(σ2)) +function GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:AbstractFloat} + return GaussianGLM(β, σ2, NoPrior()) +end function GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior) - return GaussianGLM(float.(β), float(σ2), prior) + T = float(promote_type(eltype(β), typeof(σ2))) + return GaussianGLM(convert(Vector{T}, β), convert(T, σ2), prior) end +GaussianGLM(β::AbstractVector, σ2::Real) = GaussianGLM(β, σ2, NoPrior()) DensityInterface.DensityKind(::GaussianGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( reg::GaussianGLM, y::Real; control_seq::AbstractVector{<:Real} ) - μ = dot(reg.β, control_seq) - return -0.5 * log(2π * reg.σ2) - 0.5 * ((y - μ)^2 / reg.σ2) + T = float(promote_type(eltype(reg.β), typeof(reg.σ2), eltype(control_seq), typeof(y))) + μ = zero(T) + for i in eachindex(reg.β, control_seq) + μ += reg.β[i] * control_seq[i] + end + σ2 = T(reg.σ2) + diff = T(y) - μ + return -log(T(2π) * σ2) / 2 - diff * diff / (2 * σ2) end function Random.rand( @@ -154,7 +165,14 @@ function StatsAPI.fit!( # RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. neglogprior_hess!(reg.prior, XWX, reg.β) - F = cholesky!(Symmetric(XWX)) + F = cholesky!(Symmetric(XWX); check=false) + issuccess(F) || throw( + ArgumentError( + "XᵀWX is not positive definite — `control_seq` is rank-deficient " * + "or weights are degenerate. Add a RidgePrior(λ) to regularize, or " * + "drop collinear columns from `control_seq`.", + ), + ) copyto!(reg.β, XWy) ldiv!(F, reg.β) @@ -237,6 +255,12 @@ function _bernoulli_gh!( return nothing end +#= Bound η so that exp(η) stays well below floatmax(T), with a few nats of + headroom for the w·exp(η) products that accumulate downstream. Type-aware: + ≈707 for Float64, ≈86 for Float32 — replaces the previous arbitrary ±500 + constant, which was both Float64-only and a magic number. =# +@inline _η_bound(::Type{T}) where {T<:AbstractFloat} = log(floatmax(T)) - T(2) + function _poisson_loss( β::AbstractVector{T}, y::AbstractVector, @@ -246,9 +270,10 @@ function _poisson_loss( ) where {T<:Real} n, _ = size(X) nll = zero(T) + η_max = _η_bound(T) for i in 1:n x_i = view(X, i, :) - η_i = clamp(dot(β, x_i), T(-500), T(500)) + η_i = clamp(dot(β, x_i), -η_max, η_max) nll += T(w[i]) * (exp(η_i) - T(y[i]) * η_i) end return nll + neglogprior(prior, β) @@ -266,10 +291,11 @@ function _poisson_gh!( n, p = size(X) fill!(g, zero(T)) fill!(H, zero(T)) + η_max = _η_bound(T) for i in 1:n x_i = view(X, i, :) wi = T(w[i]) - η_i = clamp(dot(β, x_i), T(-500), T(500)) + η_i = clamp(dot(β, x_i), -η_max, η_max) eη = exp(η_i) r_i = wi * (eη - T(y[i])) W_i = wi * eη @@ -450,23 +476,41 @@ to the solver. mutable struct BernoulliGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} prior::P -end -function BernoulliGLM(β::AbstractVector{T}) where {T<:Real} - return BernoulliGLM{T,NoPrior}(Vector{T}(β), NoPrior()) + function BernoulliGLM{T,P}(β::Vector{T}, prior::P) where {T<:Real,P<:AbstractPrior} + return new{T,P}(β, prior) + end end -function BernoulliGLM(β::AbstractVector{T}, prior::P) where {T<:Real,P<:AbstractPrior} +function BernoulliGLM( + β::AbstractVector{T}, prior::P +) where {T<:AbstractFloat,P<:AbstractPrior} return BernoulliGLM{T,P}(Vector{T}(β), prior) end +function BernoulliGLM(β::AbstractVector{T}) where {T<:AbstractFloat} + return BernoulliGLM(β, NoPrior()) +end + +#= Promoting fallback for integer / mixed-eltype β (e.g., Vector{Int}). The + typed constructor above is restricted to AbstractFloat because the Newton + solver allocates `Vector{T}` workspace and cannot use Int. =# +function BernoulliGLM(β::AbstractVector, prior::AbstractPrior) + T = float(eltype(β)) + return BernoulliGLM(convert(Vector{T}, β), prior) +end +BernoulliGLM(β::AbstractVector) = BernoulliGLM(β, NoPrior()) DensityInterface.DensityKind(::BernoulliGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( glm::BernoulliGLM, y::Integer; control_seq::AbstractVector{<:Real} ) - (y == 0 || y == 1) || return oftype(dot(glm.β, control_seq), -Inf) - η = dot(glm.β, control_seq) + T = float(promote_type(eltype(glm.β), eltype(control_seq))) + (y == 0 || y == 1) || return T(-Inf) + η = zero(T) + for i in eachindex(glm.β, control_seq) + η += glm.β[i] * control_seq[i] + end return y == 1 ? -log1pexp(-η) : -log1pexp(η) end @@ -523,24 +567,39 @@ composes without changes to the solver. mutable struct PoissonGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM β::Vector{T} prior::P -end -function PoissonGLM(β::AbstractVector{T}) where {T<:Real} - return PoissonGLM{T,NoPrior}(Vector{T}(β), NoPrior()) + function PoissonGLM{T,P}(β::Vector{T}, prior::P) where {T<:Real,P<:AbstractPrior} + return new{T,P}(β, prior) + end end -function PoissonGLM(β::AbstractVector{T}, prior::P) where {T<:Real,P<:AbstractPrior} +function PoissonGLM( + β::AbstractVector{T}, prior::P +) where {T<:AbstractFloat,P<:AbstractPrior} return PoissonGLM{T,P}(Vector{T}(β), prior) end +function PoissonGLM(β::AbstractVector{T}) where {T<:AbstractFloat} + return PoissonGLM(β, NoPrior()) +end + +function PoissonGLM(β::AbstractVector, prior::AbstractPrior) + T = float(eltype(β)) + return PoissonGLM(convert(Vector{T}, β), prior) +end +PoissonGLM(β::AbstractVector) = PoissonGLM(β, NoPrior()) DensityInterface.DensityKind(::PoissonGLM) = DensityInterface.HasDensity() function DensityInterface.logdensityof( glm::PoissonGLM, y::Integer; control_seq::AbstractVector{<:Real} ) - y >= 0 || return oftype(dot(glm.β, control_seq), -Inf) - η = dot(glm.β, control_seq) - return y * η - exp(η) - logfactorial(y) + T = float(promote_type(eltype(glm.β), eltype(control_seq))) + y >= 0 || return T(-Inf) + η = zero(T) + for i in eachindex(glm.β, control_seq) + η += glm.β[i] * control_seq[i] + end + return T(y) * η - exp(η) - T(logfactorial(y)) end function Random.rand(rng::AbstractRNG, glm::PoissonGLM; control_seq::AbstractVector{<:Real}) @@ -605,11 +664,23 @@ mutable struct MvGaussianGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM logdetΣ::T in_dim::Int out_dim::Int + + function MvGaussianGLM{T,P}( + B::Matrix{T}, + Σ::Matrix{T}, + prior::P, + Σ_chol::Cholesky{T,Matrix{T}}, + logdetΣ::T, + in_dim::Int, + out_dim::Int, + ) where {T<:Real,P<:AbstractPrior} + return new{T,P}(B, Σ, prior, Σ_chol, logdetΣ, in_dim, out_dim) + end end function MvGaussianGLM( B::AbstractMatrix{T}, Σ::AbstractMatrix{T}, prior::P -) where {T<:Real,P<:AbstractPrior} +) where {T<:AbstractFloat,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) @@ -617,32 +688,26 @@ function MvGaussianGLM( DimensionMismatch("Σ must be $(k)×$(k) for B with $(k) columns, got $(size(Σ))") ) - Σ_chol = try - cholesky(Symmetric(Σ, :L)) - catch - throw(ArgumentError("Σ must be positive definite")) - end + Σ_chol = cholesky(Symmetric(Σ, :L); check=false) + issuccess(Σ_chol) || throw(ArgumentError("Σ must be positive definite")) return MvGaussianGLM{T,P}( Matrix{T}(B), Matrix{T}(Σ), prior, Σ_chol, logdet(Σ_chol), p, k ) end -function MvGaussianGLM(B::AbstractMatrix{T}, Σ::AbstractMatrix{T}) where {T<:Real} +function MvGaussianGLM(B::AbstractMatrix{T}, Σ::AbstractMatrix{T}) where {T<:AbstractFloat} return MvGaussianGLM(B, Σ, NoPrior()) end -function MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix) - T = promote_type(eltype(B), eltype(Σ)) - Tf = float(T) - return MvGaussianGLM(Matrix{Tf}(B), Matrix{Tf}(Σ), NoPrior()) -end - +#= Promoting fallback: B and Σ are promoted to a common float eltype. Handles + mixed (Float32/Float64) and integer eltypes so user code can write + `MvGaussianGLM([1 0; 0 1], [1.0 0; 0 1.0])` without first converting. =# function MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix, prior::AbstractPrior) - T = promote_type(eltype(B), eltype(Σ)) - Tf = float(T) - return MvGaussianGLM(Matrix{Tf}(B), Matrix{Tf}(Σ), prior) + T = float(promote_type(eltype(B), eltype(Σ))) + return MvGaussianGLM(convert(Matrix{T}, B), convert(Matrix{T}, Σ), prior) end +MvGaussianGLM(B::AbstractMatrix, Σ::AbstractMatrix) = MvGaussianGLM(B, Σ, NoPrior()) DensityInterface.DensityKind(::MvGaussianGLM) = DensityInterface.HasDensity() @@ -669,11 +734,13 @@ function DensityInterface.logdensityof( diff = Vector{T}(undef, k) #= μ = Bᵀ x and diff = y − μ, fused as a single loop to avoid the - Adjoint*Vector temporary that BLAS would otherwise allocate. =# + Adjoint*Vector temporary that BLAS would otherwise allocate. The output + type is `eltype(B)` — inputs in higher precision are downcast on entry + and the residual buffer stays at the struct's float type. =# for j in 1:k sj = zero(T) for r in 1:p - sj += glm.B[r, j] * control_seq[r] + sj += glm.B[r, j] * T(control_seq[r]) end diff[j] = T(y[j]) - sj end @@ -683,7 +750,7 @@ function DensityInterface.logdensityof( mahal² += diff[j] * diff[j] end - return -k / 2 * log(2π) - glm.logdetΣ / 2 - mahal² / 2 + return -T(k) * log(T(2π)) / 2 - glm.logdetΣ / 2 - mahal² / 2 end function Random.rand( @@ -787,13 +854,17 @@ function StatsAPI.fit!( Pass a length-p view so RidgePrior loops over rows of XWX, not p*k. =# neglogprior_hess!(glm.prior, XWX, view(glm.B, :, 1)) - F = try - cholesky(Symmetric(XWX, :L)) - catch - min_eig = minimum(eigvals(Hermitian(XWX))) - XWX .+= (abs(min_eig) + T(1e-6)) * I - cholesky(Symmetric(XWX, :L)) - end + #= XᵀWX is PD whenever the (weighted) design matrix has full column rank. + Failure here means rank-deficient `control_seq` — surface that as a + clear error rather than silently shifting eigenvalues. =# + F = cholesky!(Symmetric(XWX, :L); check=false) + issuccess(F) || throw( + ArgumentError( + "XᵀWX is not positive definite — `control_seq` is rank-deficient " * + "or weights are degenerate. Add a RidgePrior(λ) to regularize, or " * + "drop collinear columns from `control_seq`.", + ), + ) copyto!(glm.B, XWY) ldiv!(F, glm.B) @@ -826,13 +897,17 @@ function StatsAPI.fit!( end Σ_new ./= wsum - Σ_chol_new = try - cholesky(Symmetric(Σ_new, :L)) - catch - min_eig = minimum(eigvals(Hermitian(Σ_new))) - Σ_new .+= (abs(min_eig) + T(1e-6)) * I - cholesky(Symmetric(Σ_new, :L)) - end + #= Σ_new is Σwᵢ·rᵢrᵢᵀ / Σwᵢ, which is PSD by construction and PD whenever + the residuals span ℝᵏ. Failure means a degenerate output dimension + (zero variance) — error rather than silently regularize. =# + Σ_chol_new = cholesky(Symmetric(Σ_new, :L); check=false) + issuccess(Σ_chol_new) || throw( + ArgumentError( + "Residual covariance Σ is not positive definite — at least one " * + "output dimension has zero (or perfectly collinear) residual " * + "variance. Check `obs_seq` for degenerate output columns.", + ), + ) copyto!(glm.Σ, Σ_new) glm.Σ_chol = Σ_chol_new @@ -859,20 +934,31 @@ mutable struct MvBernoulliGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM prior::P in_dim::Int out_dim::Int + + function MvBernoulliGLM{T,P}( + B::Matrix{T}, prior::P, in_dim::Int, out_dim::Int + ) where {T<:Real,P<:AbstractPrior} + return new{T,P}(B, prior, in_dim, out_dim) + end end -function MvBernoulliGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real,P<:AbstractPrior} +function MvBernoulliGLM( + B::AbstractMatrix{T}, prior::P +) where {T<:AbstractFloat,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) return MvBernoulliGLM{T,P}(Matrix{T}(B), prior, p, k) end +function MvBernoulliGLM(B::AbstractMatrix{T}) where {T<:AbstractFloat} + return MvBernoulliGLM(B, NoPrior()) +end -MvBernoulliGLM(B::AbstractMatrix{T}) where {T<:Real} = MvBernoulliGLM(B, NoPrior()) - -MvBernoulliGLM(B::AbstractMatrix) = MvBernoulliGLM(float.(B), NoPrior()) - -MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior) = MvBernoulliGLM(float.(B), prior) +function MvBernoulliGLM(B::AbstractMatrix, prior::AbstractPrior) + T = float(eltype(B)) + return MvBernoulliGLM(convert(Matrix{T}, B), prior) +end +MvBernoulliGLM(B::AbstractMatrix) = MvBernoulliGLM(B, NoPrior()) DensityInterface.DensityKind(::MvBernoulliGLM) = DensityInterface.HasDensity() @@ -1021,20 +1107,31 @@ mutable struct MvPoissonGLM{T<:Real,P<:AbstractPrior} <: AbstractGLM prior::P in_dim::Int out_dim::Int + + function MvPoissonGLM{T,P}( + B::Matrix{T}, prior::P, in_dim::Int, out_dim::Int + ) where {T<:Real,P<:AbstractPrior} + return new{T,P}(B, prior, in_dim, out_dim) + end end -function MvPoissonGLM(B::AbstractMatrix{T}, prior::P) where {T<:Real,P<:AbstractPrior} +function MvPoissonGLM( + B::AbstractMatrix{T}, prior::P +) where {T<:AbstractFloat,P<:AbstractPrior} p, k = size(B) p > 0 || throw(ArgumentError("B must have at least one row")) k > 0 || throw(ArgumentError("B must have at least one column")) return MvPoissonGLM{T,P}(Matrix{T}(B), prior, p, k) end +function MvPoissonGLM(B::AbstractMatrix{T}) where {T<:AbstractFloat} + return MvPoissonGLM(B, NoPrior()) +end -MvPoissonGLM(B::AbstractMatrix{T}) where {T<:Real} = MvPoissonGLM(B, NoPrior()) - -MvPoissonGLM(B::AbstractMatrix) = MvPoissonGLM(float.(B), NoPrior()) - -MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior) = MvPoissonGLM(float.(B), prior) +function MvPoissonGLM(B::AbstractMatrix, prior::AbstractPrior) + T = float(eltype(B)) + return MvPoissonGLM(convert(Matrix{T}, B), prior) +end +MvPoissonGLM(B::AbstractMatrix) = MvPoissonGLM(B, NoPrior()) DensityInterface.DensityKind(::MvPoissonGLM) = DensityInterface.HasDensity() @@ -1058,7 +1155,7 @@ function DensityInterface.logdensityof( for r in 1:(glm.in_dim) η += glm.B[r, j] * control_seq[r] end - lp += yj * η - exp(η) - logfactorial(yj) + lp += Tη(yj) * η - exp(η) - Tη(logfactorial(yj)) end return lp end diff --git a/src/multivariate/t.jl b/src/multivariate/t.jl index e181c69..81fd3d5 100644 --- a/src/multivariate/t.jl +++ b/src/multivariate/t.jl @@ -45,11 +45,8 @@ mutable struct MultivariateT{T<:Real} size(Σ) == (dim, dim) || throw(DimensionMismatch("Σ must be $(dim)×$(dim), got $(size(Σ))")) - Σ_chol = try - cholesky(Symmetric(Σ, :L)) - catch - throw(ArgumentError("Σ must be positive definite")) - end + Σ_chol = cholesky(Symmetric(Σ, :L); check=false) + issuccess(Σ_chol) || throw(ArgumentError("Σ must be positive definite")) logdetΣ = logdet(Σ_chol) @@ -320,14 +317,18 @@ function StatsAPI.fit!( Σ_acc[k, j] = s end - # Update Σ and Cholesky, regularizing if needed - Σ_chol_new = try - cholesky(Symmetric(Σ_acc, :L)) - catch - min_eig = minimum(eigvals(Hermitian(Σ_acc))) - Σ_acc .+= (abs(min_eig) + 1e-6) * I - cholesky(Symmetric(Σ_acc, :L)) - end + #= Σ_acc is PSD by construction and PD whenever the weighted residuals + span ℝᵈ. Failure here means a degenerate observation set (zero + variance along some axis) — surface that rather than silently + regularizing with an arbitrary eigenvalue shift. =# + Σ_chol_new = cholesky(Symmetric(Σ_acc, :L); check=false) + issuccess(Σ_chol_new) || throw( + ArgumentError( + "Σ M-step is not positive definite — observations are " * + "degenerate along at least one dimension. Inspect `obs_seq` " * + "for collinear or constant components.", + ), + ) dist.Σ .= Σ_acc dist.Σ_chol = Σ_chol_new dist.logdetΣ = logdet(dist.Σ_chol) diff --git a/test/glm/test_promotion_and_types.jl b/test/glm/test_promotion_and_types.jl new file mode 100644 index 0000000..0a1ff5f --- /dev/null +++ b/test/glm/test_promotion_and_types.jl @@ -0,0 +1,197 @@ +using EmissionModels +using Test +using LinearAlgebra +using DensityInterface +using StatsAPI + +@testset "Constructor promotion — integer and mixed eltypes" begin + @testset "GaussianGLM" begin + # Integer β/σ2 → Float64 + g = GaussianGLM([1, 2], 3) + @test g isa GaussianGLM{Float64,NoPrior} + @test eltype(g.β) === Float64 + @test g.σ2 === 3.0 + + g_ridge = GaussianGLM([1, 2], 3, RidgePrior(0.5)) + @test g_ridge isa GaussianGLM{Float64,RidgePrior{Float64}} + + # Float32 preserved + g32 = GaussianGLM(Float32[0.5, -1.0], 1.0f0) + @test g32 isa GaussianGLM{Float32,NoPrior} + @test eltype(g32.β) === Float32 + @test g32.σ2 === 1.0f0 + + # Mixed Float32 β + Float64 σ2 promotes to Float64 + g_mix = GaussianGLM(Float32[0.5, -1.0], 1.0) + @test g_mix isa GaussianGLM{Float64,NoPrior} + end + + @testset "BernoulliGLM" begin + b = BernoulliGLM([1, 2]) + @test b isa BernoulliGLM{Float64,NoPrior} + @test eltype(b.β) === Float64 + + b_ridge = BernoulliGLM([1, 2], RidgePrior(0.5)) + @test b_ridge isa BernoulliGLM{Float64,RidgePrior{Float64}} + + b32 = BernoulliGLM(Float32[0.5, -1.0]) + @test b32 isa BernoulliGLM{Float32,NoPrior} + + b32_ridge = BernoulliGLM(Float32[0.5, -1.0], RidgePrior(0.5f0)) + @test b32_ridge isa BernoulliGLM{Float32,RidgePrior{Float32}} + end + + @testset "PoissonGLM" begin + p = PoissonGLM([1, 2]) + @test p isa PoissonGLM{Float64,NoPrior} + + p_ridge = PoissonGLM([1, 2], RidgePrior(0.5)) + @test p_ridge isa PoissonGLM{Float64,RidgePrior{Float64}} + + p32 = PoissonGLM(Float32[0.5, -1.0]) + @test p32 isa PoissonGLM{Float32,NoPrior} + end + + @testset "MvGaussianGLM" begin + # All-integer B and Σ → Float64 + g = MvGaussianGLM([1 0; 0 1], [1 0; 0 1]) + @test g isa MvGaussianGLM{Float64,NoPrior} + @test eltype(g.B) === Float64 + @test eltype(g.Σ) === Float64 + + # Mixed: Int B, Float64 Σ + g_mix = MvGaussianGLM([1 0; 0 1], [1.0 0.0; 0.0 1.0]) + @test g_mix isa MvGaussianGLM{Float64,NoPrior} + + # Float32 preserved across both + g32 = MvGaussianGLM(Float32[1.0 0.0; 0.0 1.0], Float32[1.0 0.0; 0.0 1.0]) + @test g32 isa MvGaussianGLM{Float32,NoPrior} + @test eltype(g32.B) === Float32 + + # Mixed Float32 / Float64 → Float64 + g_mix32 = MvGaussianGLM(Float32[1.0 0.0; 0.0 1.0], [1.0 0.0; 0.0 1.0]) + @test g_mix32 isa MvGaussianGLM{Float64,NoPrior} + + # Integer + prior + g_ridge = MvGaussianGLM([1 0; 0 1], [1 0; 0 1], RidgePrior(0.5)) + @test g_ridge isa MvGaussianGLM{Float64,RidgePrior{Float64}} + end + + @testset "MvBernoulliGLM" begin + b = MvBernoulliGLM([1 0; 0 1]) + @test b isa MvBernoulliGLM{Float64,NoPrior} + + b_ridge = MvBernoulliGLM([1 0; 0 1], RidgePrior(0.5)) + @test b_ridge isa MvBernoulliGLM{Float64,RidgePrior{Float64}} + + b32 = MvBernoulliGLM(Float32[1.0 0.0; 0.0 1.0]) + @test b32 isa MvBernoulliGLM{Float32,NoPrior} + + b32_ridge = MvBernoulliGLM(Float32[1.0 0.0; 0.0 1.0], RidgePrior(0.5f0)) + @test b32_ridge isa MvBernoulliGLM{Float32,RidgePrior{Float32}} + end + + @testset "MvPoissonGLM" begin + p = MvPoissonGLM([1 0; 0 1]) + @test p isa MvPoissonGLM{Float64,NoPrior} + + p_ridge = MvPoissonGLM([1 0; 0 1], RidgePrior(0.5)) + @test p_ridge isa MvPoissonGLM{Float64,RidgePrior{Float64}} + + p32 = MvPoissonGLM(Float32[0.5 -1.0; 1.0 0.5]) + @test p32 isa MvPoissonGLM{Float32,NoPrior} + end +end + +@testset "logdensityof return-type stability" begin + @testset "Float32 inputs return Float32" begin + glm_g = GaussianGLM(Float32[0.5, -1.0], 1.0f0) + glm_b = BernoulliGLM(Float32[0.5, -1.0]) + glm_p = PoissonGLM(Float32[0.5, -1.0]) + glm_mg = MvGaussianGLM(Float32[0.5 -1.0; 1.0 0.5], Float32[1.0 0.3; 0.3 1.5]) + glm_mb = MvBernoulliGLM(Float32[0.5 -1.0; 1.0 0.5]) + glm_mp = MvPoissonGLM(Float32[0.5 -1.0; 0.2 0.0]) + + x32 = Float32[1.0, 2.0] + y32v = Float32[0.1, 0.2] + yi = [0, 1] + + @test @inferred(logdensityof(glm_g, 0.5f0; control_seq=x32)) isa Float32 + @test @inferred(logdensityof(glm_b, 1; control_seq=x32)) isa Float32 + @test @inferred(logdensityof(glm_b, 0; control_seq=x32)) isa Float32 + @test @inferred(logdensityof(glm_b, 2; control_seq=x32)) isa Float32 # -Inf branch + @test @inferred(logdensityof(glm_p, 3; control_seq=x32)) isa Float32 + @test @inferred(logdensityof(glm_p, -1; control_seq=x32)) isa Float32 # -Inf branch + @test @inferred(logdensityof(glm_mg, y32v; control_seq=x32)) isa Float32 + @test @inferred(logdensityof(glm_mb, yi; control_seq=x32)) isa Float32 + @test @inferred(logdensityof(glm_mp, yi; control_seq=x32)) isa Float32 + end + + @testset "Float64 inputs return Float64" begin + glm_g = GaussianGLM([0.5, -1.0], 1.0) + glm_b = BernoulliGLM([0.5, -1.0]) + glm_p = PoissonGLM([0.5, -1.0]) + glm_mg = MvGaussianGLM([0.5 -1.0; 1.0 0.5], [1.0 0.3; 0.3 1.5]) + glm_mb = MvBernoulliGLM([0.5 -1.0; 1.0 0.5]) + glm_mp = MvPoissonGLM([0.5 -1.0; 0.2 0.0]) + + x = [1.0, 2.0] + yv = [0.1, 0.2] + yi = [0, 1] + + @test @inferred(logdensityof(glm_g, 0.5; control_seq=x)) isa Float64 + @test @inferred(logdensityof(glm_b, 1; control_seq=x)) isa Float64 + @test @inferred(logdensityof(glm_p, 3; control_seq=x)) isa Float64 + @test @inferred(logdensityof(glm_mg, yv; control_seq=x)) isa Float64 + @test @inferred(logdensityof(glm_mb, yi; control_seq=x)) isa Float64 + @test @inferred(logdensityof(glm_mp, yi; control_seq=x)) isa Float64 + end +end + +@testset "fit! errors on rank-deficient design" begin + # X column 2 is constant 0 — XᵀWX is rank-deficient. + X_sing = ones(5, 2) + X_sing[:, 2] .= 0.0 + w = ones(5) + + @testset "GaussianGLM" begin + glm = GaussianGLM([0.0, 0.0], 1.0) + y = zeros(5) + @test_throws ArgumentError fit!(glm, y, w; control_seq=X_sing) + + # Same data fits fine with a RidgePrior — verifies the error message's suggestion. + glm_ridge = GaussianGLM([0.0, 0.0], 1.0, RidgePrior(1.0)) + @test fit!(glm_ridge, y, w; control_seq=X_sing) === glm_ridge + @test all(isfinite, glm_ridge.β) + end + + @testset "MvGaussianGLM" begin + # Varied obs so the Σ M-step is well-conditioned — isolate the failure + # to the (rank-deficient) XᵀWX step. + obs = [[1.0, 2.0], [3.0, -1.0], [-1.0, 0.5], [2.0, 3.0], [0.5, -2.0]] + glm = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2)) + @test_throws ArgumentError fit!(glm, obs, w; control_seq=X_sing) + + glm_ridge = MvGaussianGLM(zeros(2, 2), Matrix(1.0I, 2, 2), RidgePrior(1.0)) + @test fit!(glm_ridge, obs, w; control_seq=X_sing) === glm_ridge + @test all(isfinite, glm_ridge.B) + end +end + +@testset "PoissonGLM type-aware η bound (no Float32 overflow)" begin + #= With the old hardcoded clamp(η, ±500), exp(500) overflows Float32 to Inf. + The new bound is log(floatmax(T)) − 2, type-aware, so exp(η_max) stays + within the float range for both Float32 and Float64. =# + glm32 = PoissonGLM(Float32[1.0, 0.5]) + X = Float32[1.0 1.0; 1.0 -1.0; 1.0 0.5] + y = [1, 2, 0] + w = Float32[1.0, 1.0, 1.0] + fit!(glm32, y, w; control_seq=X) + @test all(isfinite, glm32.β) + @test eltype(glm32.β) === Float32 + + # Same for Float64 + glm64 = PoissonGLM([1.0, 0.5]) + fit!(glm64, y, w; control_seq=Float64.(X)) + @test all(isfinite, glm64.β) +end diff --git a/test/runtests.jl b/test/runtests.jl index eec4621..278fce6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -16,6 +16,7 @@ using StatsAPI @testset "GLM Models" begin include("glm/gaussian.jl") include("glm/test_bernoulli_poisson.jl") + include("glm/test_promotion_and_types.jl") end @testset "Zero-inflated models" begin include("zeroinflated/test_poisson.jl")