Source code for asymmetry.core.fitting.experiment_design

"""Bayesian experimental design for parameter trending.

Given a converged trend-model fit (``fit_parameter_model``), suggest where
the next measurement should go — and how many events it should collect —
to most constrain the model's parameters. The acquisition is the Laplace
(Gaussian-posterior) expected information gain for a single new datum:

    IG(x) = 1/2 log[1 + g(x)^T Sigma g(x) / sigma_new^2(x)]        (D-optimal)

    DeltaVar_k(x) = (Sigma g(x))_k^2 / (sigma_new^2(x) + g(x)^T Sigma g(x))
                                                                    (c-optimal)

where g(x) is the model's parameter sensitivity at the fitted values and
Sigma the fit covariance. The rank-one c-optimal update is exact for a
linear model; for strongly nonlinear models (critical exponents) it ranks
candidates correctly but *underestimates* the realised post-fit variance —
see ``docs/studies/prototype-results.md`` — so the counting recommendation
must be calibrated with ``calibrate_suggestion`` before being shown to a
user.

Everything here is GUI-free and pure numpy; refits happen only inside
``calibrate_suggestion``.
"""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass, replace

import numpy as np
from numpy.typing import NDArray
from scipy.optimize import linear_sum_assignment

from asymmetry.core.fitting.parameter_models import (
    ParameterCompositeModel,
    fit_parameter_model,
)
from asymmetry.core.fitting.parameters import ParameterSet

#: Condition number above which the covariance is flagged as ill-conditioned.
_COVARIANCE_CONDITION_WARN = 1.0e4

#: Proximity factor (in units of the larger per-candidate sigma) below which two
#: series' predicted values are flagged as at-risk of misassignment (§3.1). A
#: new datum near a predicted crossing may attach to the wrong curve, breaking
#: the rank-one update's assumption that it belongs to the series it constrains.
_ASSIGNMENT_RISK_K = 2.0

#: Utility ratio outside which the argmax is flagged as gradient-unstable
#: (recomputed with a 10x finite-difference step).
_GRADIENT_STABILITY_SPAN = (0.5, 2.0)

#: Realised/predicted variance ratio above which the analytic figure is
#: called out as optimistic in calibration warnings.
_CALIBRATION_OPTIMISM_WARN = 2.0

#: Fraction of failed calibration refits above which a warning is added.
_CALIBRATION_FAILURE_FRACTION_WARN = 0.25


