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