| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
| from typing import Any, Callable, Generic, Sequence, TypeVar |
| import math |
| import random |
|
|
| S = TypeVar("S") |
| A = TypeVar("A") |
|
|
|
|
| def normalize_logweights(logw: Sequence[float]) -> list[float]: |
| if not logw: |
| return [] |
| finite = [w for w in logw if math.isfinite(w)] |
| if not finite: |
| return [1.0 / len(logw)] * len(logw) |
| m = max(finite) |
| ws = [math.exp(w - m) if math.isfinite(w) else 0.0 for w in logw] |
| z = sum(ws) |
| if not math.isfinite(z) or z <= 0.0: |
| return [1.0 / len(logw)] * len(logw) |
| return [w / z for w in ws] |
|
|
|
|
| def effective_sample_size(probs: Sequence[float]) -> float: |
| s2 = sum(float(p) ** 2 for p in probs) |
| return 0.0 if s2 <= 0.0 else 1.0 / s2 |
|
|
|
|
| def systematic_resample(probs: Sequence[float], rng: random.Random) -> list[int]: |
| n = len(probs) |
| if n == 0: |
| return [] |
| cdf: list[float] = [] |
| total = 0.0 |
| for p in probs: |
| total += float(p) |
| cdf.append(total) |
| if total <= 0.0 or not math.isfinite(total): |
| return list(range(n)) |
| cdf = [x / total for x in cdf] |
| cdf[-1] = 1.0 |
| u0 = rng.random() / n |
| out: list[int] = [] |
| j = 0 |
| for i in range(n): |
| u = u0 + i / n |
| while j < n - 1 and u > cdf[j]: |
| j += 1 |
| out.append(j) |
| return out |
|
|
|
|
| @dataclass(frozen=True) |
| class Proposal(Generic[A]): |
| action: A |
| probability: float |
|
|
|
|
| @dataclass |
| class Particle(Generic[S]): |
| state: S |
| logw: float = 0.0 |
| history: list[Any] = field(default_factory=list) |
|
|
|
|
| @dataclass(frozen=True) |
| class SMCConfig: |
| n_particles: int = 32 |
| n_steps: int = 3 |
| ess_resample_fraction: float = 0.5 |
| seed: int = 0 |
|
|
|
|
| class ParticleFilter(Generic[S, A]): |
| """Importance-corrected SMC runtime. |
| |
| `propose(state, t)` returns a discrete proposal distribution q(a|state). |
| `transition(state, action, t)` produces the next state. |
| `log_target_increment(prev, action, nxt, t)` supplies the incremental |
| log target density. The importance update is log w += log target - log q. |
| Resampling is never performed after the final weighting step, so the final |
| posterior weights are preserved. |
| """ |
|
|
| def __init__(self, config: SMCConfig | None = None): |
| self.config = config or SMCConfig() |
|
|
| def run( |
| self, |
| initial_state: S, |
| propose: Callable[[S, int], Sequence[Proposal[A]]], |
| transition: Callable[[S, A, int], S], |
| log_target_increment: Callable[[S, A, S, int], float], |
| ) -> list[Particle[S]]: |
| cfg = self.config |
| rng = random.Random(cfg.seed) |
| particles = [Particle(state=initial_state) for _ in range(cfg.n_particles)] |
| for t in range(cfg.n_steps): |
| nxt: list[Particle[S]] = [] |
| for p in particles: |
| proposals = [x for x in propose(p.state, t) if x.probability > 0.0] |
| if not proposals: |
| raise ValueError("proposal distribution is empty") |
| z = sum(x.probability for x in proposals) |
| probs = [x.probability / z for x in proposals] |
| r = rng.random() |
| acc = 0.0 |
| j = len(proposals) - 1 |
| for i, pr in enumerate(probs): |
| acc += pr |
| if r <= acc: |
| j = i |
| break |
| selected = proposals[j] |
| q = probs[j] |
| new_state = transition(p.state, selected.action, t) |
| inc = float(log_target_increment(p.state, selected.action, new_state, t)) |
| nxt.append(Particle(new_state, p.logw + inc - math.log(q), p.history + [selected.action])) |
| particles = nxt |
| probs_now = normalize_logweights([p.logw for p in particles]) |
| if t < cfg.n_steps - 1 and effective_sample_size(probs_now) < cfg.ess_resample_fraction * len(particles): |
| ids = systematic_resample(probs_now, rng) |
| particles = [Particle(particles[i].state, 0.0, list(particles[i].history)) for i in ids] |
| final_probs = normalize_logweights([p.logw for p in particles]) |
| for p, pr in zip(particles, final_probs): |
| setattr(p, "posterior", pr) |
| return particles |
|
|
