File size: 6,878 Bytes
31d1d63 | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | 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
|