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. thin may be
38 reassigned between draws (e.g. probe at thin=1, then thin by the measured
39 autocorrelation time); stationarity is unaffected.
41 Each Gibbs conditional x_i | rest is N(0, sigma^2) truncated to the
42 interval read off the <= L constraint boxes x_i appears in. Coordinates a
43 multiple of L apart share no constraint, so each of the L "colors"
44 updates as one vectorized block.
45 """
47 def __init__(self, kernel, sigma, past, rng=None, thin=1):
48 h = np.asarray(kernel, dtype=float)
49 if h.ndim != 1 or h.size == 0:
50 raise ValueError('kernel must be a nonempty 1-D array')
51 if sigma <= 0:
52 raise ValueError('sigma must be positive')
53 if past < 1:
54 raise ValueError('past must be >= 1')
55 self.h = h
56 self.sigma = float(sigma)
57 self.thin = int(thin)
58 self.rng = np.random.default_rng() if rng is None else rng
60 L = h.size
61 M = int(past)
62 self.L, self.M = L, M
64 # The past, with its true latents as the (stationary) chain start.
65 x = self.sigma * self.rng.standard_normal(M + L - 1)
66 y = np.convolve(x, h, mode='valid')
67 self.z = np.floor(y + 0.5)
68 self.x = x
70 # Boxes and y live in padded arrays so that every coordinate x_i sees
71 # exactly L constraint rows (rows outside the data are unconstrained).
72 P = L - 1
73 self.ypad = np.zeros(M + 2 * P)
74 self.lo = np.full(M + 2 * P, -np.inf)
75 self.hi = np.full(M + 2 * P, np.inf)
76 self.lo[P:P + M] = self.z - 0.5
77 self.hi[P:P + M] = self.z + 0.5
79 # Row i+j (padded) carries coefficient h[j] for coordinate i.
80 self.classes = [np.arange(c0, M + L - 1, L) for c0 in range(L)]
81 self.rowmats = [idx[:, None] + np.arange(L)[None, :]
82 for idx in self.classes]
83 self.nonzero = h != 0
84 # z_{M+1} = round(hr_head @ x[M:] + h[0] * x_free), x_free fresh.
85 self.hr_head = h[::-1][:-1]
87 def _sweep(self):
88 h, x, sigma = self.h, self.x, self.sigma
89 M, L = self.M, self.L
90 P = L - 1
91 self.ypad[P:P + M] = np.convolve(x, h, mode='valid') # kill fp drift
92 for idx, rows in zip(self.classes, self.rowmats):
93 r = self.ypad[rows] - np.outer(x[idx], h)
94 with np.errstate(divide='ignore', invalid='ignore'):
95 b1 = (self.lo[rows] - r) / h[None, :]
96 b2 = (self.hi[rows] - r) / h[None, :]
97 xlo = np.where(h[None, :] > 0, b1, b2)
98 xhi = np.where(h[None, :] > 0, b2, b1)
99 xlo[:, ~self.nonzero] = -np.inf
100 xhi[:, ~self.nonzero] = np.inf
101 xlo = xlo.max(axis=1)
102 xhi = xhi.min(axis=1)
103 xnew = truncated_std_normal(xlo / sigma, xhi / sigma, self.rng) * sigma
104 self.ypad[rows] += (xnew - x[idx])[:, None] * h[None, :]
105 x[idx] = xnew
107 def draw(self, k):
108 """The next k samples of z_{M+1}, continuing the chain."""
109 out = np.empty(k, dtype=np.int64)
110 M = self.M
111 for i in range(k):
112 for _ in range(self.thin):
113 self._sweep()
114 c = float(self.hr_head @ self.x[M:]) if self.L > 1 else 0.0
115 y_next = c + self.sigma * self.h[0] * self.rng.standard_normal()
116 out[i] = int(np.floor(y_next + 0.5))
117 return out