"""
Canonical Stage 2b τ calibration settings.
``TauCalibrationSettings`` is the single source of truth for the Stage 2b
parameters across every surface:
* the public API signatures (``Pipeline.calibrate_tau`` (with
``shape="gaussian"`` for the τ_G variant) / ``Pipeline.recommend_shape`` and
the matching ``ftmwpipeline.api`` functions),
* the CLI ``--preset`` flag plus the existing per-knob flags,
* the resolution chain ``explicit > persisted > preset > recommended >
hard default``,
* the persisted canonical record in ``processing_parameters/stage2b_tau``,
* the YAML preset interchange format.
The dataclass mirrors the Stage 5 :class:`~ftmwpipeline.core.stage_fit_settings.StageFitSettings`
template: every field is ``Optional`` with ``None`` meaning *unset* (fall
through the resolution chain). A *resolved* instance (produced by
:func:`resolve`) has every field filled with a hard default if no layer
supplied a value.
The dataclass is structured into six sub-dataclasses grouping the knobs by
what they configure: ``stft``, ``polish``, ``aggregation``, ``band``,
``gaussian``, and ``recommendation``. The grouping maps 1:1 to HDF5
subgroups under ``processing_parameters/stage2b_tau`` so each sub-block is
independently inspectable. ``gaussian`` and ``recommendation`` overlap in
three fields by design (``snr_min`` / bounds / seeds); each consumer reads
from its own block so the Gaussian-twin τ calibration and the 3-way
shape-recommendation hook can evolve independent operating points.
The *recommended* layer of :func:`resolve` is reserved but currently
unused for Stage 2b -- Stage 2b is the originator of recommendations
(it produces the ``recommended_shape`` attr the Stage 5 resolver
consumes), not a downstream consumer of any upstream recommendation. The
layer is kept in the signature so a future upstream recommender (e.g. a
Stage 2 σ-driven ``snr_min`` suggestion) can land without API churn.
The ``_HARD_DEFAULTS`` nested dict mirrors the ``DEFAULT_*`` constants in
:mod:`ftmwpipeline.fitting.tau_calibration`. Those constants are still
imported by the kernel functions as their parameter defaults; once every
consumer reads from a resolved ``TauCalibrationSettings``, the constants
become docstring-only and can be removed.
This module is dependency-free within the package (stdlib + PyYAML for
preset interchange) so it can be imported from ``core`` without cycles.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Tuple, Union
from . import settings_framework as sf
from .knob_metadata import knob_field
# ---------------------------------------------------------------------------
# Sub-dataclasses (one per HDF5 subgroup / YAML block)
# ---------------------------------------------------------------------------
@dataclass
class StftSubSettings:
"""Sliding-active-window STFT knobs (shared by every Stage 2b consumer)."""
n_seg: Optional[int] = knob_field(
help="Number of non-overlapping STFT frames (window = T_full / n_seg).",
tier="primary",
inst_sensitivity="Y",
grid=(6, 8, 10, 14, 20),
cli=True,
argtype=int,
)
t_sigma: Optional[float] = knob_field(
help="Above-threshold SNR gate for per-frame signal detection "
"(contributor floor).",
tier="primary",
inst_sensitivity="Y",
grid=(3.0, 4.0, 5.0, 6.0, 8.0),
cli=True,
argtype=float,
)
tau_max_us: Optional[float] = knob_field(
help="Hard upper clip on recovered tau (saturation -> spur candidate); "
"unset -> derived.",
inst_sensitivity="maybe",
grid=(20.0, 40.0, 80.0),
cli=True,
argtype=float,
)
tau_max_factor: Optional[float] = knob_field(
help="tau_max as a multiple of the full-record duration when "
"tau_max_us is unset.",
inst_sensitivity="maybe",
grid=(3.0, 5.0, 8.0, 12.0),
)
rss_gate_factor: Optional[float] = knob_field(
help="Bad-fit gate strength (relative-or-absolute residual hybrid).",
inst_sensitivity="maybe",
grid=(3.0, 5.0, 8.0, 12.0),
cli=True,
argtype=float,
)
relative_gate_fraction: Optional[float] = knob_field(
help="Relative-RSS fraction below which a per-frame fit is accepted.",
inst_sensitivity="maybe",
grid=(0.02, 0.05, 0.10, 0.20),
)
sigma_x_full: Optional[float] = None
sigma_time: Optional[float] = knob_field(
help="Time-domain sigma_t override; default measures from the FID "
"active-region tail.",
cli=True,
argtype=float,
)
@dataclass
class PolishSubSettings:
"""Pure-exp polish step (consumed by ``calibrate_tau`` only)."""
polish: Optional[bool] = None
polish_n_iter: Optional[int] = knob_field(
help="Gauss-Newton polish iterations per eligible contributor.",
inst_sensitivity="N",
grid=(1, 2, 3),
)
polish_snr_cap: Optional[float] = knob_field(
help="SNR above which the Gauss-Newton polish is skipped (avoid "
"over-correction).",
tier="primary",
inst_sensitivity="Y",
grid=(5.0, 7.0, 9.0, 12.0),
)
polish_noise_debias: Optional[bool] = knob_field(
help="Apply Rician-unbiased magnitude on high-SNR frames (removes "
"residual bias).",
inst_sensitivity="Y",
grid=(False, True),
)
@dataclass
class AggregationSubSettings:
"""Majority-vote + acceptance pre-conditions (shared by both τ twins)."""
min_contributors: Optional[int] = knob_field(
help="Minimum contributor count for the calibration to pass " "preconditions.",
inst_sensitivity="N",
grid=(100, 200, 400, 800),
cli=True,
argtype=int,
)
sigma_tau_fraction_max: Optional[float] = knob_field(
help="Max sigma_tau/tau_maj for the calibration to pass preconditions.",
inst_sensitivity="N",
grid=(0.10, 0.20, 0.30),
cli=True,
argtype=float,
)
bimodality_dominant_fraction: Optional[float] = knob_field(
help="Dominant-mode fraction above which a bimodal histogram still " "passes.",
inst_sensitivity="N",
grid=(0.6, 0.7, 0.8),
cli=True,
argtype=float,
)
sigma_tau_floor_us: Optional[float] = knob_field(
help="Floor on the reported sigma_tau (guards against over-tight " "spreads).",
inst_sensitivity="maybe",
grid=(0.0, 0.5, 1.0),
)
spur_cluster_multiplier: Optional[float] = knob_field(
help="Scale on the spur-cluster width (wider -> more bins flagged as "
"spurs).",
inst_sensitivity="maybe",
grid=(1.0, 1.5, 2.0),
)
@dataclass
class BandSubSettings:
"""Per-band majority routing (shared by both τ twins)."""
compute_band_majorities: Optional[bool] = knob_field(
help="Compute per-band tau majorities (the tau-vs-frequency band " "steps).",
inst_sensitivity="Y",
grid=(False, True),
)
band_edges_mhz: Optional[Tuple[float, ...]] = None
band_labels: Optional[Tuple[str, ...]] = None
min_contributors_per_band: Optional[int] = knob_field(
help="Min contributors for a band to use its own tau majority (else "
"band-wide).",
tier="primary",
inst_sensitivity="Y",
grid=(25, 50, 100, 200),
)
@dataclass
class GaussianSubSettings:
"""Knobs consumed only by ``calibrate_tau(shape="gaussian")``.
``min_contributors`` is distinct from
:attr:`AggregationSubSettings.min_contributors` -- the Gaussian variant
has a smaller hard default (50 vs 200) because the eligible Gaussian
pool is naturally smaller after the Δχ²ᵣ filter. The two live on
different sub-blocks to avoid the name collision; the impl routes a
bare ``--min-contributors`` onto this block for a Gaussian run.
"""
snr_min: Optional[float] = knob_field(
help="Gaussian tau_G: per-bin SNR floor for a contributor to enter " "the fit.",
tier="primary",
inst_sensitivity="Y",
grid=(10.0, 15.0, 20.0, 30.0),
cli=True,
argtype=float,
)
tau_G_bound_lo: Optional[float] = knob_field(
help="Gaussian tau_G lower fit bound (us).",
inst_sensitivity="maybe",
grid=(0.2, 0.5, 1.0),
cli=True,
argtype=float,
flag="--tau-g-bound-lo",
)
tau_G_bound_hi: Optional[float] = knob_field(
help="Gaussian tau_G upper fit bound (us).",
inst_sensitivity="maybe",
grid=(50.0, 100.0, 200.0),
cli=True,
argtype=float,
flag="--tau-g-bound-hi",
)
tau_G_seeds: Optional[Tuple[float, ...]] = None
delta_chi2r_min: Optional[float] = knob_field(
help="Min chi2r improvement of the Gaussian over the exp fit to count "
"a bin.",
inst_sensitivity="maybe",
grid=(0.5, 1.0, 2.0),
cli=True,
argtype=float,
)
tau_G_upper_fraction: Optional[float] = knob_field(
help="Fraction of the tau_G bound above which a fit is treated as " "railed.",
inst_sensitivity="maybe",
grid=(0.5, 0.7, 0.9),
cli=True,
argtype=float,
flag="--tau-g-upper-fraction",
)
min_contributors: Optional[int] = knob_field(
help="Minimum Gaussian-eligible contributor count for tau_G " "preconditions.",
inst_sensitivity="maybe",
grid=(25, 50, 100),
)
@dataclass
class RecommendationSubSettings:
"""``recommend_shape``-only knobs plus the calibrate_tau auto-run flag.
Overlaps with :class:`GaussianSubSettings` in four fields
(``snr_min`` / ``tau_bound_lo`` / ``tau_bound_hi`` / ``tau_G_seeds``)
by design: the shape-recommendation hook and the production τ_G
calibration are conceptually independent and may legitimately ship
with different operating points (the recommender's contributor pool
can be wider or narrower than the calibration's).
``auto_recommend`` controls whether ``calibrate_tau`` (either shape)
invokes :func:`recommend_shape` automatically
after the primary calibration writes; default ``True`` so the
Stage 5 resolver's *recommended* layer fires on every fresh Stage 2b
run without a second explicit user step.
"""
snr_min: Optional[float] = knob_field(
help="Shape vote: per-bin SNR floor for a contributor to vote.",
inst_sensitivity="maybe",
grid=(10.0, 15.0, 20.0, 30.0),
)
tau_bound_lo: Optional[float] = knob_field(
help="Shape vote: lower tau fit bound shared by the per-bin model "
"fits (us).",
inst_sensitivity="maybe",
grid=(0.2, 0.5, 1.0),
)
tau_bound_hi: Optional[float] = knob_field(
help="Shape vote: upper tau fit bound shared by the per-bin model "
"fits (us).",
inst_sensitivity="maybe",
grid=(50.0, 100.0, 200.0),
)
tau_G_seeds: Optional[Tuple[float, ...]] = None
pure_margin_threshold: Optional[float] = knob_field(
help="Min SNR-weighted vote margin for a pure shape to win (else " "'none').",
tier="primary",
inst_sensitivity="maybe",
grid=(0.05, 0.10, 0.15, 0.20),
)
auto_recommend: Optional[bool] = None
[docs]
@dataclass
class TauCalibrationSettings:
"""Stage 2b τ calibration settings (see module docstring)."""
stft: StftSubSettings = field(default_factory=StftSubSettings)
polish: PolishSubSettings = field(default_factory=PolishSubSettings)
aggregation: AggregationSubSettings = field(default_factory=AggregationSubSettings)
band: BandSubSettings = field(default_factory=BandSubSettings)
gaussian: GaussianSubSettings = field(default_factory=GaussianSubSettings)
recommendation: RecommendationSubSettings = field(
default_factory=RecommendationSubSettings
)
[docs]
def is_empty(self) -> bool:
"""True if no field is set across any sub-dataclass."""
return not sf.any_field_set(self, _SUB_NAMES)
# Sub-dataclass field names on TauCalibrationSettings, in HDF5/YAML order.
_SUB_NAMES = (
"stft",
"polish",
"aggregation",
"band",
"gaussian",
"recommendation",
)
# Hard defaults per sub-dataclass. These mirror the ``DEFAULT_*`` constants
# in ``fitting/tau_calibration.py``. Kept as inline literals (rather than
# imported from ``fitting/``) to keep ``core`` dependency-free from
# ``fitting``; the fitting module's constants are the readable canonical
# source and these must track them.
_HARD_DEFAULTS: Dict[str, Dict[str, Any]] = {
"stft": {
"n_seg": 10,
"t_sigma": 5.0,
"tau_max_factor": 5.0, # tau_max = factor * T_full when tau_max_us unset
"rss_gate_factor": 5.0,
"relative_gate_fraction": 0.05,
# ``tau_max_us``, ``sigma_x_full``, and ``sigma_time`` legitimately
# stay None: tau_max_us derives at runtime from T_full * tau_max_factor;
# sigma_x_full is an opt-in spectral-noise override; sigma_time is
# measured from the FID tail when unset.
},
"polish": {
"polish": True,
"polish_n_iter": 1,
"polish_snr_cap": 9.0,
"polish_noise_debias": False,
},
"aggregation": {
"min_contributors": 200,
"sigma_tau_fraction_max": 0.20,
"bimodality_dominant_fraction": 0.70,
"sigma_tau_floor_us": 0.5,
"spur_cluster_multiplier": 1.0,
},
"band": {
"compute_band_majorities": True,
"min_contributors_per_band": 50,
# ``band_edges_mhz`` legitimately stays None (default arithmetic
# three-way split based on the Stage 1 trim range); ``band_labels``
# legitimately stays None (defaults to ("low", "mid", "high") inside
# ``compute_band_majorities``).
},
"gaussian": {
"snr_min": 20.0,
"tau_G_bound_lo": 0.5,
"tau_G_bound_hi": 100.0,
"tau_G_seeds": (100.0, 50.0, 20.0, 10.0, 5.0, 3.0),
"delta_chi2r_min": 1.0,
"tau_G_upper_fraction": 0.7,
"min_contributors": 50,
},
"recommendation": {
"snr_min": 20.0,
"tau_bound_lo": 0.5,
"tau_bound_hi": 100.0,
"tau_G_seeds": (100.0, 50.0, 20.0, 10.0, 5.0, 3.0),
"pure_margin_threshold": 0.10,
"auto_recommend": True,
},
}
# Fields that must round-trip as tuples (not lists / arrays). Mirror the
# typed declarations on the sub-dataclasses above.
_TUPLE_FIELDS = {"tau_G_seeds", "band_edges_mhz", "band_labels"}
def _coerce_tuple_element(field_name: str, value: Any) -> Any:
if field_name == "band_labels":
if isinstance(value, bytes):
return value.decode("utf-8")
return str(value)
return float(value)
# Value codecs: tuple-valued fields encode as lists (HDF5 1-D arrays) and
# decode back to plain tuples; everything else uses the framework defaults.
def _encode_value(field_name: str, value: Any) -> Any:
if value is None:
return sf.NONE
if isinstance(value, tuple):
return list(value)
return value
def _decode_value(field_name: str, value: Any) -> Any:
if isinstance(value, bytes):
value = value.decode("utf-8")
if isinstance(value, str) and value == sf.NONE:
return None
if field_name in _TUPLE_FIELDS and value is not None:
# HDF5 reads tuples back as numpy arrays; coerce to plain tuples.
return tuple(_coerce_tuple_element(field_name, v) for v in value)
return value
def _yaml_encode(field_name: str, value: Any) -> Any:
if isinstance(value, tuple):
return list(value)
return value
def _yaml_coerce(field_name: str, value: Any) -> Any:
if field_name in _TUPLE_FIELDS and value is not None:
return tuple(_coerce_tuple_element(field_name, v) for v in value)
return value
# ---------------------------------------------------------------------------
# Resolution chain
# ---------------------------------------------------------------------------
def resolve(
explicit: Optional[TauCalibrationSettings] = None,
preset: Optional[TauCalibrationSettings] = None,
persisted: Optional[TauCalibrationSettings] = None,
recommended: Optional[TauCalibrationSettings] = None,
) -> TauCalibrationSettings:
"""Merge the four layers by precedence into a resolved ``TauCalibrationSettings``.
Per-field precedence: ``explicit > persisted > preset > recommended``,
then any remaining ``None`` field falls back to the matching value in
:data:`_HARD_DEFAULTS`. A value persisted in the ``.ftmw`` outranks a
``.yml`` preset, so the preset only seeds fields the file has not fixed
and a shared experiment reproduces from the file alone. The
``recommended`` layer is reserved for a future upstream recommender; it
is currently always passed ``None`` by Stage 2b call sites, and the slot
is kept here so the resolver shape stays uniform with Stage 5's.
"""
layers = (explicit, persisted, preset, recommended)
return sf.fill_resolved_subblocks(
TauCalibrationSettings(),
TauCalibrationSettings,
_SUB_NAMES,
_HARD_DEFAULTS,
layers,
)
# ---------------------------------------------------------------------------
# Dict <-> dataclass round-trip (drives both HDF5 and YAML serialization)
# ---------------------------------------------------------------------------
def to_attrs(settings: TauCalibrationSettings) -> Dict[str, Any]:
"""Nested attrs dict (one top-level key per sub-dataclass).
Sub-dataclass values use ``__None__`` for unset fields; tuples are
encoded as lists so HDF5 can persist them as 1-D arrays.
"""
return sf.subblocks_to_attrs(settings, _SUB_NAMES, _encode_value)
def from_attrs(attrs: Dict[str, Any]) -> TauCalibrationSettings:
"""Inverse of :func:`to_attrs` (tolerant of missing sub-blocks)."""
return sf.subblocks_from_attrs(
TauCalibrationSettings(),
TauCalibrationSettings,
_SUB_NAMES,
attrs,
_decode_value,
)
# ---------------------------------------------------------------------------
# YAML interchange
# ---------------------------------------------------------------------------
def to_yaml_dict(settings: TauCalibrationSettings) -> Dict[str, Any]:
"""Sparse nested dict suitable for ``yaml.safe_dump`` (omits ``None``)."""
return sf.subblocks_to_yaml_dict(settings, _SUB_NAMES, _yaml_encode)
def from_yaml_dict(data: Optional[Mapping[str, Any]]) -> TauCalibrationSettings:
"""Build a :class:`TauCalibrationSettings` from a YAML-shaped mapping.
Each sub-dataclass block is a mapping of field name -> value; unknown
keys raise ``ValueError`` so typos surface loudly. Preset-metadata
keys ``name`` and ``description`` at the top level are accepted but
ignored.
"""
if data is None:
return TauCalibrationSettings()
if not isinstance(data, dict):
raise ValueError(f"preset YAML root must be a mapping; got {type(data)}")
return sf.subblocks_from_yaml_dict(
TauCalibrationSettings(),
TauCalibrationSettings,
_SUB_NAMES,
data,
_yaml_coerce,
)
def from_yaml(source: Union[str, Path]) -> TauCalibrationSettings:
"""Load a :class:`TauCalibrationSettings` from a YAML file path or text."""
return from_yaml_dict(sf.load_yaml_source(source))
def load_preset(name_or_path: Union[str, Path]) -> TauCalibrationSettings:
"""Load a Stage 2b preset by bare name or by filesystem path.
Bare names resolve against the packaged ``ftmwpipeline.presets``
resources (e.g. ``"defaults"`` ->
``ftmwpipeline/presets/defaults.yaml``); paths load directly.
Preset YAML wraps the Stage 2b settings inside a top-level ``stage2b:``
block (alongside an optional ``stage5:`` block for Stage 5 settings).
The ``name:`` and ``description:`` metadata fields are accepted but
ignored by the settings parser -- they're documentation for the preset
author.
Returns an empty :class:`TauCalibrationSettings` (no fields set) when
the preset carries no ``stage2b:`` block, so a Stage-5-only preset
loads cleanly without producing spurious Stage 2b overrides.
Parameters
----------
name_or_path :
Bare preset name (no extension) or a path to a YAML file.
Returns
-------
TauCalibrationSettings
The parsed preset; unset fields stay ``None`` so the resolver can
fall through to higher-precedence layers.
Raises
------
FileNotFoundError
If a bare name does not match any packaged preset, or the
supplied path does not exist.
"""
return sf.load_subblock_preset(
TauCalibrationSettings, "stage2b", name_or_path, from_yaml_dict
)
def to_yaml(settings: TauCalibrationSettings) -> str:
"""Serialize to a YAML string (sparse; omits unset fields)."""
return sf.dump_yaml(to_yaml_dict(settings))
__all__ = [
"StftSubSettings",
"PolishSubSettings",
"AggregationSubSettings",
"BandSubSettings",
"GaussianSubSettings",
"RecommendationSubSettings",
"TauCalibrationSettings",
"resolve",
"to_attrs",
"from_attrs",
"to_yaml",
"from_yaml",
"to_yaml_dict",
"from_yaml_dict",
"load_preset",
]