Link with HiddenMarkovModels.jl?

[gdalle]

Hey @mchernys and @AntonOresten,
Following up on the discussion in JuliaRegistries/General#119495 to avoid cluttering that PR, what are $N$ and $M$ in the complexity $O(NM)$ that you mention for inference routines?

[murrellb]

N is sequence length and M is state space/alphabet size.

[gdalle]

If your transition matrix is sparse you can use a SparseMatrixCSC in HiddenMarkovModels and recover complexity $O(NC)$ where $C$ is the number of nonzero coefficients. In your case I assume $C = \theta(M)$?
Of course a dedicated algorithm is probably faster, I just wanted to point it out in case it's easier for you to offload maintenance of inference routines.

[murrellb]

Our transition matrix is not sparse. Any state can switch to any other state.

[gdalle]

Out of curiosity, what's the transition structure? Just for the fun of it I'm curious to see how hard it would be to get the right complexity on my end

[murrellb]

There are a few variants, but the main thing to exploit is that, if you switch states, the distribution over the state you switch to does not depend on the state you switch from.

[gdalle]

Ok so it's always rank-1 and all the rows are the same?

[murrellb]

No (I think) the matrix is full-rank. You can have high self-transition rates, and the identity transition matrix is the limiting case as the switching rate goes to zero. The distribution over the states you switch to is the same (excluding the current state) conditioned on a switching event.

[gdalle]

Oh right, so rank-1 + diagonal. Interesting. I'll give it a try

[gdalle]

With the latest release of my package, you can now specify the transition matrix as a lazy operator instead of an actual matrix. As a result, procedures like forward! should become allocation-free even in your case, where creating the transition matrix would be wasteful.
If you deem this interesting, we can think about adaptation to Viterbi's algorithm.

begin
	using Distributions
	using HiddenMarkovModels
	import HiddenMarkovModels as HMMs
	using LinearAlgebra
end

@kwdef struct Rank1HMM{T,D} <: AbstractHMM
	init::Vector{T}
	trans_self::T  # probability of staying in the same state
	trans_other::Vector{T}  # probability of transitioning to each state if we don't stay
	dists::Vector{D}
end

HMMs.initialization(hmm::Rank1HMM) = hmm.init

HMMs.obs_distributions(hmm::Rank1HMM) = hmm.dists

function HMMs.transition_matrix(hmm::Rank1HMM)
	return (
		hmm.trans_self .* I(length(hmm)) .+
		(1-hmm.trans_self) .* transpose(hmm.trans_other)
	)
end

function HMMs.predict_next_state!(
	p2::AbstractVector{<:Real}, hmm::Rank1HMM, p1::AbstractVector{<:Real}, ::Nothing
)
	p2 .= hmm.trans_self .* p2 .+ (1 - hmm.trans_self) .* hmm.trans_other
end

hmm = Rank1HMM(
	init=[0.3, 0.4, 0.3],
	trans_self=0.5,
	trans_other=[0.1, 0.5, 0.4],
	dists=[Normal(1.0), Normal(2.0), Normal(3.0)]
)

(; state_seq, obs_seq) = rand(hmm, 100)

forward(hmm, obs_seq)