Importance Sampling with Mixture Proposals

CompetingClocks can compute path log-likelihoods under multiple distributions simultaneously, enabling importance sampling with mixture proposals for continuous-time stochastic processes. This page describes how to use the likelihood_cnt feature and integrate it with Gen.jl for Bayesian inference.

Path-Space Importance Sampling

Let $X$ be an event path generated by a continuous-time stochastic system (SIR, assembly line, reliability, etc.). In CompetingClocks, a path is a sequence of (clock key, firing time) pairs:

\[x = \bigl\{ (k_1, t_1), (k_2, t_2), \ldots, (k_n, t_n) \bigr\}.\]

Define:

• Target path law $p_\theta(x)$: the GSMP dynamics you want to characterize (parameterized by $\theta$)
• Proposal path law $q_\phi(x)$: biased dynamics that make rare events more frequent

Standard path-space importance sampling:

1. Draw $x \sim q_\phi$
2. Compute the importance weight:

   \[w(x) = \frac{p_\theta(x)}{q_\phi(x)}, \quad \log w(x) = \log p_\theta(x) - \log q_\phi(x)\]

3. Estimate expectations under $p_\theta$ via weighted averages $\sum_i w_i f(x_i) / \sum_i w_i$

CompetingClocks supports this directly:

• Create a sampler with path_likelihood=true and likelihood_cnt=K
• For each enable! call, pass a vector of $K$ distributions for that clock
• The sampler simulates using one distribution (selected via sample_from_distribution!) but accumulates log-likelihoods for the realized path under all $K$ distributions
• Call pathloglikelihood(sampler, end_time) to retrieve the $K$ log path-densities

Mixture Proposals

Introduce $M$ proposals $q_1, \ldots, q_M$ and form a mixture:

\[q_{\text{mix}}(x) = \sum_{m=1}^M \alpha_m q_m(x), \quad \alpha_m \ge 0, \quad \sum_m \alpha_m = 1\]

To sample $x \sim q_{\text{mix}}$:

1. Draw component index $J \sim \text{Categorical}(\alpha)$
2. Draw $x \sim q_J$

The mixture importance weight is:

\[\log w(x) = \log p_\theta(x) - \log q_{\text{mix}}(x) = \log p_\theta(x) - \log\Bigl( \sum_{m=1}^M \alpha_m e^{\log q_m(x)} \Bigr)\]

Compute the denominator numerically via logsumexp:

\[\log q_{\text{mix}}(x) = \text{logsumexp}_m\bigl(\log \alpha_m + \log q_m(x)\bigr)\]

CompetingClocks implementation:

• Set likelihood_cnt = 1 + M (target plus $M$ proposals)
• For each clock, pass a distribution vector where index 1 is the target $p_\theta$ and indices 2 through $M+1$ are the proposals
• Call sample_from_distribution!(sampler, 1 + j) to sample from proposal $j$
• After simulation, pathloglikelihood returns $1 + M$ log-likelihoods

Example: Poisson Process with Mixture Proposals

Consider a homogeneous Poisson process on $[0, T]$ with rate $\lambda$. We define:

• Target: rate $\lambda$
• Proposals: rates $\lambda \cdot m_1, \ldots, \lambda \cdot m_M$ where each $m_i$ is a multiplicative bias
• Mixture weights: uniform $\alpha = (1/M, \ldots, 1/M)$

The following function simulates a path under the mixture proposal and returns both the event count and the log importance weight:

using Random
using Statistics: mean, var
using Distributions: Exponential, Categorical, UnivariateDistribution, Gamma, Poisson, logpdf
using CompetingClocks
using StatsFuns: logsumexp

const ClockKey = Int

"""
    simulate_poisson_path_with_mixture(λ_target, multipliers, α, T_end, rng)

Simulate a Poisson process path up to time `T_end` using a mixture of biased
rates. Returns `(count, log_weight)` where:
- `count`: number of events by time `T_end`
- `log_weight`: log p(path | λ_target) - log q_mix(path)
"""
function simulate_poisson_path_with_mixture(
    λ_target::Float64,
    multipliers::Vector{Float64},
    α::Vector{Float64},
    T_end::Float64,
    rng::AbstractRNG
)
    M = length(multipliers)
    @assert length(α) == M
    @assert isapprox(sum(α), 1.0)
    K = 1 + M  # target + M proposals

    # Create sampler with K path likelihoods
    sampler = SamplingContext(ClockKey, Float64, rng;
        method=FirstToFireMethod(),
        path_likelihood=true,
        likelihood_cnt=K
    )

    # Select which proposal to sample from
    j = rand(rng, Categorical(α))
    sample_from_distribution!(sampler, 1 + j)

    # Distribution vector: [target, proposal_1, ..., proposal_M]
    dists = Vector{UnivariateDistribution}(undef, K)
    dists[1] = Exponential(1.0 / λ_target)
    for m in 1:M
        dists[1 + m] = Exponential(1.0 / (λ_target * multipliers[m]))
    end

    enable!(sampler, 1, dists)

    count = 0
    when, which = next(sampler)
    while !isnothing(which) && when <= T_end
        fire!(sampler, which, when)
        count += 1
        enable!(sampler, 1, dists)
        when, which = next(sampler)
    end

    # Extract log-likelihoods
    logliks = pathloglikelihood(sampler, T_end)
    log_p = logliks[1]
    log_qs = logliks[2:end]

    # Mixture weight via logsumexp
    log_q_mix = logsumexp(log.(α) .+ collect(log_qs))
    log_weight = log_p - log_q_mix

    return count, log_weight
