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/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/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 6b78852..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) @@ -12,7 +15,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..b6d1782 --- /dev/null +++ b/docs/src/distributions.md @@ -0,0 +1,100 @@ +# 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. + +## 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 new file mode 100644 index 0000000..1f74600 --- /dev/null +++ b/docs/src/glm.md @@ -0,0 +1,101 @@ +# 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)) +``` + +## 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): + +``\ell(β) = -\sum_{i=1}^{n} w_i \, \log p(y_i \mid x_i, β) + \text{neglogprior}(β)`` + +```julia +# 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/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..4b2316e --- /dev/null +++ b/docs/src/priors.md @@ -0,0 +1,81 @@ +# 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. + +## API Reference + +```@docs +AbstractPrior +NoPrior +RidgePrior +``` diff --git a/src/EmissionModels.jl b/src/EmissionModels.jl index 0491f6b..e059dcc 100644 --- a/src/EmissionModels.jl +++ b/src/EmissionModels.jl @@ -1,9 +1,9 @@ module EmissionModels -using Distributions: Poisson, Chisq +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 @@ -13,10 +13,15 @@ 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, BernoulliGLM, PoissonGLM +export MvGaussianGLM, MvBernoulliGLM, MvPoissonGLM +export AbstractPrior, NoPrior, RidgePrior +export neglogprior, neglogprior_grad!, neglogprior_hess! end diff --git a/src/glms/glm.jl b/src/glms/glm.jl index 0221454..3a1d237 100644 --- a/src/glms/glm.jl +++ b/src/glms/glm.jl @@ -4,9 +4,1259 @@ 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 + +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 + +""" + 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. 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 +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 + 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, prior::P +) where {T<:AbstractFloat,P<:AbstractPrior} + return GaussianGLM{T,P}(Vector{T}(β), σ2, prior) +end +function GaussianGLM(β::AbstractVector{T}, σ2::T) where {T<:AbstractFloat} + return GaussianGLM(β, σ2, NoPrior()) +end + +function GaussianGLM(β::AbstractVector, σ2::Real, prior::AbstractPrior) + 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} +) + 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( + 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{T}, + obs_seq::AbstractVector{<:Real}, + weights::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(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 + + # RidgePrior(λ) accumulates λI into XᵀWX, giving (XᵀWX + λI)β = XᵀWy. + neglogprior_hess!(reg.prior, XWX, reg.β) + + 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.β) + + 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 +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] + +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 + +#= 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, + w::AbstractVector{<:Real}, + X::AbstractMatrix{<:Real}, + prior::AbstractPrior, +) 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), -η_max, η_max) + 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)) + η_max = _η_bound(T) + for i in 1:n + x_i = view(X, i, :) + wi = T(w[i]) + η_i = clamp(dot(β, x_i), -η_max, η_max) + 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 a hand-rolled Newton solver with +backtracking line search, 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 + + 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<: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} +) + 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 + +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, max_iter=50, gtol=1e-8, max_backtrack=20) + +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, + obs_seq::AbstractVector, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, + 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 + +""" + 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 +a hand-rolled Newton solver with backtracking line search. 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 + + 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<: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} +) + 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}) + return rand(rng, Poisson(exp(dot(glm.β, control_seq)))) +end + +""" + fit!(glm::PoissonGLM, obs_seq, weight_seq; + control_seq, max_iter=50, gtol=1e-8, max_backtrack=20) + +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, + obs_seq::AbstractVector, + weight_seq::AbstractVector{<:Real}; + control_seq::AbstractMatrix{<:Real}, + 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()`) +""" +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 + + 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<: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")) + size(Σ) == (k, k) || throw( + DimensionMismatch("Σ must be $(k)×$(k) for B with $(k) columns, got $(size(Σ))") + ) + + Σ_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<:AbstractFloat} + return MvGaussianGLM(B, Σ, 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 = 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() + +""" + logdensityof(glm::MvGaussianGLM, y::AbstractVector; control_seq) + +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} +) + 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 = 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. 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] * T(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 -T(k) * log(T(2π)) / 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 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, + 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 + + 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 + end + 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"), + ) + 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 + end + + #= 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)) + + #= 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) + + #= Σ 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 + 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 + end + Σ_new ./= wsum + + #= Σ_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 + 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 + + 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<: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 + +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() + +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 + +""" + 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 + + 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<: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 + +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() + +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 += Tη(yj) * η - exp(η) - Tη(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 2eff54d..81fd3d5 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,20 +45,15 @@ 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 - throw(ArgumentError("Σ must be positive definite")) - end + Σ_chol = cholesky(Symmetric(Σ, :L); check=false) + issuccess(Σ_chol) || throw(ArgumentError("Σ must be positive definite")) 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 +91,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 +111,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 +131,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 +145,29 @@ 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-μ) + #= 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) + mahal² += diff[i] * diff[i] + end - # 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 +176,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 +195,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 +206,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 +252,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 +259,89 @@ 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.Σ) + #= Σ_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) - # 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 +350,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 +361,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 +376,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 +448,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 +455,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 +531,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 +542,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 +557,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) 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 new file mode 100644 index 0000000..2636de1 --- /dev/null +++ b/test/allocations.jl @@ -0,0 +1,261 @@ +using EmissionModels +using Test +using Random +using LinearAlgebra +using DensityInterface +using StatsAPI + +#= + 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): + - 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 +) + +@testset "Allocations (steady state)" begin + rng = Random.MersenneTwister(0) + REPS = 1000 + + @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]) + 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) + + # 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]) + 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 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)) ≤ 256 * REPS + + 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 new file mode 100644 index 0000000..0564bf0 --- /dev/null +++ b/test/glm/gaussian.jl @@ -0,0 +1,292 @@ + +using Test +using Random +using LinearAlgebra +using Distributions +using DensityInterface +using StatsAPI +using EmissionModels + +rng = Random.MersenneTwister(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 = 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 "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 + + @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 + +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 = 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 + +@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 diff --git a/test/glm/test_bernoulli_poisson.jl b/test/glm/test_bernoulli_poisson.jl new file mode 100644 index 0000000..e3722e5 --- /dev/null +++ b/test/glm/test_bernoulli_poisson.jl @@ -0,0 +1,536 @@ +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 + +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/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 1c6a94a..278fce6 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,18 +1,23 @@ 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 + + @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") end @@ -21,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