File size: 4,322 Bytes
198bd6d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | 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
|