end

Verifying Correctness

The weighted sample mean should match the true expectation $E[N_T] = \lambda T$:

λ_target = 1.5
T_end = 10.0
multipliers = [0.5, 1.0, 2.0]
α = fill(1.0 / length(multipliers), length(multipliers))

n_samples = 1000
counts = zeros(Int, n_samples)
log_weights = zeros(Float64, n_samples)

rng = Xoshiro(42)
for i in 1:n_samples
    counts[i], log_weights[i] = simulate_poisson_path_with_mixture(
        λ_target, multipliers, α, T_end, rng
    )
end

# Normalize weights
max_lw = maximum(log_weights)
weights = exp.(log_weights .- max_lw)
weights ./= sum(weights)

weighted_mean = sum(weights .* counts)  # ≈ 15.0
true_mean = λ_target * T_end            # = 15.0

The unweighted mean will differ from the true mean (due to the biased proposals), but the importance-weighted mean converges to the correct value.

Integration with Gen.jl

Gen.jl provides a probabilistic programming framework where you can combine CompetingClocks simulations with parameter inference. The key insight: treat the CompetingClocks simulation as a black box that returns a summary and an importance weight, then incorporate that weight into Gen's trace score.

Model Structure

Consider a model with:

• Parameter $\theta$ with prior $p(\theta)$
• Latent path $X$ with law $p_\theta(x)$
• Observation $y$ with likelihood $p(y \mid x, \theta)$

Rather than representing each event as a Gen choice, let CompetingClocks handle the path simulation. For each trace:

1. Sample $\theta$ from the prior (Gen choice)
2. Simulate $x$ using the mixture proposal via CompetingClocks
3. Compute $\log w(x) = \log p_\theta(x) - \log q_{\text{mix}}(x)$
4. Add the weight to Gen's score so inference sees the correct target density

Implementation

using Gen

@gen function poisson_inference_model(
    T_end::Float64,
    multipliers::Vector{Float64},
    α::Vector{Float64},
    rng::AbstractRNG
)
    # Prior on target rate
    λ ~ gamma(2.0, 1.0)

    # Simulate path with mixture proposal
    count, log_w = simulate_poisson_path_with_mixture(λ, multipliers, α, T_end, rng)

    return (count=count, log_weight=log_w, λ=λ)
end

Inference Loop

Since the path importance weight must be incorporated into the particle weights, implement a manual importance sampling loop:

# Generate synthetic observation
true_λ = 1.5
obs_count = rand(Xoshiro(999), Poisson(true_λ * T_end))

# Importance sampling over λ
n_particles = 1000
rng = Xoshiro(42)

λ_samples = zeros(Float64, n_particles)
total_log_weights = zeros(Float64, n_particles)

for i in 1:n_particles
    # Sample from prior
    λ = rand(rng, Gamma(2.0, 1.0))
    λ_samples[i] = λ

    # Simulate with mixture proposal
    sim_count, log_w_path = simulate_poisson_path_with_mixture(
        λ, multipliers, α, T_end, rng
    )

    # Observation likelihood
    log_obs = logpdf(Poisson(max(sim_count, 1)), obs_count)

    # Total weight: path correction + observation
    total_log_weights[i] = log_w_path + log_obs
end

# Normalize and compute posterior statistics
max_lw = maximum(total_log_weights)
weights = exp.(total_log_weights .- max_lw)
weights ./= sum(weights)

posterior_mean_λ = sum(weights .* λ_samples)
ess = 1.0 / sum(weights.^2)  # Effective sample size

This gives two layers of importance sampling:

1. Inner layer (CompetingClocks): corrects from $q_{\text{mix}}$ to $p_\lambda$
2. Outer layer (Gen/manual): corrects from prior $p(\lambda)$ to posterior $p(\lambda \mid y)$

Variance Reduction

Mixture proposals can substantially improve effective sample size compared to a single proposal. The benefit comes from adaptive coverage: when one proposal generates unlikely paths under the target, others may assign higher probability.

Compare effective sample sizes:

The mixture achieves roughly 4× higher ESS by hedging across multiple rate biases.

Generalizing to Other Models

To adapt this pattern to your CompetingClocks model:

1. Identify target parameters $\theta$ that define $p_\theta$ — these become Gen random choices or inference targets

2. For each clock, construct a distribution vector [p_dist, q1_dist, ..., qM_dist] and call the vectorized enable!(sampler, key, dists)

3. Select the sampling distribution with sample_from_distribution!(sampler, index) before the simulation loop

4. After simulation, call pathloglikelihood(sampler, end_time) and compute the mixture weight via logsumexp

5. Wrap in a helper that returns (path_summary, log_weight)

6. In Gen or your inference code:

   • Call the helper
   • Incorporate log_weight into particle weights or Gen factors
   • Add observation likelihoods
   • Use any inference algorithm (importance sampling, SMC, MCMC)

This separation keeps variance-reduction logic (multiple proposals, mixture weights, common random numbers) in CompetingClocks while exposing a clean interface to probabilistic programming frameworks.