[docs] @dataclass(frozen=True) class NextPointSuggestion: """Utility curve and recommendation for the next measurement. ``utility`` is ``DeltaVar_k`` (variance reduction of the target parameter) in c-optimal mode, or the information gain ``IG`` in D-optimal mode (``target is None``). Two degenerate shapes exist, both explained by ``warnings``: unusable inputs (bad covariance, no noise model, ...) yield empty arrays with ``best_x`` NaN, while a valid setup where no candidate carries information keeps the full curve (all-zero ``utility``) with ``best_x`` NaN. """ x_candidates: NDArray[np.float64] utility: NDArray[np.float64] extrapolated: NDArray[np.bool_] best_x: float target: str | None sigma_new: NDArray[np.float64] #: Predicted post-measurement sigma of the target at ``best_x`` for a #: reference-statistics point (c-optimal mode only). Rank-one estimate — #: optimistic near critical points; calibrate before display. predicted_post_sigma: float | None = None #: N/N_ref needed at ``best_x`` to reach ``sigma_goal`` (rank-one #: estimate). ``0.0`` means the goal is already met without a new point. events_factor_to_target: float | None = None target_unreachable: bool = False #: sigma of the target in the N -> infinity limit of a single new point #: at ``best_x`` — the floor set by the other parameters' uncertainty. floor_sigma: float | None = None #: Per-candidate assignment-risk flag (multi-series acquisition only). ``True`` #: where two series' predicted values lie within ``_ASSIGNMENT_RISK_K`` sigma #: of each other, so a new datum there may attach to the wrong curve. ``None`` #: for single-series suggestions, which have no crossing to worry about. risk_mask: NDArray[np.bool_] | None = None warnings: tuple[str, ...] = ()
[docs] @dataclass(frozen=True) class SuggestionCalibration: """Monte-Carlo check of a suggestion: simulate the point, refit, measure. ``realized_post_sigma`` is the median post-fit sigma of the target over the successful trials; the p16/p84 band shows its spread. ``calibration_ratio`` is realised variance / rank-one predicted variance (> 1 means the analytic figure was optimistic). """ x_new: float events_factor: float target: str n_trials: int n_failed: int realized_post_sigma: float realized_post_sigma_p16: float realized_post_sigma_p84: float predicted_post_sigma: float | None = None calibration_ratio: float | None = None warnings: tuple[str, ...] = ()
def _empty_suggestion(target: str | None, warnings: list[str]) -> NextPointSuggestion: empty = np.zeros(0, dtype=float) return NextPointSuggestion( x_candidates=empty, utility=empty, extrapolated=np.zeros(0, dtype=bool), best_x=float("nan"), target=target, sigma_new=empty, warnings=tuple(warnings), ) def _validated_covariance( covariance: tuple[list[str], list[list[float]]] | None, warnings: list[str], ) -> tuple[list[str], NDArray[np.float64]] | None: if not covariance: warnings.append("No covariance matrix available from the fit.") return None names, matrix = covariance cov = np.asarray(matrix, dtype=float) if cov.ndim != 2 or cov.shape != (len(names), len(names)) or not names: warnings.append("Covariance matrix shape does not match its parameter names.") return None if not np.all(np.isfinite(cov)): warnings.append("Covariance matrix contains non-finite entries.") return None cov = 0.5 * (cov + cov.T) eigenvalues = np.linalg.eigvalsh(cov) max_eig = float(eigenvalues[-1]) min_eig = float(eigenvalues[0]) if max_eig <= 0.0 or min_eig < -1e-12 * max(max_eig, 1.0): warnings.append("Covariance matrix is not positive definite.") return None if min_eig <= 0.0 or max_eig / min_eig > _COVARIANCE_CONDITION_WARN: warnings.append( "Fit covariance is ill-conditioned (strong parameter correlations); " "utilities are approximate." ) return list(names), cov def _empirical_sigma( x_data: NDArray[np.float64], y_err: NDArray[np.float64], warnings: list[str], ) -> tuple[NDArray[np.float64], NDArray[np.float64]] | None: """Return (sorted x, sigma at those x) from the finite positive errors.""" xd = np.asarray(x_data, dtype=float) ed = np.asarray(y_err, dtype=float) ok = np.isfinite(xd) & np.isfinite(ed) & (ed > 0.0) if not np.any(ok): warnings.append("No finite positive y-errors to build a noise model from.") return None order = np.argsort(xd[ok]) return xd[ok][order], ed[ok][order] def _sensitivities( model: ParameterCompositeModel, values: dict[str, float], free_names: list[str], x: NDArray[np.float64], step_scale: float = 1.0e-6, ) -> NDArray[np.float64]: """Central-difference d f/d theta_j over the whole candidate grid. Returns shape ``(len(free_names), len(x))``. Non-finite entries are zeroed — a candidate in a region where the model is undefined simply carries no information about that parameter. """ grads = np.zeros((len(free_names), len(x)), dtype=float) for j, name in enumerate(free_names): theta = float(values[name]) h = step_scale * max(1.0, abs(theta)) hi = dict(values) lo = dict(values) hi[name] = theta + h lo[name] = theta - h f_hi = np.asarray(model.function(x, **hi), dtype=float) f_lo = np.asarray(model.function(x, **lo), dtype=float) grads[j] = (f_hi - f_lo) / (2.0 * h) return np.nan_to_num(grads, nan=0.0, posinf=0.0, neginf=0.0) def _utility_curve( grads: NDArray[np.float64], cov: NDArray[np.float64], sigma_new: NDArray[np.float64], target_index: int | None, ) -> NDArray[np.float64]: projected = cov @ grads # (n_params, n_candidates): Sigma g per candidate quadratic = np.einsum("jc,jc->c", grads, projected) # g^T Sigma g variance = sigma_new**2 if target_index is None: with np.errstate(invalid="ignore", divide="ignore"): utility = 0.5 * np.log1p(quadratic / variance) else: utility = projected[target_index] ** 2 / (variance + quadratic) return np.nan_to_num(utility, nan=0.0, posinf=0.0, neginf=0.0) def _candidate_grid( x_data: NDArray[np.float64], y_err: NDArray[np.float64], x_min: float, x_max: float, n_candidates: int, warnings: list[str], ) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64], NDArray[np.bool_]] | None: """Shared candidate-grid + empirical-noise construction. Returns ``(candidates, x_meas, sigma_new, extrapolated)`` or ``None`` if degenerate (appending to ``warnings`` in that case). ``candidates`` includes the in-range measured x values (so "re-measure" competes) plus a uniform grid of ``n_candidates`` points over ``[x_min, x_max]``. """ noise = _empirical_sigma(x_data, y_err, warnings) if noise is None: return None x_meas, sigma_meas = noise if not (np.isfinite(x_min) and np.isfinite(x_max)) or x_max <= x_min: warnings.append("Invalid candidate range.") return None grid = np.linspace(float(x_min), float(x_max), int(n_candidates)) in_range = x_meas[(x_meas >= x_min) & (x_meas <= x_max)] candidates = np.unique(np.concatenate([grid, in_range])) sigma_new = np.interp(candidates, x_meas, sigma_meas) extrapolated = (candidates < x_meas[0]) | (candidates > x_meas[-1]) if np.any(extrapolated): warnings.append( "Candidate range extends beyond the measured points; suggestions " "there rely entirely on the assumed model form." ) return candidates, x_meas, sigma_new, extrapolated
[docs] def suggest_next_point( model: ParameterCompositeModel, parameters: ParameterSet, covariance: tuple[list[str], list[list[float]]] | None, x_data: NDArray[np.float64], y_err: NDArray[np.float64], x_min: float, x_max: float, *, target: str | None = None, sigma_goal: float | None = None, n_candidates: int = 257, ) -> NextPointSuggestion: """Rank candidate x values by how much a new point there constrains the fit. ``covariance`` is the ``(names, matrix)`` pair from ``ParameterModelFitResult.covariance``; ``x_data``/``y_err`` are the measured series (the per-point errors define the empirical noise model for the hypothetical new point at reference statistics). ``target`` selects c-optimality for that parameter; ``None`` means D-optimal. ``sigma_goal`` (c-optimal only) additionally solves for the event-count factor N/N_ref needed at the best x — a rank-one planning figure that should be calibrated with :func:`calibrate_suggestion` before display. Never raises on degenerate inputs: returns an empty suggestion with an explanatory warning instead. """ warnings: list[str] = [] validated = _validated_covariance(covariance, warnings) if validated is None: return _empty_suggestion(target, warnings) free_names, cov = validated if target is not None and target not in free_names: warnings.append(f"Target parameter {target!r} is not a free fitted parameter.") return _empty_suggestion(target, warnings) noise = _empirical_sigma(x_data, y_err, warnings) if noise is None: return _empty_suggestion(target, warnings) x_meas_all, _sigma_meas = noise if not (np.isfinite(x_min) and np.isfinite(x_max)) or x_max <= x_min: warnings.append("Invalid candidate range.") return _empty_suggestion(target, warnings) if len(x_meas_all) <= len(free_names) + 1: warnings.append( "Fit is barely constrained (few points for the number of free " "parameters); the suggestion is unreliable — consider a coarse scan." ) grid_result = _candidate_grid(x_data, y_err, x_min, x_max, n_candidates, warnings) assert grid_result is not None # already validated noise + range above candidates, x_meas, sigma_new, extrapolated = grid_result values = {p.name: float(p.value) for p in parameters} missing = [n for n in free_names if n not in values] if missing: warnings.append(f"Fitted values missing for parameters: {missing}.") return _empty_suggestion(target, warnings) target_index = free_names.index(target) if target is not None else None grads = _sensitivities(model, values, free_names, candidates) utility = _utility_curve(grads, cov, sigma_new, target_index) if not np.any(utility > 0.0): warnings.append("No candidate carries information about the fit parameters.") return NextPointSuggestion( x_candidates=candidates, utility=utility, extrapolated=extrapolated, best_x=float("nan"), target=target, sigma_new=sigma_new, warnings=tuple(warnings), ) best_index = int(np.argmax(utility)) best_x = float(candidates[best_index]) # Gradient-stability probe at the winner only: near a model domain # boundary (e.g. x ~ fitted T_c) the finite-difference sensitivity is # finite but step-dependent, so a 10x step materially moving the utility # marks the suggestion as approximate. probe_x = np.array([best_x], dtype=float) probe_sigma = sigma_new[best_index : best_index + 1] coarse = _utility_curve( _sensitivities(model, values, free_names, probe_x, step_scale=1.0e-5), cov, probe_sigma, target_index, )[0] fine = utility[best_index] low, high = _GRADIENT_STABILITY_SPAN if coarse <= 0.0 or not (low <= coarse / fine <= high): warnings.append( "The suggested point sits near a model domain boundary; its " "utility is step-sensitive and approximate." ) predicted_post_sigma: float | None = None events_factor: float | None = None unreachable = False floor_sigma: float | None = None if target_index is not None: g_best = grads[:, best_index] projected = cov @ g_best quadratic = float(g_best @ projected) prior_var = float(cov[target_index, target_index]) gain = float(projected[target_index]) ** 2 post_var = prior_var - gain / (float(sigma_new[best_index]) ** 2 + quadratic) predicted_post_sigma = float(np.sqrt(max(post_var, 0.0))) if quadratic > 0.0: floor_var = prior_var - gain / quadratic floor_sigma = float(np.sqrt(max(floor_var, 0.0))) if sigma_goal is not None and np.isfinite(sigma_goal) and sigma_goal > 0.0: goal_var = float(sigma_goal) ** 2 if prior_var <= goal_var: events_factor = 0.0 elif quadratic <= 0.0 or gain <= 0.0: unreachable = True else: needed = gain / (prior_var - goal_var) - quadratic if needed <= 0.0: unreachable = True else: events_factor = float(sigma_new[best_index]) ** 2 / needed if unreachable: warnings.append( "The precision goal cannot be reached with a single new " "point at the suggested x — the other parameters' " "uncertainty sets the floor." ) return NextPointSuggestion( x_candidates=candidates, utility=utility, extrapolated=extrapolated, best_x=best_x, target=target, sigma_new=sigma_new, predicted_post_sigma=predicted_post_sigma, events_factor_to_target=events_factor, target_unreachable=unreachable, floor_sigma=floor_sigma, warnings=tuple(warnings), )
[docs] @dataclass(frozen=True) class SeriesSpec: """One series' inputs to a multi-series acquisition. Mirrors :func:`suggest_next_point`'s per-series argument types: a fitted ``model``/``parameters`` pair, the ``(names, matrix)`` fit ``covariance`` (``None`` when the fit produced none), and the measured ``x_data``/``y_err`` that define the empirical new-point noise model. One physical measurement at a candidate ``x`` is assumed to yield one datum for *every* series (§3.1). """ model: ParameterCompositeModel parameters: ParameterSet covariance: tuple[list[str], list[list[float]]] | None x_data: NDArray[np.float64] y_err: NDArray[np.float64]
def _series_label(labels: Sequence[str] | None, index: int) -> str: if labels is not None and 0 <= index < len(labels): return str(labels[index]) return f"series {index}" def _assignment_risk_mask( series: Sequence[SeriesSpec], candidates: NDArray[np.float64], ) -> tuple[NDArray[np.bool_], list[str]]: """Flag candidates where two series' predictions risk misassignment (§3.1). A candidate is at-risk when any two series' predicted values lie within ``_ASSIGNMENT_RISK_K * max(sigma_m, sigma_n)`` of each other — near a predicted crossing the new datum's components are unlabelled and may attach to the wrong curve, so the rank-one information gain (which assumes correct attachment) is optimistic there. Series without a usable noise model are excluded (no sigma to compare against); with fewer than two usable series the mask is all-``False`` (a lone curve has nothing to cross). """ predictions: list[NDArray[np.float64]] = [] sigmas: list[NDArray[np.float64]] = [] for spec in series: noise = _empirical_sigma(spec.x_data, spec.y_err, []) if noise is None: continue values = {p.name: float(p.value) for p in spec.parameters} raw = np.asarray(spec.model.function(candidates, **values), dtype=float) x_meas, sigma_meas = noise predictions.append(np.nan_to_num(raw, nan=np.inf)) sigmas.append(np.interp(candidates, x_meas, sigma_meas)) risk = np.zeros(candidates.shape, dtype=bool) warnings: list[str] = [] if len(predictions) < 2: return risk, warnings for a in range(len(predictions)): for b in range(a + 1, len(predictions)): threshold = _ASSIGNMENT_RISK_K * np.maximum(sigmas[a], sigmas[b]) gap = np.abs(predictions[a] - predictions[b]) risk |= np.isfinite(gap) & (gap <= threshold) if np.any(risk): flagged = candidates[risk] warnings.append( f"Predicted curves approach within {_ASSIGNMENT_RISK_K:g}σ over " f"x ∈ [{float(flagged.min()):.4g}, {float(flagged.max()):.4g}]; a new " f"point there risks feeding the wrong curve (misassignment)." ) return risk, warnings
[docs] def suggest_next_point_multi( series: Sequence[SeriesSpec], x_min: float, x_max: float, *, target: tuple[int, str] | None = None, sigma_goal: float | None = None, n_candidates: int = 257, labels: Sequence[str] | None = None, ) -> NextPointSuggestion: """Rank candidate x by how much a shared new point constrains several series. Each entry of ``series`` is a fitted :class:`SeriesSpec`; one measurement at ``x`` adds one datum to *every* series. With ``target is None`` the utility is the D-optimal information-gain sum over series (§3.1): U(x) = Σ_m ½ log[1 + g_m(x)ᵀ Σ_m g_m(x) / σ_m²(x)] With ``target = (series_index, param_name)`` the objective is that one series' c-optimal variance reduction for ``param_name`` — the other series still receive the datum but do not enter the objective — and the ``sigma_goal``/event-count solve applies to the target exactly as in :func:`suggest_next_point`. Either way the returned suggestion carries a per-candidate :attr:`NextPointSuggestion.risk_mask` (§3.1): the rank-one update assumes the new datum attaches to the right curve, so candidates near a predicted crossing are flagged. A series whose covariance is invalid (:func:`_validated_covariance`) contributes nothing to the utility and adds a warning naming it (by ``labels`` when supplied, else by index); the suggestion still succeeds if at least one series is usable. All-unusable inputs degrade to an empty suggestion. Never raises. """ warnings: list[str] = [] if not series: warnings.append("No series supplied for multi-series acquisition.") return _empty_suggestion(None, warnings) if target is not None: target_index, target_param = target if not (0 <= target_index < len(series)): warnings.append(f"Target series index {target_index} is out of range.") return _empty_suggestion(target_param, warnings) spec = series[target_index] base = suggest_next_point( spec.model, spec.parameters, spec.covariance, spec.x_data, spec.y_err, x_min, x_max, target=target_param, sigma_goal=sigma_goal, n_candidates=n_candidates, ) if base.x_candidates.size == 0: return base risk, risk_warnings = _assignment_risk_mask(series, base.x_candidates) return replace( base, risk_mask=risk, warnings=base.warnings + tuple(risk_warnings), ) # D-optimal: sum the per-series information gains over the usable series. usable: list[ tuple[list[str], NDArray[np.float64], dict[str, float], SeriesSpec, NDArray, NDArray] ] = [] for i, spec in enumerate(series): cov_warnings: list[str] = [] validated = _validated_covariance(spec.covariance, cov_warnings) if validated is None: warnings.append(f"{_series_label(labels, i)} has no usable covariance and is skipped.") continue noise = _empirical_sigma(spec.x_data, spec.y_err, []) if noise is None: warnings.append(f"{_series_label(labels, i)} has no noise model and is skipped.") continue free_names, cov = validated values = {p.name: float(p.value) for p in spec.parameters} if any(name not in values for name in free_names): warnings.append(f"{_series_label(labels, i)} is missing fitted parameter values.") continue usable.append((free_names, cov, values, spec, noise[0], noise[1])) if not usable: warnings.append("No usable series for the multi-series acquisition.") return _empty_suggestion(None, warnings) if not (np.isfinite(x_min) and np.isfinite(x_max)) or x_max <= x_min: warnings.append("Invalid candidate range.") return _empty_suggestion(None, warnings) grid = np.linspace(float(x_min), float(x_max), int(n_candidates)) in_range = [ x_meas[(x_meas >= x_min) & (x_meas <= x_max)] for _fn, _c, _v, _s, x_meas, _sm in usable ] candidates = np.unique(np.concatenate([grid, *in_range])) utility = np.zeros_like(candidates) extrapolated = np.zeros(candidates.shape, dtype=bool) sigma_stack: list[NDArray[np.float64]] = [] for free_names, cov, values, spec, x_meas, sigma_meas in usable: sigma_new = np.interp(candidates, x_meas, sigma_meas) sigma_stack.append(sigma_new) extrapolated |= (candidates < x_meas[0]) | (candidates > x_meas[-1]) grads = _sensitivities(spec.model, values, free_names, candidates) utility = utility + _utility_curve(grads, cov, sigma_new, None) if np.any(extrapolated): warnings.append( "Candidate range extends beyond the measured points; suggestions " "there rely entirely on the assumed model form." ) sigma_new_mean = np.mean(np.vstack(sigma_stack), axis=0) risk, risk_warnings = _assignment_risk_mask(series, candidates) warnings.extend(risk_warnings) if not np.any(utility > 0.0): warnings.append("No candidate carries information about the fit parameters.") return NextPointSuggestion( x_candidates=candidates, utility=utility, extrapolated=extrapolated, best_x=float("nan"), target=None, sigma_new=sigma_new_mean, risk_mask=risk, warnings=tuple(warnings), ) best_index = int(np.argmax(utility)) return NextPointSuggestion( x_candidates=candidates, utility=utility, extrapolated=extrapolated, best_x=float(candidates[best_index]), target=None, sigma_new=sigma_new_mean, risk_mask=risk, warnings=tuple(warnings), )
[docs] def calibrate_suggestion( model: ParameterCompositeModel, parameters: ParameterSet, x_data: NDArray[np.float64], y_data: NDArray[np.float64], y_err: NDArray[np.float64], x_new: float, *, target: str, events_factor: float = 1.0, n_trials: int = 50, seed: int = 0, predicted_post_sigma: float | None = None, ) -> SuggestionCalibration: """Monte-Carlo-calibrate a suggestion by simulating the point and refitting. Each trial draws the hypothetical new datum from the fitted model plus Gaussian noise at the empirical sigma for ``x_new`` (scaled by ``1/sqrt(events_factor)``), appends it to the measured series, refits from the fitted values, and records the post-fit sigma of ``target``. The median over trials is the calibrated post-measurement sigma — the honest counterpart of the rank-one ``predicted_post_sigma``, which is known to be optimistic near critical points (see ``docs/studies/prototype-results.md``). ``y_err`` must be the effective per-point sigmas the trend fit was weighted with (already resolved through the fit's error mode). """ warnings: list[str] = [] xd = np.asarray(x_data, dtype=float) yd = np.asarray(y_data, dtype=float) ed = np.asarray(y_err, dtype=float) noise = _empirical_sigma(xd, ed, warnings) if noise is None or not np.isfinite(x_new) or events_factor <= 0.0 or n_trials < 1: warnings.append("Calibration inputs are degenerate; nothing simulated.") return SuggestionCalibration( x_new=float(x_new), events_factor=float(events_factor), target=target, n_trials=0, n_failed=0, realized_post_sigma=float("nan"), realized_post_sigma_p16=float("nan"), realized_post_sigma_p84=float("nan"), predicted_post_sigma=predicted_post_sigma, warnings=tuple(warnings), ) x_meas, sigma_meas = noise sigma_new = float(np.interp(x_new, x_meas, sigma_meas)) / float(np.sqrt(events_factor)) values = {p.name: float(p.value) for p in parameters} y_model_new = float(np.asarray(model.function(np.array([x_new]), **values))[0]) realized: list[float] = [] n_failed = 0 for trial in range(int(n_trials)): rng = np.random.default_rng(seed + trial) y_new = y_model_new + rng.normal(0.0, sigma_new) result = fit_parameter_model( np.append(xd, x_new), np.append(yd, y_new), np.append(ed, sigma_new), model, parameters, ) err = result.uncertainties.get(target) if result.success else None if err is not None and np.isfinite(err) and err > 0.0: realized.append(float(err)) else: n_failed += 1 if not realized: warnings.append("All calibration refits failed; no realised sigma available.") return SuggestionCalibration( x_new=float(x_new), events_factor=float(events_factor), target=target, n_trials=int(n_trials), n_failed=n_failed, realized_post_sigma=float("nan"), realized_post_sigma_p16=float("nan"), realized_post_sigma_p84=float("nan"), predicted_post_sigma=predicted_post_sigma, warnings=tuple(warnings), ) sigmas = np.array(realized, dtype=float) median = float(np.median(sigmas)) ratio: float | None = None if predicted_post_sigma is not None and predicted_post_sigma > 0.0: ratio = float(median**2 / predicted_post_sigma**2) if ratio > _CALIBRATION_OPTIMISM_WARN: warnings.append( "The analytic post-fit sigma is substantially optimistic here " "(strong model nonlinearity); trust the calibrated figure." ) if n_failed > n_trials * _CALIBRATION_FAILURE_FRACTION_WARN: warnings.append("A large fraction of calibration refits failed.") return SuggestionCalibration( x_new=float(x_new), events_factor=float(events_factor), target=target, n_trials=int(n_trials), n_failed=n_failed, realized_post_sigma=median, realized_post_sigma_p16=float(np.percentile(sigmas, 16.0)), realized_post_sigma_p84=float(np.percentile(sigmas, 84.0)), predicted_post_sigma=predicted_post_sigma, calibration_ratio=ratio, warnings=tuple(warnings), )
[docs] def calibrate_events_for_goal( model: ParameterCompositeModel, parameters: ParameterSet, x_data: NDArray[np.float64], y_data: NDArray[np.float64], y_err: NDArray[np.float64], x_new: float, *, target: str, sigma_goal: float, initial_events_factor: float = 1.0, n_trials: int = 30, seed: int = 0, max_events_factor: float = 1024.0, max_iterations: int = 6, ) -> tuple[float | None, SuggestionCalibration]: """Find the event-count factor whose *calibrated* post-sigma meets the goal. Multiplicative search starting from ``initial_events_factor`` (use the rank-one ``events_factor_to_target`` when available): each step scales the factor by the shortfall ``(median/goal)^2`` — exact if the new point's variance dominates, conservative otherwise. Returns ``(factor, last_calibration)``; factor is ``None`` when the goal is not met by ``max_events_factor`` (the single-point floor applies). """ factor = max(float(initial_events_factor), 1.0e-3) calibration: SuggestionCalibration | None = None for iteration in range(max(1, int(max_iterations))): calibration = calibrate_suggestion( model, parameters, x_data, y_data, y_err, x_new, target=target, events_factor=factor, n_trials=n_trials, seed=seed + 1000 * iteration, ) median = calibration.realized_post_sigma if not np.isfinite(median): return None, calibration if median <= sigma_goal: return factor, calibration if factor >= max_events_factor: break factor = min(factor * min((median / sigma_goal) ** 2, 8.0), max_events_factor) assert calibration is not None return None, calibration
#: Relative-to-noise threshold below which the discrimination utility is #: considered "everywhere within noise" (candidate models agree). _DISCRIMINATION_NOISE_FLOOR = 1.0e-12
[docs] def suggest_discriminating_point( lead_model: ParameterCompositeModel, lead_parameters: ParameterSet, alternatives: Sequence[tuple[ParameterCompositeModel, ParameterSet]], x_data: NDArray[np.float64], y_err: NDArray[np.float64], x_min: float, x_max: float, *, n_candidates: int = 257, ) -> NextPointSuggestion: """Rank candidate x values by how well they discriminate between models. Scores the ``lead_model`` (the currently favoured fit) against **every** entry in ``alternatives`` (each a fitted ``(model, parameters)`` pair over the same series), taking the elementwise max over alternatives — per the TAS-AI lock-in lesson (§4/§8.1 of the study), a leader must be tested against all live competitors, not just the runner-up: U_disc(x) = max_i [f_lead(x) - f_i(x)]^2 / (2 sigma_new^2(x)) Returns a :class:`NextPointSuggestion` with ``target=None`` (this mode has no single target parameter) and no event-count/goal fields set. Never raises: degenerate inputs (no alternatives, no noise model, bad range) degrade to an empty suggestion with an explanatory warning, and a utility that is everywhere within noise of zero is flagged rather than silently returning a meaningless argmax. """ warnings: list[str] = [] if not alternatives: warnings.append("No alternative models supplied to discriminate against.") return _empty_suggestion(None, warnings) grid_result = _candidate_grid(x_data, y_err, x_min, x_max, n_candidates, warnings) if grid_result is None: return _empty_suggestion(None, warnings) candidates, _x_meas, sigma_new, extrapolated = grid_result lead_values = {p.name: float(p.value) for p in lead_parameters} non_finite_seen = False def _evaluate(model: ParameterCompositeModel, values: dict[str, float]) -> NDArray[np.float64]: nonlocal non_finite_seen raw = np.asarray(model.function(candidates, **values), dtype=float) if not np.all(np.isfinite(raw)): non_finite_seen = True return np.nan_to_num(raw, nan=0.0, posinf=0.0, neginf=0.0) f_lead = _evaluate(lead_model, lead_values) variance = sigma_new**2 utility = np.zeros_like(candidates) for alt_model, alt_parameters in alternatives: alt_values = {p.name: float(p.value) for p in alt_parameters} f_alt = _evaluate(alt_model, alt_values) with np.errstate(invalid="ignore", divide="ignore"): pairwise = (f_lead - f_alt) ** 2 / (2.0 * variance) pairwise = np.nan_to_num(pairwise, nan=0.0, posinf=0.0, neginf=0.0) utility = np.maximum(utility, pairwise) if non_finite_seen: warnings.append( "One or more candidate models produced non-finite values on the " "candidate grid; those points were treated as carrying no " "discriminating information." ) noise_scale = float(np.max(variance)) if variance.size else 0.0 relative_floor = _DISCRIMINATION_NOISE_FLOOR * max(noise_scale, 1.0) if not np.any(utility > relative_floor): warnings.append( "candidate models agree within noise everywhere in this range — no discriminating point" ) return NextPointSuggestion( x_candidates=candidates, utility=utility, extrapolated=extrapolated, best_x=float("nan"), target=None, sigma_new=sigma_new, warnings=tuple(warnings), ) best_index = int(np.argmax(utility)) best_x = float(candidates[best_index]) return NextPointSuggestion( x_candidates=candidates, utility=utility, extrapolated=extrapolated, best_x=best_x, target=None, sigma_new=sigma_new, warnings=tuple(warnings), )
[docs] def set_matching_divergence( hypothesis_a: Sequence[tuple[ParameterCompositeModel, ParameterSet]], hypothesis_b: Sequence[tuple[ParameterCompositeModel, ParameterSet]], sigma: Sequence[NDArray[np.float64]], x_candidates: NDArray[np.float64], ) -> NDArray[np.float64]: """Per-candidate min-cost set-matching divergence between two hypotheses (§3.3). Each hypothesis is an equal-length list of ``(model, parameters)`` curves. At every candidate ``x`` the two hypotheses' predicted ``N``-vectors are matched one-to-one (Hungarian) to minimise the total weighted cost, with per-pair cost [f_i^A(x) − f_j^B(x)]² / (2 σ_i²(x)) (``sigma[i]`` is series ``i`` of hypothesis A's per-candidate noise). Because a new run's components are *unlabelled*, two assignment hypotheses are only distinguishable through the sets of values they predict — this is that set-level discrepancy. It is zero where the hypotheses predict the same value set (e.g. at a crossing where the labellings coincide) and peaks where they imply genuinely different curve sets. For ``N = 1`` it reduces exactly to the labelled discrimination integrand of :func:`suggest_discriminating_point`. Returns ``U(x)`` over ``x_candidates``. Non-finite predictions or zero sigmas are treated as a large finite cost so the matching stays well-posed. """ x = np.asarray(x_candidates, dtype=float) n = len(hypothesis_a) if n != len(hypothesis_b): raise ValueError("hypothesis_a and hypothesis_b must have the same length") if n == 0 or x.size == 0: return np.zeros(x.shape, dtype=float) def _stack(hypothesis: Sequence[tuple[ParameterCompositeModel, ParameterSet]]) -> NDArray: rows = [] for model, params in hypothesis: values = {p.name: float(p.value) for p in params} rows.append(np.asarray(model.function(x, **values), dtype=float)) return np.vstack(rows) # (N, G) f_a = _stack(hypothesis_a) f_b = _stack(hypothesis_b) sig = np.vstack([np.asarray(s, dtype=float) for s in sigma]) # (N, G) variance = 2.0 * sig**2 large = 1e18 utility = np.zeros(x.shape, dtype=float) for g in range(x.size): diff = f_a[:, g][:, None] - f_b[:, g][None, :] # (N, N) with np.errstate(invalid="ignore", divide="ignore"): cost = diff**2 / variance[:, g][:, None] cost = np.nan_to_num(cost, nan=large, posinf=large, neginf=large) rows, cols = linear_sum_assignment(cost) utility[g] = float(cost[rows, cols].sum()) return utility
[docs] def aic_weights(chi_squareds: Sequence[float], n_free_params: Sequence[int]) -> list[float]: """Akaike-weight a set of candidate models from their fit chi-squareds. ``AIC_i = chi2_i + 2 p_i``; weights are ``exp(-(AIC_i - min AIC)/2)`` normalised to sum to 1 (the minimum is subtracted before exponentiating for numerical safety, per Burnham & Anderson). BIC is deliberately not used here — its ``p * ln(n)`` penalty drifts under the sequential-n growth of an in-progress trend series (§4 of the study). A model with non-finite chi-squared gets weight ``0.0`` and is excluded from the normalisation; if every model is non-finite, returns a list of ``0.0`` (there is nothing to rank). Raises ``ValueError`` if the two input lists differ in length — a mismatch here is a programming error, not a user-data problem. """ if len(chi_squareds) != len(n_free_params): raise ValueError( f"chi_squareds and n_free_params must have the same length " f"({len(chi_squareds)} != {len(n_free_params)})." ) if not chi_squareds: return [] aic = np.array( [c + 2.0 * p for c, p in zip(chi_squareds, n_free_params, strict=True)], dtype=float ) finite = np.isfinite(aic) if not np.any(finite): return [0.0] * len(aic) min_aic = float(np.min(aic[finite])) raw_weights = np.zeros_like(aic) raw_weights[finite] = np.exp(-(aic[finite] - min_aic) / 2.0) total = float(np.sum(raw_weights)) if total <= 0.0: return [0.0] * len(aic) return (raw_weights / total).tolist()
[docs] def cost_weighted_utility( x_candidates: NDArray[np.float64], utility: NDArray[np.float64], x_current: float, *, count_time: float, up_rate: float, down_rate: float, gamma: float = 0.7, ) -> NDArray[np.float64]: """Weight a utility curve by movement + counting cost (TAS-AI's IG^gamma/t). ``move_time(x) = up_rate * (x - x_current)`` for ``x > x_current``, else ``down_rate * (x_current - x)``; all rates/times share a unit (hours). Returns ``clip(utility, 0)**gamma / (count_time + move_time)``. Pure array-in/array-out with defensive guards rather than raising: an invalid cost model (``count_time <= 0``, a negative rate, or a non-finite ``x_current``) must not silently distort the ranking, so the original ``utility`` is returned unchanged in that case. The GUI is responsible for validating these inputs before calling; this function does not emit warnings. """ x = np.asarray(x_candidates, dtype=float) u = np.asarray(utility, dtype=float) if count_time <= 0.0 or up_rate < 0.0 or down_rate < 0.0 or not np.isfinite(x_current): return u move_time = np.where(x > x_current, up_rate * (x - x_current), down_rate * (x_current - x)) denominator = count_time + move_time return np.clip(u, 0.0, None) ** gamma / denominator