1"""The process and its conditional sampler.
3Model: x iid N(0, sigma^2) -> y = h * x (causal FIR, kernel length L)
4-> z = round(y). A zero-phase or otherwise shifted application of the same
5kernel gives the same law, so causal convolution loses no generality.
7ConditionalChain targets H(z_{M+1} | z_1..z_M): it fixes an observed past
8z_1..z_M and Gibbs-samples the latent x under the box constraints
9y_t in [z_t - 1/2, z_t + 1/2), emitting one exact draw of z_{M+1} per step.
10"""
12import numpy as np
13from scipy.special import ndtr, ndtri
16def truncated_std_normal(lo, hi, rng):
17 """Standard normal truncated to [lo, hi], elementwise, by inverse CDF.
18 Mirrored into the lower tail so the CDF differences keep precision."""
19 flip = lo > -hi # midpoint above 0 (robust to (-inf, inf) intervals)
20 a = np.where(flip, -hi, lo)
21 b = np.where(flip, -lo, hi)
22 fa = ndtr(a)
23 fb = ndtr(b)
24 u = fa + (fb - fa) * rng.random(a.shape)
25 x = ndtri(np.clip(u, 1e-300, 1 - 1e-16))
26 x = np.where(flip, -x, x)
27 return np.clip(x, lo, hi)
30class ConditionalChain:
31 """Stationary chain of exact draws of z_{M+1} given a fixed past z_1..z_M.
33 The constructor draws the past from the prior; the generating latents are
34 themselves an exact draw from p(x | z), so the Gibbs chain starts in
35 stationarity — no burn-in bias, only autocorrelation. draw(k) advances the
36 chain k steps (thin sweeps each) and returns the k sampled z_{M+1} values,
37 each marginally distributed exactly as z_{M+1} | z_1..z_M.
39 Each Gibbs conditional x_i | rest is N(0, sigma^2) truncated to the
40 interval read off the <= L constraint boxes x_i appears in. Coordinates a
41 multiple of L apart share no constraint, so each of the L "colors"
42 updates as one vectorized block.
43 """
45 def __init__(self, kernel, sigma, past, rng=None, thin=1):
46 h = np.asarray(kernel, dtype=float)
47 if h.ndim != 1 or h.size == 0:
48 raise ValueError('kernel must be a nonempty 1-D array')
49 if sigma <= 0:
50 raise ValueError('sigma must be positive')
51 if past < 1:
52 raise ValueError('past must be >= 1')
53 self.h = h
54 self.sigma = float(sigma)
55 self.thin = int(thin)
56 self.rng = np.random.default_rng() if rng is None else rng
58 L = h.size
59 M = int(past)
60 self.L, self.M = L, M
62 # The past, with its true latents as the (stationary) chain start.
63 x = self.sigma * self.rng.standard_normal(M + L - 1)
64 y = np.convolve(x, h, mode='valid')
65 self.z = np.floor(y + 0.5)
66 self.x = x
68 # Boxes and y live in padded arrays so that every coordinate x_i sees
69 # exactly L constraint rows (rows outside the data are unconstrained).
70 P = L - 1
71 self.ypad = np.zeros(M + 2 * P)
72 self.lo = np.full(M + 2 * P, -np.inf)
73 self.hi = np.full(M + 2 * P, np.inf)
74 self.lo[P:P + M] = self.z - 0.5
75 self.hi[P:P + M] = self.z + 0.5
77 # Row i+j (padded) carries coefficient h[j] for coordinate i.
78 self.classes = [np.arange(c0, M + L - 1, L) for c0 in range(L)]
79 self.rowmats = [idx[:, None] + np.arange(L)[None, :]
80 for idx in self.classes]
81 self.nonzero = h != 0
82 # z_{M+1} = round(hr_head @ x[M:] + h[0] * x_free), x_free fresh.
83 self.hr_head = h[::-1][:-1]
85 def _sweep(self):
86 h, x, sigma = self.h, self.x, self.sigma
87 M, L = self.M, self.L
88 P = L - 1
89 self.ypad[P:P + M] = np.convolve(x, h, mode='valid') # kill fp drift
90 for idx, rows in zip(self.classes, self.rowmats):
91 r = self.ypad[rows] - np.outer(x[idx], h)
92 with np.errstate(divide='ignore', invalid='ignore'):
93 b1 = (self.lo[rows] - r) / h[None, :]
94 b2 = (self.hi[rows] - r) / h[None, :]
95 xlo = np.where(h[None, :] > 0, b1, b2)
96 xhi = np.where(h[None, :] > 0, b2, b1)
97 xlo[:, ~self.nonzero] = -np.inf
98 xhi[:, ~self.nonzero] = np.inf
99 xlo = xlo.max(axis=1)
100 xhi = xhi.min(axis=1)
101 xnew = truncated_std_normal(xlo / sigma, xhi / sigma, self.rng) * sigma
102 self.ypad[rows] += (xnew - x[idx])[:, None] * h[None, :]
103 x[idx] = xnew
105 def draw(self, k):
106 """The next k samples of z_{M+1}, continuing the chain."""
107 out = np.empty(k, dtype=np.int64)
108 M = self.M
109 for i in range(k):
110 for _ in range(self.thin):
111 self._sweep()
112 c = float(self.hr_head @ self.x[M:]) if self.L > 1 else 0.0
113 y_next = c + self.sigma * self.h[0] * self.rng.standard_normal()
114 out[i] = int(np.floor(y_next + 0.5))
115 return out