| from __future__ import annotations | |
| import numpy as np | |
| class InfeasibleBridgeError(ValueError): | |
| """The requested starting distribution has no conditioned path law.""" | |
| def transition_matrix(P): | |
| P = np.asarray(P, dtype=float) | |
| if P.ndim != 2 or P.shape[0] != P.shape[1] or not len(P): | |
| raise ValueError("P must be a nonempty square transition matrix") | |
| if not np.isfinite(P).all() or (P < 0).any(): | |
| raise ValueError("P must contain finite nonnegative probabilities") | |
| sums = P.sum(axis=1) | |
| if not np.all((sums == 0) | np.isclose(sums, 1.0, rtol=1e-8, atol=1e-12)): | |
| raise ValueError("Each nonzero row of P must sum to one") | |
| return P | |
| def horizon_length(N): | |
| if isinstance(N, (bool, np.bool_)) or not isinstance(N, (int, np.integer)) or N < 1: | |
| raise ValueError("The sampling horizon must be a positive integer") | |
| return int(N) | |
| def state_indices(states, size, name): | |
| values = np.asarray(list(states)) | |
| if values.ndim != 1 or not values.size or values.dtype.kind not in "iu": | |
| raise ValueError(f"{name} must be a nonempty sequence of integer state indices") | |
| if ((values < 0) | (values >= size)).any(): | |
| raise ValueError(f"{name} contains an out-of-range state index") | |
| return values.astype(int) | |
| def weights_array(weights, N, size): | |
| if weights is None: | |
| weights = np.ones(size) | |
| weights = np.asarray(weights, dtype=float) | |
| if weights.shape == (size,): | |
| weights = np.broadcast_to(weights, (N, size)) | |
| if weights.shape != (N, size): | |
| raise ValueError(f"Weights must have shape ({size},) or ({N}, {size})") | |
| if not np.isfinite(weights).all() or (weights < 0).any(): | |
| raise ValueError("Weights must be finite and nonnegative") | |
| return weights | |
| def _log_nonnegative(values): | |
| out = np.full(values.shape, -np.inf) | |
| np.log(values, out=out, where=values > 0) | |
| return out | |
| def _from_log(log_values): | |
| with np.errstate(under="ignore", over="ignore"): | |
| values = np.exp(log_values) | |
| if not np.isfinite(values).all() or np.any(np.isfinite(log_values) & (values == 0)): | |
| raise FloatingPointError( | |
| "Backward values exceed floating-point range; no fallback or support floor was applied" | |
| ) | |
| return values | |
| def backward(P, weights, N, terminal, *, rescale=True): | |
| """Return h, with optional per-time scaling that leaves kernels unchanged.""" | |
| P = transition_matrix(P) | |
| N = horizon_length(N) | |
| weights = weights_array(weights, N, len(P)) | |
| terminal = np.asarray(terminal, dtype=float) | |
| if ( | |
| terminal.shape != (len(P),) | |
| or not np.isfinite(terminal).all() | |
| or (terminal < 0).any() | |
| ): | |
| raise ValueError("Terminal weights must be a finite nonnegative vector") | |
| if not np.any(terminal > 0): | |
| raise ValueError("At least one terminal state must have positive weight") | |
| rows, cols = np.nonzero(P) | |
| log_edges = np.log(P[rows, cols]) | |
| log_weights = _log_nonnegative(weights) | |
| log_h = np.full((N + 1, len(P)), -np.inf) | |
| log_h[N] = _log_nonnegative(terminal) | |
| for n in range(N - 1, -1, -1): | |
| edge_terms = log_edges + log_weights[n, cols] + log_h[n + 1, cols] | |
| np.logaddexp.at(log_h[n], rows, edge_terms) | |
| if rescale: | |
| finite = np.isfinite(log_h[n]) | |
| if finite.any(): | |
| log_h[n, finite] -= np.logaddexp.reduce(log_h[n, finite]) | |
| return _from_log(log_h) | |
| def kernel(P, weights, h_next): | |
| """Normalize feasible rows; return exact zero rows for infeasible states.""" | |
| P = transition_matrix(P) | |
| weights = weights_array(weights, 1, len(P))[0] | |
| h_next = np.asarray(h_next, dtype=float) | |
| if h_next.shape != (len(P),) or not np.isfinite(h_next).all() or (h_next < 0).any(): | |
| raise ValueError("The backward function must be a finite nonnegative vector") | |
| rows, cols = np.nonzero(P) | |
| scores = ( | |
| np.log(P[rows, cols]) | |
| + _log_nonnegative(weights)[cols] | |
| + _log_nonnegative(h_next)[cols] | |
| ) | |
| normalizers = np.full(len(P), -np.inf) | |
| np.logaddexp.at(normalizers, rows, scores) | |
| valid = np.isfinite(scores) & np.isfinite(normalizers[rows]) | |
| M = np.zeros_like(P) | |
| with np.errstate(under="ignore"): | |
| M[rows[valid], cols[valid]] = np.exp(scores[valid] - normalizers[rows[valid]]) | |
| sums = M.sum(axis=1, keepdims=True) | |
| np.divide(M, sums, out=M, where=sums > 0) | |
| return M | |
| def build(P, weights, targets, N, *, starts=None): | |
| P = transition_matrix(P) | |
| N = horizon_length(N) | |
| targets = np.unique(state_indices(targets, len(P), "Target set")) | |
| weights = weights_array(weights, N, len(P)) | |
| terminal = np.zeros(len(P)) | |
| terminal[targets] = 1.0 / len(targets) | |
| h = backward(P, weights, N, terminal) | |
| if starts is not None: | |
| starts = state_indices(starts, len(P), "Starting set") | |
| infeasible = starts[h[0, starts] == 0] | |
| if infeasible.size: | |
| raise InfeasibleBridgeError( | |
| f"No positive-weight path reaches the terminal set at horizon {N} " | |
| f"from starting states {infeasible.tolist()}" | |
| ) | |
| return [kernel(P, weights[n], h[n + 1]) for n in range(N)], h | |
| def checked_starts(kernels, starts): | |
| if not len(kernels): | |
| raise ValueError("At least one transition kernel is required") | |
| first = np.asarray(kernels[0]) | |
| starts = state_indices(starts, len(first), "Starting set") | |
| rows = first[starts] | |
| if not np.isfinite(rows).all() or (rows < 0).any(): | |
| raise ValueError("Invalid conditioned probabilities for a starting state") | |
| bad = starts[rows.sum(axis=1) == 0] | |
| if bad.size: | |
| raise InfeasibleBridgeError( | |
| f"Starting states {bad.tolist()} have h_0=0 at the requested horizon; " | |
| "the prescribed starting distribution cannot be sampled" | |
| ) | |
| return starts | |
| def sample(kernels, starts, n_paths=200, seed=0, *, rng=None): | |
| starts = checked_starts(kernels, starts) | |
| if ( | |
| isinstance(n_paths, (bool, np.bool_)) | |
| or not isinstance(n_paths, (int, np.integer)) | |
| or n_paths < 0 | |
| ): | |
| raise ValueError("n_paths must be a nonnegative integer") | |
| rng = np.random.default_rng(seed) if rng is None else rng | |
| size = len(kernels[0]) | |
| paths = [] | |
| for _ in range(n_paths): | |
| i = int(rng.choice(starts)) | |
| path = [i] | |
| for n, matrix in enumerate(kernels): | |
| p = np.asarray(matrix[i], dtype=float) | |
| total = p.sum() | |
| if not np.isfinite(p).all() or (p < 0).any() or not total > 0: | |
| raise InfeasibleBridgeError( | |
| f"Invalid or infeasible conditioned row at step {n}, state {i}" | |
| ) | |
| i = int(rng.choice(size, p=p / total)) | |
| path.append(i) | |
| paths.append(np.asarray(path, dtype=int)) | |
| return paths | |