| |
| |
|
|
| from __future__ import annotations |
|
|
| import math |
| from dataclasses import dataclass |
|
|
| import numpy as np |
|
|
|
|
| P01_TWO_SIDED = 2.576 |
|
|
|
|
| def standard_error_from_tailstd( |
| tailstd: float, |
| *, |
| n_systems: int, |
| batch_size: int, |
| ) -> float: |
| if tailstd is None or not math.isfinite(float(tailstd)): |
| return float("nan") |
| denominator = max(1, int(n_systems) * int(batch_size)) |
| return float(tailstd) / math.sqrt(denominator) |
|
|
|
|
| def welch_band_walkermean( |
| tailstd_a: float, |
| tailstd_b: float, |
| *, |
| n_systems_a: int, |
| batch_size_a: int, |
| n_systems_b: int, |
| batch_size_b: int, |
| z: float = P01_TWO_SIDED, |
| ) -> float: |
| se_a = standard_error_from_tailstd( |
| tailstd_a, |
| n_systems=n_systems_a, |
| batch_size=batch_size_a, |
| ) |
| se_b = standard_error_from_tailstd( |
| tailstd_b, |
| n_systems=n_systems_b, |
| batch_size=batch_size_b, |
| ) |
| if not (math.isfinite(se_a) and math.isfinite(se_b)): |
| return float("nan") |
| return float(z) * math.sqrt(se_a * se_a + se_b * se_b) |
|
|
|
|
| def select_winner_per_physical( |
| R, |
| tailstd, |
| beam_logp, |
| batch_per_candidate, |
| *, |
| z: float = P01_TWO_SIDED, |
| ucb_z: float = 2.0, |
| ): |
| R = np.asarray(R, dtype=float) |
| tailstd = np.asarray(tailstd, dtype=float) |
| beam_logp = np.asarray(beam_logp, dtype=float) |
| P, K = R.shape |
| Bp = int(batch_per_candidate) |
| winner_idx = np.zeros(P, dtype=np.int64) |
| tie_mask = np.zeros((P, K), dtype=bool) |
| bands = np.full((P, K), np.nan, dtype=float) |
| se = np.array( |
| [ |
| [ |
| standard_error_from_tailstd( |
| tailstd[p, k], |
| n_systems=1, |
| batch_size=Bp, |
| ) |
| for k in range(K) |
| ] |
| for p in range(P) |
| ], |
| dtype=float, |
| ) |
| reason: list[str] = [] |
| for p in range(P): |
| Rp, sp, lp = R[p], tailstd[p], beam_logp[p] |
| finite = np.isfinite(Rp) |
| if not finite.any(): |
| w = int(np.argmax(lp)) |
| winner_idx[p] = w |
| tie_mask[p, w] = True |
| reason.append("energy_unavailable") |
| continue |
| istar = int(np.argmin(np.where(finite, Rp, np.inf))) |
| for k in range(K): |
| band = welch_band_walkermean( |
| sp[k], |
| sp[istar], |
| n_systems_a=1, |
| batch_size_a=Bp, |
| n_systems_b=1, |
| batch_size_b=Bp, |
| z=z, |
| ) |
| bands[p, k] = band |
| if k == istar: |
| tie_mask[p, k] = True |
| elif finite[k] and math.isfinite(band) and abs(Rp[k] - Rp[istar]) < band: |
| tie_mask[p, k] = True |
| tie_idx = np.where(tie_mask[p])[0] |
| se_tie = se[p, tie_idx] |
| ucb = Rp[tie_idx] + ucb_z * np.where(np.isfinite(se_tie), se_tie, np.inf) |
| w = int(tie_idx[int(np.argmin(ucb))]) if np.any(np.isfinite(ucb)) else istar |
| winner_idx[p] = w |
| reason.append("argmin_energy" if w == istar else "tie_broken_by_ucb") |
| return winner_idx, tie_mask, reason, bands, se |
|
|
|
|
| @dataclass(frozen=True, slots=True) |
| class ChannelMetrics: |
| mean: float |
| standard_error: float |
| walker_tail_std: float |
| local_energy_std: float |
| lag1_autocorrelation: float | None |
|
|
| def as_dict(self) -> dict[str, float | None]: |
| return { |
| "mean": self.mean, |
| "standard_error": self.standard_error, |
| "walker_tail_std": self.walker_tail_std, |
| "local_energy_std": self.local_energy_std, |
| "lag1_autocorrelation": self.lag1_autocorrelation, |
| } |
|
|
|
|
| class EnergyWindow: |
| def __init__(self, steps: int, systems: int, batch_size: int): |
| self.steps = int(steps) |
| self.systems = int(systems) |
| self.batch_size = int(batch_size) |
| shape = (self.steps, self.systems, self.batch_size) |
| self.values = { |
| "total": np.empty(shape, dtype=np.float32), |
| "exchange": np.empty(shape, dtype=np.float32), |
| "field": np.empty(shape, dtype=np.float32), |
| } |
| self.sums = { |
| "total": np.zeros((self.systems, self.batch_size), dtype=np.float32), |
| "exchange": np.zeros((self.systems, self.batch_size), dtype=np.float32), |
| "field": np.zeros((self.systems, self.batch_size), dtype=np.float32), |
| } |
| self.count = 0 |
|
|
| def push(self, total, exchange, field) -> None: |
| if self.count >= self.steps: |
| raise IndexError("energy window is full") |
| for name, value in ( |
| ("total", total), |
| ("exchange", exchange), |
| ("field", field), |
| ): |
| array = np.asarray(value).real.astype(np.float32, copy=False) |
| expected = (self.systems, self.batch_size) |
| if array.shape != expected: |
| raise ValueError( |
| f"{name} energy must have shape {expected}, got {array.shape}" |
| ) |
| self.values[name][self.count] = array |
| self.sums[name] += array |
| self.count += 1 |
|
|
| def tail_means(self, channel: str) -> np.ndarray: |
| if self.count < 1: |
| raise RuntimeError("energy window is empty") |
| return (self.sums[channel] / float(self.count)).astype(np.float32, copy=False) |
|
|
| def tail_mean(self, channel: str, system: int) -> float: |
| return float(np.mean(self.tail_means(channel)[system])) |
|
|
| def tail_std(self, channel: str, system: int) -> float: |
| per_walker = np.asarray(self.tail_means(channel)[system]).ravel() |
| if per_walker.size < 2: |
| return float("nan") |
| return float(np.std(per_walker, ddof=1)) |
|
|
| def metrics(self, channel: str) -> ChannelMetrics: |
| if self.systems != 1: |
| raise ValueError("bare eval metrics require one physical system") |
| samples = self.values[channel][: self.count, 0] |
| per_walker = self.tail_means(channel)[0] |
| mean = float(np.mean(per_walker)) |
| tailstd = ( |
| float(np.std(per_walker, ddof=1)) if per_walker.size >= 2 else float("nan") |
| ) |
| se = standard_error_from_tailstd( |
| tailstd, |
| n_systems=1, |
| batch_size=self.batch_size, |
| ) |
| step_std = np.std(samples, axis=1) |
| gap_sq = float(np.nanmean(step_std**2)) |
| local_std = ( |
| math.sqrt(gap_sq) |
| if math.isfinite(gap_sq) and gap_sq >= 0.0 |
| else float("nan") |
| ) |
| step_means = np.mean(samples, axis=1) |
| finite = step_means[np.isfinite(step_means)] |
| autocorrelation = None |
| if len(finite) >= 20: |
| centered = finite - np.mean(finite) |
| denominator = float(np.sum(centered * centered)) |
| if denominator > 0.0: |
| autocorrelation = float( |
| np.sum(centered[:-1] * centered[1:]) / denominator |
| ) |
| return ChannelMetrics( |
| mean=mean, |
| standard_error=se, |
| walker_tail_std=tailstd, |
| local_energy_std=local_std, |
| lag1_autocorrelation=autocorrelation, |
| ) |
|
|
