Add duration distribution helpers

docs/src/api.md

@@ -112,6 +112,8 @@ HiddenMarkovModels.log_initialization
 HiddenMarkovModels.log_transition_matrix
 HiddenMarkovModels.mul_rows_cols!
 HiddenMarkovModels.argmaxplus_transmul!
+HiddenMarkovModels.duration_logdensityof
+HiddenMarkovModels.rand_duration
 ```
 
 ## Index

src/HiddenMarkovModels.jl

@@ -38,6 +38,7 @@ include("utils/fit.jl")
 include("utils/lightdiagnormal.jl")
 include("utils/lightcategorical.jl")
 include("utils/limits.jl")
+include("utils/duration.jl")
 
 include("inference/predict.jl")
 include("inference/forward.jl")

src/utils/duration.jl

@@ -0,0 +1,30 @@
+"""
+$(SIGNATURES)
+
+Log-probability that a state with sojourn-time distribution `dist` is held for exactly `k`
+consecutive timesteps.
+
+`dist` is supplied by the user with support on `{0, 1, 2, ...}` and is interpreted as the law
+of `(sojourn time - 1)`. The sojourn time itself lives on `{1, 2, 3, ...}`, so this returns
+`logdensityof(dist, k - 1)`. Durations `k < 1` fall outside the user's support and yield
+`-Inf` through `dist` (e.g. a `Distributions.Distribution` returns `-Inf` for a negative
+argument).
+
+See also [`rand_duration`](@ref).
+"""
+duration_logdensityof(dist, k::Integer) = logdensityof(dist, k - one(k))
+
+"""
+$(SIGNATURES)
+
+Sample a sojourn time on `{1, 2, 3, ...}` from the duration distribution `dist`.
+
+`dist` is interpreted as the law of `(sojourn time - 1)` (see [`duration_logdensityof`](@ref)),
+so this returns `rand(rng, dist) + 1`.
+"""
+function rand_duration(rng::AbstractRNG, dist)
+    t = rand(rng, dist)
+    return t + one(t)
+end
+
+rand_duration(dist) = rand_duration(default_rng(), dist)

test/duration.jl

@@ -0,0 +1,69 @@
+using HiddenMarkovModels
+using HiddenMarkovModels: duration_logdensityof, rand_duration
+using DensityInterface: DensityInterface, DensityKind, HasDensity, logdensityof
+using Distributions: Geometric, NegativeBinomial, Poisson
+using Random: AbstractRNG
+using StableRNGs: StableRNG
+using Statistics: mean
+using Test
+
+# A minimal user-defined duration distribution that only implements the required interface,
+# to confirm we never reach into `Distributions`-specific internals.
+struct ConstDuration{T} end
+DensityInterface.DensityKind(::ConstDuration) = HasDensity()
+DensityInterface.logdensityof(::ConstDuration, x) = iszero(x) ? 0.0 : -Inf
+Base.rand(::AbstractRNG, ::ConstDuration{T}) where {T} = zero(T)
+
+@testset "Duration convention" begin
+    @testset "Shift relationship" begin
+        for dist in (Poisson(2.0), Geometric(0.3), NegativeBinomial(4.0, 0.5))
+            for k in -2:20
+                @test duration_logdensityof(dist, k) == logdensityof(dist, k - 1)
+            end
+            # Support is pushed onto the strictly positive integers.
+            @test duration_logdensityof(dist, 0) == -Inf
+            @test duration_logdensityof(dist, -1) == -Inf
+            @test duration_logdensityof(dist, 1) > -Inf
+        end
+    end
+
+    @testset "Normalization over sojourn times" begin
+        for dist in (Poisson(2.0), Geometric(0.3), NegativeBinomial(4.0, 0.5))
+            pmf_sum = sum(exp(duration_logdensityof(dist, k)) for k in 1:5000)
+            @test pmf_sum ≈ 1.0 atol = 1e-6
+        end
+    end
+
+    @testset "Sampling is shifted by one" begin
+        rng = StableRNG(63)
+        # mean(dist) + 1, since the sojourn is (the sampled value) + 1.
+        for (dist, expected_mean) in (
+            (Poisson(2.0), 3.0),        # mean 2 -> 3
+            (Geometric(0.25), 4.0),     # mean 3 -> 4
+            (NegativeBinomial(4.0, 0.5), 5.0),  # mean 4 -> 5
+        )
+            samples = [rand_duration(rng, dist) for _ in 1:10_000]
+            @test all(s isa Integer for s in samples)
+            @test all(s >= 1 for s in samples)
+            @test mean(samples) ≈ expected_mean rtol = 0.05
+        end
+    end
+
+    @testset "Sample type is preserved" begin
+        # The `+1` shift must keep the sampled type.
+        for T in (Int, Int32, Float32)
+            d = ConstDuration{T}()
+            @test rand_duration(StableRNG(1), d) === one(T)
+            @test rand_duration(d) === one(T)  # default rng overload
+        end
+    end
+
+    @testset "Custom user distribution (interface only)" begin
+        d = ConstDuration{Int}()
+        @test duration_logdensityof(d, 1) == 0.0
+        @test duration_logdensityof(d, 2) == -Inf
+        @test duration_logdensityof(d, 0) == -Inf
+        @test rand_duration(StableRNG(1), d) == 1
+        @test rand_duration(d) == 1  # default rng overload
+    end
+end

test/runtests.jl

@@ -76,4 +76,8 @@ end
     @testset verbose = true "Miscellaneous" begin
         include("misc.jl")
     end
+
+    @testset verbose = true "Duration convention" begin
+        include("duration.jl")
+    end
 end