Blocked Sampler Configuration

Implement a helper that builds a small Poisson switch model and runs sampling with explicit blocking: the two continuous rate variables are updated together in one block, and a discrete switch variable is updated with a caller-chosen random-walk proposal instead of relying on automatic step assignment.

Capabilities

Manual blocked sampling for the switch model

• Running on counts [2, 4, 6, 9, 3, 5] with default arguments returns acceptance entries for both "rates" and "switch" blocks, each between 0 and 1, and block_assignments lists ["rate_a", "rate_b"] under "rates" and ["switch"] under "switch". @test

Configurable discrete proposals

• Passing proposal="gaussian" runs the sampler and yields a result whose proposal field is "gaussian" while keeping both blocks present in block_assignments. @test
• Passing any proposal other than "heavy-tailed" or "gaussian" raises a ValueError that mentions the allowed options. @test

Reproducible outputs with seeding

• Running run_blocked_sampler twice with seed=7 on counts [2, 4, 6, 9, 3, 5] yields identical posterior_means for "rate_a" and "rate_b" to within 1e-6. @test

Implementation

@generates

API

from typing import Dict, List, Literal, Optional, TypedDict

class BlockedSamplingResult(TypedDict):
    posterior_means: Dict[str, float]
    acceptance: Dict[str, float]
    block_assignments: Dict[str, List[str]]
    proposal: Literal["heavy-tailed", "gaussian"]

def run_blocked_sampler(
    counts: List[int],
    draws: int = 500,
    tune: int = 300,
    proposal: Literal["heavy-tailed", "gaussian"] = "heavy-tailed",
    seed: Optional[int] = None,
) -> BlockedSamplingResult:
    """
    Build and sample a Poisson-rate switch model with manual blocked steps.
    Returns posterior means computed from the retained draws, acceptance rates
    for each sampling block, the variable grouping per block, and the proposal
    label actually used for the discrete switch.
    """

Dependencies { .dependencies }

pymc3 { .dependency }

Probabilistic programming library used to define the model and run custom blocked sampling with user-chosen proposals.

numpy { .dependency }

Used for numerical inputs, random seeding, and summary statistics.