diff --git a/papers/cosmo_val/config/config.yaml b/papers/cosmo_val/config/config.yaml index 945f2f8f..98409371 100644 --- a/papers/cosmo_val/config/config.yaml +++ b/papers/cosmo_val/config/config.yaml @@ -68,7 +68,7 @@ cosmo_val: cosebis: min_sep_int: 0.9 max_sep_int: 300 - nbins_int: 2000 + nbins_int: 1000 npatch: 100 nmodes: 20 scale_cuts: [ diff --git a/src/sp_validation/cosmo_val/pure_eb.py b/src/sp_validation/cosmo_val/pure_eb.py index 7524074e..71f5d76a 100644 --- a/src/sp_validation/cosmo_val/pure_eb.py +++ b/src/sp_validation/cosmo_val/pure_eb.py @@ -27,7 +27,7 @@ def calculate_pure_eb( nbins=None, min_sep_int=0.08, max_sep_int=300, - nbins_int=100, + nbins_int=1000, npatch=256, var_method="jackknife", cov_path_int=None, @@ -57,7 +57,7 @@ def calculate_pure_eb( max_sep_int : float, optional Maximum separation for the integration binning. Defaults to 300. nbins_int : int, optional - Number of bins for the integration binning. Defaults to 100. + Number of bins for the integration binning. Defaults to 1000. npatch : int, optional Number of patches for the jackknife or bootstrap resampling. Defaults to the value in self.npatch if not provided. @@ -143,7 +143,7 @@ def plot_pure_eb( nbins=None, min_sep_int=0.08, max_sep_int=300, - nbins_int=100, + nbins_int=1000, npatch=None, var_method="jackknife", cov_path_int=None, @@ -175,7 +175,7 @@ def plot_pure_eb( Binning parameters for reporting scale. Uses treecorr_config if None. min_sep_int, max_sep_int, nbins_int : float, float, int Binning parameters for integration scale - (default: 0.08-300 arcmin, 100 bins) + (default: 0.08-300 arcmin, 1000 bins) npatch : int, optional Number of patches for jackknife covariance. Uses self.npatch if None. var_method : str diff --git a/src/sp_validation/sacc_io.py b/src/sp_validation/sacc_io.py new file mode 100644 index 00000000..a16079f9 --- /dev/null +++ b/src/sp_validation/sacc_io.py @@ -0,0 +1,989 @@ +"""SACC_IO. + +:Name: sacc_io.py + +:Description: Read/write the standard SACC data-product layout for the + weak-lensing validation package. One file describes each + catalogue version: + + - ``{version}.sacc`` — NZ tracers, reporting-grid ξ±, pseudo-Cℓ + (EE/BB/EB) with bandpower windows, COSEBIs, pure E/B, ρ/τ PSF + diagnostics, and the fine ξ± integration input for + COSEBIs / pure-EB (``grid='integration'`` tagged points). The + covariance is assembled block-diagonally from the + per-statistic covariances (zero cross-blocks, never + materialized): the analysis blocks first, then a dense + per-pair integration-ξ block (the analytic integration-binning + covariance when it exists — it feeds derived-statistic error + propagation — or the TreeCorr ``varxip``/``varxim`` diagonal as + degraded fallback). ``assemble_covariance`` hands sacc a list + of blocks, which stores a ``BlockDiagonalCovariance`` (one + FITS table per block, Σ block² on disk rather than a dense + N²) — cost scales with the integration grid's size, a + parameter set by the caller, not baked into the layout. + + Insertion order is load-bearing. A Sacc is a flat list of data + points in the order ``add_data_point`` was called, and row/column + ``i`` of the covariance refers to the ``i``-th inserted point — + there is no other linkage between a point and its covariance + entry. SACC preserves that order bitwise through FITS save/load, + so writers define the covariance layout by their insertion + sequence. Writers below insert in the canonical order — ξ+ then + ξ−, Cℓ (ee, bb, eb), COSEBIs (all Eₙ then all Bₙ), pure E/B + (xip_E, xim_E, xip_B, xim_B, xip_amb, xim_amb — matching + ``b_modes._EB_KEYS``), ρ, then τ — but readers never assume + global order: they resolve indices through + ``Sacc.indices(dtype, tracers, **tags)``. + + Tag filters are plain keyword arguments to ``indices`` / + ``get_data_points`` / ``get_tag``; the ``tags={...}`` form + silently selects nothing and must never be used. + + Tomographic ξ ordering: each ``add_xi`` call inserts one tracer + pair as ``[xip; xim]``, so a multi-pair vector is *pair-major* + — ``[pair_0 xip; pair_0 xim; pair_1 xip; …]`` — not type-major + (``[all xip; all xim]``). A tomographic ξ covariance with + cross-pair correlations is therefore supplied to + ``assemble_covariance`` as ONE contiguous block spanning the + consecutive ``add_xi`` calls, ordered pair-by-pair to match + insertion. Writers must call ``add_xi`` in the same pair order + the covariance was built in. Converters that need a type-major + layout (e.g. the DES 2pt-FITS convention) permute explicitly via + ``s.indices`` rather than assuming global order. + +Optionality: a file's contents are flexible about which components of + a statistic it actually has. ``add_pseudo_cl`` requires EE (the + bandpower window's reference series) but BB and EB are optional + — an analysis that never computed EB simply omits it. + ``add_cosebis`` requires Eₙ but Bₙ is optional. ``add_pure_eb`` + requires xip_E/xim_E but the B and ambiguous-mode blocks are + each optional, independently (B and amb are unrelated + computations). Everything else a writer takes — θ/ℓ grids, + tracer/bin identifiers, the value array for a component you ARE + adding — is structurally necessary and has no default; + supplying it partially would desynchronise the covariance + layout, so it is refused rather than degraded. Readers mirror + this split: a composite reader (``get_pseudo_cl``, + ``get_cosebis``, ``get_pure_eb``) returns ``None`` (or omits the + key) for a component the file doesn't carry, but a selection + naming that component explicitly (``s.indices``, ``_mean``, + ``extract``) still fails loud on no match — silence is reserved + for "this file doesn't have that optional piece", never for + "you asked for something specific and it isn't there". +""" + +import numpy as np +import sacc + +PSF_TRACER = "psf_stars" + +# Standard SACC data-type strings. +XI_PLUS = "galaxy_shear_xi_plus" +XI_MINUS = "galaxy_shear_xi_minus" +CL_EE = "galaxy_shear_cl_ee" +CL_BB = "galaxy_shear_cl_bb" +CL_EB = "galaxy_shear_cl_eb" +COSEBI_EE = "galaxy_shear_cosebi_ee" +COSEBI_BB = "galaxy_shear_cosebi_bb" + +# Custom data-type strings (all parse under sacc.parse_data_type_name). +PURE_TYPES = { + "xip_E": "galaxy_shear_xiPureE_plus", + "xim_E": "galaxy_shear_xiPureE_minus", + "xip_B": "galaxy_shear_xiPureB_plus", + "xim_B": "galaxy_shear_xiPureB_minus", + "xip_amb": "galaxy_shear_xiPureAmb_plus", + "xim_amb": "galaxy_shear_xiPureAmb_minus", +} +# PURE_TYPES key order is the insertion order of the six pure-EB blocks — +# matches b_modes._EB_KEYS, whose order is the [xip_E; xim_E; xip_B; xim_B; +# xip_amb; xim_amb] layout of the treecorr/MC pure-EB covariance +# (b_modes.calculate_eb_statistics, ~L392). +PURE_KEYS = tuple(PURE_TYPES) + +RHO_PLUS = "psf_rho{k}_xi_plus" +RHO_MINUS = "psf_rho{k}_xi_minus" +TAU_PLUS = "galaxyPsf_tau{k}_xi_plus" +TAU_MINUS = "galaxyPsf_tau{k}_xi_minus" + + +def source_name(i): + """SACC tracer name for source redshift bin ``i`` (0-based). + + Kept as the single definition of the ``source_{i}`` naming contract — + external consumers address tracers through it rather than the f-string. + """ + return f"source_{i}" + + +def new_sacc(nz, metadata=None): + """Create a Sacc with the survey's NZ (and PSF) tracers. + + Parameters + ---------- + nz : dict or sequence + Redshift distributions, one per source bin. Either a mapping + ``{i: (z, nz)}`` keyed by 0-based bin index, or a sequence of + ``(z, nz)`` array pairs (bin index = position). Tracers are named + ``source_{i}``. + metadata : dict, optional + Key/value pairs stored on ``s.metadata``. + + Returns + ------- + sacc.Sacc + Sacc holding the ``source_{i}`` NZ tracers and the ``psf_stars`` + Misc tracer (needed by ρ/τ diagnostics). + """ + items = nz.items() if isinstance(nz, dict) else enumerate(nz) + s = sacc.Sacc() + for i, (z, nz_i) in items: + s.add_tracer("NZ", source_name(i), np.asarray(z), np.asarray(nz_i)) + # The PSF star sample sits alongside the source bins because a Sacc has a + # single tracer namespace — every data point references tracers from one + # flat list. This is bookkeeping, not physics: psf_stars is a Misc tracer + # (no n(z)) that exists only so ρ/τ points have something to reference. + s.add_tracer("Misc", PSF_TRACER) + for key, value in (metadata or {}).items(): + s.metadata[key] = value + return s + + +def _pair(bins): + """Resolve a ``(i, j)`` bin pair to the ``(source_i, source_j)`` names. + + The pair is normalised to ``i <= j``: shear-shear statistics are symmetric + in the tracer pair, and SACC stores each pair under one ordering, so + ``(1, 0)`` must address the same points as ``(0, 1)``. Kept as a helper: + every ``add_*``/``get_*`` function relies on this normalisation. + """ + i, j = sorted(bins) + return (source_name(i), source_name(j)) + + +def _check_ascending(name, values): + """Require ``values`` to be strictly ascending; else raise ValueError. + + Insertion order is the covariance (and bandpower-window) order, and readers + return points in insertion order, so an out-of-order grid would silently + desynchronise a data vector from its covariance. Enforce monotonicity at + write time instead. + """ + values = np.asarray(values) + if not np.all(np.diff(values) > 0): + raise ValueError( + f"{name} must be strictly ascending (insertion order is the " + f"covariance order); got {values.tolist()}" + ) + + +def _add_theta_series(s, dtype, tracers, theta, values, **tags): + """Insert one theta-tagged series, one point per (theta, value) pair.""" + for th, value in zip(theta, values): + s.add_data_point(dtype, tracers, float(value), theta=float(th), **tags) + + +def add_xi( + s, + bins, + theta, + xip, + xim, + *, + grid, + theta_nom=None, + npairs=None, + weight=None, +): + """Add a real-space shear 2PCF (ξ+ then ξ−) for one tracer pair. + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + bins : tuple of int + Source bin pair ``(i, j)``. + theta : array_like + Angular separations (arcmin) — TreeCorr ``meanr``. + xip, xim : array_like + ξ+ and ξ− at ``theta``. + grid : {'reporting', 'integration'} + Stored as the ``grid`` tag on every point. The reporting (analysis) + ξ± and the fine COSEBIs/pure-EB integration input share the same + data type and tracer pair, so this tag is the only thing that + disambiguates them in the file. + theta_nom : array_like, optional + Nominal bin centres — TreeCorr ``rnom`` — stored as ``theta_nom``. + npairs, weight : array_like, optional + TreeCorr pair counts and weights, stored per point. + """ + _check_ascending("theta", theta) + tracers = _pair(bins) + optional = {"theta_nom": theta_nom, "npairs": npairs, "weight": weight} + extras = {key: arr for key, arr in optional.items() if arr is not None} + for dtype, xi in ((XI_PLUS, xip), (XI_MINUS, xim)): + for n, th in enumerate(theta): + tags = { + "theta": float(th), + "grid": grid, + **{key: float(arr[n]) for key, arr in extras.items()}, + } + s.add_data_point(dtype, tracers, float(xi[n]), **tags) + + +def add_pseudo_cl( + s, + bins, + ell_eff, + cl_ee, + cl_bb=None, + cl_eb=None, + *, + window_ells, + window_weights, + grid="reporting", +): + """Add pseudo-Cℓ EE (required) plus whichever of BB/EB were computed. + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + bins : tuple of int + Source bin pair ``(i, j)``. + ell_eff : array_like + Effective multipole of each bandpower. + cl_ee : array_like + EE bandpowers at ``ell_eff``. + cl_bb, cl_eb : array_like, optional + BB and/or EB bandpowers at ``ell_eff``. Each defaults to ``None`` and + is then simply not written — EB in particular is often not computed + at all. + window_ells : array_like + Multipoles spanned by the bandpower window matrix (shape ``(nell,)``). + window_weights : array_like + Window matrix ``W`` of shape ``(nell, nbp)`` — one column per + bandpower — from NaMaster ``get_bandpower_windows``. One + ``sacc.BandpowerWindow`` is built and shared across every component + written. + grid : str, optional + Stored as the ``grid`` tag on every point (default ``'reporting'``), + joining ``merge``'s ℓ-consistency group; variant ℓ binnings belong + under different tag values. + """ + _check_ascending("ell_eff", ell_eff) + tracers = _pair(bins) + window = sacc.BandpowerWindow(np.asarray(window_ells), np.asarray(window_weights)) + # add_ell_cl accepts no extra tags, so inline its per-point insertion + # (ell + shared window + window_ind column index) plus the grid tag. + components = [(CL_EE, cl_ee)] + components += [ + (dtype, cl) for dtype, cl in ((CL_BB, cl_bb), (CL_EB, cl_eb)) if cl is not None + ] + for dtype, cl in components: + for n, (ell, value) in enumerate(zip(ell_eff, cl)): + s.add_data_point( + dtype, + tracers, + float(value), + ell=float(ell), + window=window, + window_ind=n, + grid=grid, + ) + + +def add_cosebis(s, bins, En, scale_cut, Bn=None): + """Add COSEBIs Eₙ (required) and Bₙ (optional) for one scale cut. + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + bins : tuple of int + Source bin pair ``(i, j)``. + En : array_like + E-mode COSEBI amplitudes, one per logarithmic mode ``n`` (1-based). + scale_cut : tuple of float + ``(theta_min, theta_max)`` in arcmin, stored on every point as the + ``theta_min``/``theta_max`` tags; multiple cuts coexist in one file, + told apart by these tags. + Bn : array_like, optional + B-mode COSEBI amplitudes at the same ``n``. Defaults to ``None`` and + is then simply not written. The ``[En; Bn]`` layout, when both are + present, matches the COSEBI covariance. + """ + tracers = _pair(bins) + theta_min, theta_max = scale_cut + components = [(COSEBI_EE, En)] + if Bn is not None: + components.append((COSEBI_BB, Bn)) + for dtype, modes in components: + for n, value in enumerate(modes, start=1): + s.add_data_point( + dtype, + tracers, + float(value), + n=n, + theta_min=float(theta_min), + theta_max=float(theta_max), + ) + + +def add_pure_eb( + s, + bins, + theta, + xip_E, + xim_E, + xip_B=None, + xim_B=None, + xip_amb=None, + xim_amb=None, + *, + grid="reporting", +): + """Add pure E-mode (required) plus whichever of B/ambiguous were computed. + + Blocks are inserted in ``PURE_KEYS`` order (xip_E, xim_E, xip_B, xim_B, + xip_amb, xim_amb), matching ``b_modes._EB_KEYS`` and the pure-EB + covariance layout — whichever subset is present. + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + bins : tuple of int + Source bin pair ``(i, j)``. + theta : array_like + Angular separations (arcmin), shared by every block written. + xip_E, xim_E : array_like + The pure E-mode correlation functions at ``theta``. + xip_B, xim_B : array_like, optional + The pure B-mode correlation functions. Both default to ``None``; the + pair is written together or not at all — supply both or neither. + xip_amb, xim_amb : array_like, optional + The ambiguous-mode correlation functions. Both default to ``None``; + same both-or-neither rule as B. + grid : str, optional + Stored as the ``grid`` tag on every point (default ``'reporting'``), + joining ξ's theta-consistency group in ``merge``'s guard. + """ + _check_ascending("theta", theta) + tracers = _pair(bins) + pairs = { + "B": (xip_B, xim_B), + "amb": (xip_amb, xim_amb), + } + for label, (p, m) in pairs.items(): + if (p is None) != (m is None): + raise ValueError( + f"add_pure_eb: xip_{label} and xim_{label} must both be " + "given or both omitted" + ) + values = { + "xip_E": xip_E, + "xim_E": xim_E, + "xip_B": xip_B, + "xim_B": xim_B, + "xip_amb": xip_amb, + "xim_amb": xim_amb, + } + for key in PURE_KEYS: + arr = values[key] + if arr is not None: + _add_theta_series(s, PURE_TYPES[key], tracers, theta, arr, grid=grid) + + +def add_rho(s, k, theta, rho_p, rho_m, *, grid="reporting"): + """Add a ρ_k PSF statistic (ρ+ then ρ−) on the ``psf_stars`` tracer. + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + k : int + ρ index (0…5). + theta : array_like + Angular separations (arcmin). + rho_p, rho_m : array_like + ρ_k+ and ρ_k− at ``theta``. + grid : str, optional + Stored as the ``grid`` tag on every point (default ``'reporting'``), + joining ξ's theta-consistency group in ``merge``'s guard. + """ + _check_ascending("theta", theta) + tracers = (PSF_TRACER, PSF_TRACER) + _add_theta_series(s, RHO_PLUS.format(k=k), tracers, theta, rho_p, grid=grid) + _add_theta_series(s, RHO_MINUS.format(k=k), tracers, theta, rho_m, grid=grid) + + +def add_tau(s, bins, k, theta, tau_p, tau_m, *, grid="reporting"): + """Add a τ_k PSF-leakage statistic (τ+ then τ−). + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + bins : tuple of int + Source bin ``i`` and PSF; the τ tracers are ``(source_i, psf_stars)``. + Only ``bins[0]`` is used. + k : int + τ index (0, 2 or 5). + theta : array_like + Angular separations (arcmin). + tau_p, tau_m : array_like + τ_k+ and τ_k− at ``theta``. + grid : str, optional + Stored as the ``grid`` tag on every point (default ``'reporting'``), + joining ξ's theta-consistency group in ``merge``'s guard. + """ + _check_ascending("theta", theta) + tracers = (source_name(bins[0]), PSF_TRACER) + _add_theta_series(s, TAU_PLUS.format(k=k), tracers, theta, tau_p, grid=grid) + _add_theta_series(s, TAU_MINUS.format(k=k), tracers, theta, tau_m, grid=grid) + + +def assemble_covariance(s, blocks): + """Assemble a ``BlockDiagonalCovariance`` from per-statistic blocks. + + Each block is validated against the current insertion order: its indices + must be contiguous and ascending, the blocks must tile ``0…len(s.mean)`` + exactly (no gap, no overlap), and each block must be square with a size + matching its index span. Any violation raises ``ValueError`` naming the + mismatch. Cross-blocks are zero and implicit — never materialized — + because the blocks are passed to ``add_covariance`` as a list, which + ``sacc.BaseCovariance.make`` turns into a ``BlockDiagonalCovariance`` + (one FITS table per block, Σ block² on disk rather than a dense N² file). + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place via ``add_covariance``. + blocks : sequence + Ordered ``(selector, cov)`` pairs (or a mapping of the same). Each + ``selector`` is either an index array, or a ``(data_type, tracers)`` + / ``(data_type, tracers, tags)`` tuple resolved through + ``s.indices``; ``cov`` is the block's dense covariance. + + Returns + ------- + sacc.Sacc + ``s``, with the assembled ``BlockDiagonalCovariance`` attached. + """ + items = blocks.items() if isinstance(blocks, dict) else blocks + ntot = len(s.mean) + ordered_blocks = [] + cursor = 0 + for selector, cov in items: + idx = _resolve_indices(s, selector) + cov = np.asarray(cov) + if not np.array_equal(idx, np.arange(idx[0], idx[0] + len(idx))): + raise ValueError( + f"covariance block {selector!r} resolves to non-contiguous " + f"or non-ascending indices {idx.tolist()}" + ) + if idx[0] != cursor: + raise ValueError( + f"covariance block {selector!r} starts at index {idx[0]} but " + f"the previous blocks cover through {cursor} — blocks must tile " + "the data vector with no gap or overlap" + ) + if cov.ndim != 2 or cov.shape[0] != cov.shape[1]: + raise ValueError( + f"covariance block {selector!r} must be square; got shape {cov.shape}" + ) + if cov.shape[0] != len(idx): + raise ValueError( + f"covariance block {selector!r} has size {cov.shape[0]} but " + f"spans {len(idx)} data points" + ) + ordered_blocks.append(cov) + cursor = idx[-1] + 1 + if cursor != ntot: + raise ValueError( + f"covariance blocks cover {cursor} of {ntot} data points — the " + "blocks must tile the whole data vector" + ) + s.add_covariance(ordered_blocks) + return s + + +def _resolve_indices(s, selector): + """Resolve a covariance-block selector to a sorted index array.""" + if isinstance(selector, (np.ndarray, list, tuple, range)) and not ( + len(selector) in (2, 3) and isinstance(selector[0], str) + ): + return np.asarray(selector, dtype=int) + data_type, tracers = selector[0], selector[1] + tags = selector[2] if len(selector) == 3 else {} + return _indices(s, data_type, tuple(tracers), **tags) + + +def add_diagonal_covariance(s, variances): + """Attach a ``DiagonalCovariance`` from a 1-D variance array. + + The 1-D array is passed straight to ``add_covariance`` (never + ``np.diag``), which is what makes SACC store a ``DiagonalCovariance``. + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + variances : array_like + Per-point variances, ``len == len(s.mean)``. + + Returns + ------- + sacc.Sacc + ``s``, with the ``DiagonalCovariance`` attached. + """ + s.add_covariance(np.asarray(variances)) + return s + + +def get_nz(s, i): + """Return ``(z, nz)`` for source bin ``i``.""" + tracer = s.tracers[source_name(i)] + return tracer.z, tracer.nz + + +def _get_pm(s, dtype_p, dtype_m, tracers, **tags): + """Return ``(theta, plus, minus)`` for a +/− data-type pair.""" + return ( + _tag(s, dtype_p, tracers, "theta", **tags), + _mean(s, dtype_p, tracers, **tags), + _mean(s, dtype_m, tracers, **tags), + ) + + +def get_xi(s, bins, *, grid): + """Return ``(theta, xip, xim)`` for one tracer pair and grid.""" + return _get_pm(s, XI_PLUS, XI_MINUS, _pair(bins), grid=grid) + + +def get_pseudo_cl(s, bins): + """Return ``(ell_eff, cl_ee, cl_bb, cl_eb, window)`` for one tracer pair. + + ``cl_bb``/``cl_eb`` come back ``None`` if the file doesn't carry that + component (``add_pseudo_cl`` makes both optional). ``window`` is the + shared ``sacc.BandpowerWindow`` recovered via ``get_bandpower_windows``; + its columns are in the same insertion order as the returned + ``ell_eff``/``cl`` arrays, so window column ``j`` corresponds to + ``ell_eff[j]``. + """ + tracers = _pair(bins) + window = s.get_bandpower_windows(s.indices(CL_EE, tracers)) + return ( + _tag(s, CL_EE, tracers, "ell"), + _mean(s, CL_EE, tracers), + _mean_optional(s, CL_BB, tracers), + _mean_optional(s, CL_EB, tracers), + window, + ) + + +def get_cosebis(s, bins, scale_cut=None): + """Return ``(n, En, Bn)`` for one tracer pair. + + ``Bn`` comes back ``None`` if the file doesn't carry it (``add_cosebis`` + makes it optional). + + Parameters + ---------- + scale_cut : tuple of float, optional + ``(theta_min, theta_max)`` to select when several cuts share the file. + """ + tracers = _pair(bins) + if scale_cut is None: + cuts = { + (s.data[i].tags["theta_min"], s.data[i].tags["theta_max"]) + for i in _indices(s, COSEBI_EE, tracers) + } + if len(cuts) > 1: + raise ValueError( + f"several COSEBIs scale cuts share the file ({sorted(cuts)}) " + "— pass scale_cut=(theta_min, theta_max) to pick one" + ) + tags = dict(zip(("theta_min", "theta_max"), map(float, scale_cut or ()))) + modes = _tag(s, COSEBI_EE, tracers, "n", **tags) + return ( + modes.astype(int), + _mean(s, COSEBI_EE, tracers, **tags), + _mean_optional(s, COSEBI_BB, tracers, **tags), + ) + + +def get_pure_eb(s, bins): + """Return ``(theta, {key: array})`` for whichever pure-EB blocks exist. + + xip_E/xim_E are always present (``add_pure_eb`` requires them); the dict + holds whichever of the B and ambiguous-mode keys (out of ``PURE_KEYS``: + xip_E, xim_E, xip_B, xim_B, xip_amb, xim_amb) the file actually carries — + a key absent from the file is simply absent from the dict, not mapped to + ``None``. + """ + tracers = _pair(bins) + theta = _tag(s, PURE_TYPES["xip_E"], tracers, "theta") + arrays = {} + for key in PURE_KEYS: + values = _mean_optional(s, PURE_TYPES[key], tracers) + if values is not None: + arrays[key] = values + return theta, arrays + + +def get_rho(s, k): + """Return ``(theta, rho_p, rho_m)`` for ρ index ``k``.""" + tracers = (PSF_TRACER, PSF_TRACER) + return _get_pm(s, RHO_PLUS.format(k=k), RHO_MINUS.format(k=k), tracers) + + +def get_tau(s, bins, k): + """Return ``(theta, tau_p, tau_m)`` for τ index ``k`` and source bin.""" + tracers = (source_name(bins[0]), PSF_TRACER) + return _get_pm(s, TAU_PLUS.format(k=k), TAU_MINUS.format(k=k), tracers) + + +def _indices(s, data_type, tracers, **tag_filters): + """``Sacc.indices`` that fails loud instead of selecting nothing. + + ``Sacc.indices`` returns an *empty array* (warning only) when a selection + matches no point — e.g. a typo'd tag value, or a float tag filter that is + not bitwise-identical to the stored one. Every reader here funnels through + this guard so an unmatched selection raises instead of propagating empty + arrays downstream. + """ + idx = np.asarray(s.indices(data_type, tracers, **tag_filters), dtype=int) + if len(idx) == 0: + raise ValueError( + f"selection matched no points: ({data_type}, {tracers}, " + f"{tag_filters}) — note float tags match by exact equality" + ) + return idx + + +def _mean(s, data_type, tracers, **tag_filters): + """Mean values for a selection, in ``s.indices`` (insertion) order. + + Never re-sort: insertion order is the covariance and bandpower-window + order, so returning in ``s.indices`` order keeps every reader aligned with + the covariance for any file (and ascending for canonically-written files, + which the writers enforce). + """ + return s.mean[_indices(s, data_type, tracers, **tag_filters)] + + +def _tag(s, data_type, tracers, tag, **tag_filters): + """Values of ``tag`` for a selection, in insertion order.""" + idx = _indices(s, data_type, tracers, **tag_filters) + return np.array([s.data[i].tags[tag] for i in idx]) + + +def _mean_optional(s, data_type, tracers, **tag_filters): + """Mean values for a selection, or ``None`` if the file has none. + + Used by composite readers (``get_pseudo_cl``, ``get_cosebis``, + ``get_pure_eb``) for the components a writer made optional (BB/EB, + COSEBI Bₙ, pure B/ambiguous): absence is a legitimate "this file doesn't + have that piece", not a typo to fail loud on — unlike ``_mean``/ + ``_indices``, used for selections that name a component explicitly. + """ + idx = np.asarray(s.indices(data_type, tracers, **tag_filters), dtype=int) + return s.mean[idx] if len(idx) else None + + +def extract(s, data_type=None, tracers=None, **tag_filters): + """Extract a sub-Sacc (points + aligned covariance sub-block). + + A copy is made and everything *not* matching the selection is removed, so + the covariance sub-block comes out correctly aligned and the original is + untouched. + + Parameters + ---------- + s : sacc.Sacc + Source, left unmodified. + data_type : str, optional + Data type to keep. + tracers : tuple, optional + Tracer pair to keep, as SACC tracer **names** (e.g. + ``("source_0", "source_0")`` or ``("source_0", "psf_stars")``) — *not* + integer bin indices. This differs deliberately from the ``add_*`` / + ``get_*`` interface, whose ``bins`` argument takes integer pairs: + ``extract`` is the generic selection escape hatch, mirroring + ``Sacc.keep_selection`` and addressing non-source tracers uniformly. + **tag_filters + Tag filters (plain kwargs, e.g. ``grid='integration'``). + + Returns + ------- + sacc.Sacc + New Sacc holding only the selected points. + """ + sub = s.copy() + tracer_filter = {"tracers": tuple(tracers)} if tracers is not None else {} + sub.keep_selection(data_type, **tracer_filter, **tag_filters) + if len(sub.mean) == 0: + raise ValueError( + f"extract selected no points: ({data_type}, {tracers}, {tag_filters})" + ) + return sub + + +def merge(saccs): + """Merge several per-statistic Sacc objects into one file's worth. + + A thin wrapper around ``sacc.concatenate_data_sets``: data points + concatenate in input order, tracers shared by several inputs (the + ``source_i`` NZ tracers, ``psf_stars``) are stored once, and the + covariance combines block-diagonally in the same order — the library + requires either **all** inputs to carry a covariance or **none**, and + raises otherwise (cross-statistic covariance assembly beyond + block-diagonal is out of scope here; see ``assemble_covariance``). + + Metadata must be consistent: keys present in several inputs must carry + equal values (a ``type: data`` file cannot merge with a ``type: mock`` + file), and the union lands on the result. This deliberately replaces the + library's clash behaviour, which mangles clashing keys by appending + labels. + + Grid consistency follows tagging semantics: the ``grid`` tag declares + which binning a set of points lives on, so all same-length theta (or + ell) arrays under one tag value must be bitwise identical — sacc itself + never validates angles across data types/tracers, and a grid that + differs only at floating-point level chokes CosmoSIS downstream instead + of failing loud here. Different lengths within a tag group pass + (scale-cut subsets are legitimate); grids under different tag values + are unconstrained (``reporting`` vs ``integration`` differ by design); + θ and ℓ are separate domains, each checked against itself only. ℓ + series sharing a bitwise-equal grid must also share the bandpower + window (series without windows skip that check). + + Parameters + ---------- + saccs : sequence of sacc.Sacc + The per-statistic data sets, in the insertion order the merged file + should have. Inputs are left unmodified. + + Returns + ------- + sacc.Sacc + The merged data set. + + Raises + ------ + ValueError + If metadata conflicts, a shared tracer differs across inputs, two + same-length theta or ell arrays under the same ``grid`` tag value + are not bitwise identical, or two ℓ series sharing a grid carry + different bandpower windows. + """ + saccs = list(saccs) + metadata = {} + for s in saccs: + for key, value in s.metadata.items(): + if key in metadata and metadata[key] != value: + raise ValueError( + f"conflicting metadata across merge inputs: {key!r} is " + f"{metadata[key]!r} in one input and {value!r} in another" + ) + metadata[key] = value + # Strip metadata before concatenating (the library "resolves" clashing + # keys by renaming them), then restore the validated union. + stripped = [s.copy() for s in saccs] + for s in stripped: + s.metadata.clear() + seen, shared = set(), set() # tracers appearing in more than one input + for s in saccs: + shared |= seen & set(s.tracers) + seen |= set(s.tracers) + same_tracers = sorted(shared) + # The library keeps the FIRST input's tracer on a name clash with no + # equality check — verify shared tracers really are the same object. + for name in same_tracers: + first, *rest = [s.tracers[name] for s in saccs if name in s.tracers] + for other in rest: + if type(other) is not type(first) or not all( + np.array_equal(getattr(first, a, None), getattr(other, a, None)) + for a in ("z", "nz") + ): + raise ValueError( + f"shared tracer {name!r} differs across merge inputs — " + "the merged file would silently keep the first" + ) + merged = sacc.concatenate_data_sets(*stripped, same_tracers=same_tracers) + for key, value in metadata.items(): + merged.metadata[key] = value + _check_grid_consistency(merged, "theta") + _check_grid_consistency(merged, "ell") + return merged + + +def _grid_groups(s, angle): + """Nested map ``grid-tag-value -> (data_type, tracers) -> point indices``. + + One entry per ``(data_type, tracers)`` series carrying an ``angle`` + (``'theta'`` or ``'ell'``) tag, in each series' own insertion order + (never re-sorted), nested under the ``grid`` tag value it lives on + (``None`` for untagged series) — the shape ``merge``'s consistency + guard checks within each tag value. Indices (not angle values) are + kept so the ℓ check can also recover each series' bandpower window. + """ + groups = {} + for i, point in enumerate(s.data): + if angle in point.tags: + groups.setdefault(point.tags.get("grid"), {}).setdefault( + (point.data_type, point.tracers), [] + ).append(i) + return groups + + +def _check_grid_consistency(s, angle): + """Raise unless same-length ``angle`` arrays under one ``grid`` tag match. + + Consistency follows tagging semantics: the ``grid`` tag declares which + binning a series lives on, so all same-length theta (or ell) arrays + sharing a tag value must be bitwise identical — sacc never validates + angles across data types/tracers, and a grid diverging at + floating-point level chokes CosmoSIS downstream instead of failing + loud here. Different lengths within a tag value pass (scale-cut + subsets are legitimate); series under different tag values are + unconstrained (``reporting`` vs ``integration`` differ by design). + θ and ℓ are separate domains, each checked against itself only. For + ℓ, two series on a bitwise-equal grid must also share the bandpower + window (equal window ells and weight matrix); series without windows + (foreign files) skip the window check. + """ + for tag, by_series in _grid_groups(s, angle).items(): + series = [ + (key, np.array([s.data[i].tags[angle] for i in idx]), idx) + for key, idx in by_series.items() + ] + for i, (key_a, arr_a, idx_a) in enumerate(series): + for key_b, arr_b, idx_b in series[i + 1 :]: + if len(arr_a) != len(arr_b): + continue + if not np.array_equal(arr_a, arr_b): + max_diff = np.max(np.abs(arr_a - arr_b)) + raise ValueError( + f"{angle} grids under the same grid tag ({tag!r}) " + f"differ; harmonize upstream — groups {key_a!r} and " + f"{key_b!r} (max abs diff {max_diff:.3e})" + ) + if angle != "ell" or any( + "window" not in s.data[idx[0]].tags for idx in (idx_a, idx_b) + ): + continue + win_a, win_b = map(s.get_bandpower_windows, (idx_a, idx_b)) + if not ( + np.array_equal(win_a.values, win_b.values) + and np.array_equal(win_a.weight, win_b.weight) + ): + raise ValueError( + f"bandpower windows differ between series sharing an " + f"ell grid under grid tag {tag!r}; harmonize upstream " + f"— groups {key_a!r} and {key_b!r}" + ) + + +def update_statistic(s, sub): + """Overwrite the values of ``s``'s points that match ``sub``'s, in place. + + The merge-back half of the extract → conceal → merge blinding flow + (PRD #241 §4): each point of ``sub`` is matched to exactly one point of + ``s`` by ``(data_type, tracers, tags)``, and that point's *value* is + replaced. Nothing else changes — insertion order, tags, windows and the + covariance are untouched (blinding shifts the mean only), so ``sub``'s + own covariance (e.g. the sub-block ``extract`` attaches) is deliberately + not consulted. A ``sub`` point with no match, or with several, raises + ``ValueError``. + + Parameters + ---------- + s : sacc.Sacc + Target, mutated in place. + sub : sacc.Sacc + The replacement block, e.g. ``extract(s, ...)`` after concealment. + """ + claimed = set() + for point in sub.data: + idx = s.indices(point.data_type, point.tracers, **point.tags) + if len(idx) != 1: + raise ValueError( + f"update_statistic: {len(idx)} points in the target match " + f"({point.data_type}, {point.tracers}, {point.tags}) — need " + "exactly one" + ) + if idx[0] in claimed: + raise ValueError( + f"update_statistic: two sub points match the same target " + f"point ({point.data_type}, {point.tracers}, {point.tags})" + ) + claimed.add(idx[0]) + s.data[idx[0]].value = point.value + + +def save(s, path, *, type): + """Write ``s`` to ``path`` (FITS), overwriting any existing file. + + Parameters + ---------- + s : sacc.Sacc + Data set to write; its metadata is stamped in place. + type : {'data', 'mock'} + Provenance of the underlying catalogue, stored as the required + ``type`` metadata tag (PRD #241 §4, "Mocks vs data"). The caller — + the pipeline computing the data vector — knows whether its input + catalogue is a mock; there is deliberately no default. ``load`` + refuses ``type='data'`` files that are not blinded. + """ + if type not in ("data", "mock"): + raise ValueError(f"type must be 'data' or 'mock'; got {type!r}") + if s.metadata.get("type", type) != type: + raise ValueError( + f"Sacc metadata already carries type={s.metadata['type']!r}; " + f"refusing to re-stamp as {type!r}" + ) + s.metadata["type"] = type + s.save_fits(path, overwrite=True) + + +def load(path, *, allow_unblinded=False): + """Load a Sacc from ``path`` (FITS), failing closed on unblinded data. + + Every sacc_io file carries a ``type: data|mock`` metadata tag (stamped by + ``save``); blinded files are additionally stamped ``concealed=True`` by + Smokescreen. A ``type='data'`` file without that stamp is real, unblinded + data, and loading it raises — skipping the blind can never silently + expose the measured vector (PRD #241 §4). Mocks load freely, blinded or + not. + + Parameters + ---------- + path : str + File to load. + allow_unblinded : bool, optional + Escape hatch for the two legitimate consumers of unblinded data: + the blinding step itself (which must read the true vector to conceal + it) and the unblinding/verification tooling. Nothing else — no + analysis, plotting or inference code — may pass ``True``. + + Returns + ------- + sacc.Sacc + The loaded data set. + """ + s = sacc.Sacc.load_fits(path) + if ( + s.metadata["type"] == "data" + and not s.metadata.get("concealed", False) + and not allow_unblinded + ): + raise ValueError( + f"{path} holds real data (type='data') without the " + "concealed=True blinding stamp — refusing to load an unblinded " + "data vector. Only the blinding/unblinding tooling may pass " + "allow_unblinded=True." + ) + return s diff --git a/src/sp_validation/tests/test_sacc_io.py b/src/sp_validation/tests/test_sacc_io.py new file mode 100644 index 00000000..8f8564ae --- /dev/null +++ b/src/sp_validation/tests/test_sacc_io.py @@ -0,0 +1,1184 @@ +"""Tests for :mod:`sp_validation.sacc_io`. + +All synthetic, all fast: build in memory, round-trip through ``tmp_path``, +and assert arrays/tags/windows/covariances come back bitwise-identical and +correctly aligned. No cluster paths. +""" + +import numpy as np +import pytest + +from sp_validation import sacc_io as sio + + +# --------------------------------------------------------------------------- # +# Synthetic builders +# --------------------------------------------------------------------------- # +def _nz(seed, n=50): + rng = np.random.default_rng(seed) + z = np.linspace(0.01, 2.0, n) + return z, rng.uniform(0.1, 1.0, n) + + +def _theta(nbins=6): + return np.geomspace(1.0, 100.0, nbins) + + +def _spd(n, seed): + """Symmetric positive-definite matrix of size ``n``.""" + a = np.random.default_rng(seed).normal(size=(n, n)) + return a @ a.T + n * np.eye(n) + + +def _base_sacc(nbins=1): + """A Sacc with ``nbins`` NZ source tracers plus the PSF tracer.""" + return sio.new_sacc({i: _nz(i) for i in range(nbins)}) + + +def _add_xi(s, bins=(0, 0), scale=1.0, grid="reporting"): + """Write the default synthetic ξ± block; returns ``(xip, xim)``.""" + xip, xim = np.arange(6) * scale * 1e-5, np.arange(6) * scale * 2e-5 + sio.add_xi(s, bins, _theta(), xip, xim, grid=grid) + return xip, xim + + +# Canonical per-statistic index blocks (concatenated in insertion order). +def _xi_block(s, tr, **tags): + return np.concatenate( + [s.indices(sio.XI_PLUS, tr, **tags), s.indices(sio.XI_MINUS, tr, **tags)] + ) + + +def _cl_block(s, tr): + return np.concatenate( + [s.indices(sio.CL_EE, tr), s.indices(sio.CL_BB, tr), s.indices(sio.CL_EB, tr)] + ) + + +def _cosebi_block(s, tr): + return np.concatenate([s.indices(sio.COSEBI_EE, tr), s.indices(sio.COSEBI_BB, tr)]) + + +# --------------------------------------------------------------------------- # +# 1. Per-writer round-trip (arrays / tags / windows / NZ bitwise) +# --------------------------------------------------------------------------- # +def _roundtrip(s, tmp_path, name="rt"): + path = tmp_path / f"{name}.sacc" + sio.save(s, str(path), type="mock") + return sio.load(str(path)) + + +def test_nz_roundtrip(tmp_path): + z, nz = _nz(3) + s = sio.new_sacc({0: (z, nz)}, metadata={"version": "v1.4.6.3"}) + s2 = _roundtrip(s, tmp_path, "nz") + z2, nz2 = sio.get_nz(s2, 0) + assert np.array_equal(z2, z) + assert np.array_equal(nz2, nz) + assert s2.metadata["version"] == "v1.4.6.3" + assert sio.PSF_TRACER in s2.tracers + + +def test_xi_roundtrip(tmp_path): + theta = _theta() + xip, xim = np.arange(6) * 1e-5, np.arange(6) * 2e-5 + npairs, weight = np.arange(6) * 1e3, np.arange(6) * 1.5 + s = _base_sacc() + sio.add_xi( + s, + (0, 0), + theta, + xip, + xim, + grid="reporting", + theta_nom=theta * 1.01, + npairs=npairs, + weight=weight, + ) + s2 = _roundtrip(s, tmp_path, "xi") + th, p, m = sio.get_xi(s2, (0, 0), grid="reporting") + assert np.array_equal(th, theta) + assert np.array_equal(p, xip) + assert np.array_equal(m, xim) + # extra tags survive + idx = s2.indices(sio.XI_PLUS, ("source_0", "source_0"), grid="reporting") + tags = s2.data[idx[0]].tags + assert tags["grid"] == "reporting" + assert set(tags) >= {"theta", "theta_nom", "npairs", "weight", "grid"} + + +def test_pseudo_cl_roundtrip(tmp_path): + ell_eff = np.array([30.0, 120.0, 210.0, 300.0]) + nell, nbp = 50, len(ell_eff) + window_ells = np.arange(2, 2 + nell).astype(float) + W = np.random.default_rng(5).uniform(size=(nell, nbp)) + ee, bb, eb = np.arange(nbp) * 1e-9, np.arange(nbp) * 2e-9, np.arange(nbp) * 3e-9 + s = _base_sacc() + sio.add_pseudo_cl( + s, + (0, 0), + ell_eff, + ee, + bb, + eb, + window_ells=window_ells, + window_weights=W, + ) + s2 = _roundtrip(s, tmp_path, "cl") + ell, cl_ee, cl_bb, cl_eb, window = sio.get_pseudo_cl(s2, (0, 0)) + assert np.array_equal(ell, ell_eff) + assert np.array_equal(cl_ee, ee) + assert np.array_equal(cl_bb, bb) + assert np.array_equal(cl_eb, eb) + assert np.array_equal(window.weight, W) + assert np.array_equal(window.values, window_ells) + + +def test_cosebis_roundtrip(tmp_path): + En, Bn = np.arange(1, 11) * 1e-6, np.arange(1, 11) * 1e-7 + s = _base_sacc() + sio.add_cosebis(s, (0, 0), En, (1.0, 100.0), Bn=Bn) + s2 = _roundtrip(s, tmp_path, "cosebi") + n, E, B = sio.get_cosebis(s2, (0, 0)) + assert np.array_equal(n, np.arange(1, 11)) + assert np.array_equal(E, En) + assert np.array_equal(B, Bn) + idx = s2.indices(sio.COSEBI_EE, ("source_0", "source_0")) + assert s2.data[idx[0]].tags["theta_min"] == 1.0 + assert s2.data[idx[0]].tags["theta_max"] == 100.0 + + +def test_pure_eb_roundtrip(tmp_path): + theta = _theta() + arrays = {key: np.arange(6) * (i + 1) * 1e-6 for i, key in enumerate(sio.PURE_KEYS)} + s = _base_sacc() + sio.add_pure_eb(s, (0, 0), theta, **arrays) + s2 = _roundtrip(s, tmp_path, "pureeb") + th, back = sio.get_pure_eb(s2, (0, 0)) + assert np.array_equal(th, theta) + for key in sio.PURE_KEYS: + assert np.array_equal(back[key], arrays[key]) + + +def test_rho_roundtrip(tmp_path): + theta = _theta() + s = _base_sacc() + for k in range(6): + sio.add_rho( + s, k, theta, np.arange(6) * (k + 1) * 1e-6, np.arange(6) * (k + 1) * 2e-6 + ) + s2 = _roundtrip(s, tmp_path, "rho") + for k in range(6): + th, p, m = sio.get_rho(s2, k) + assert np.array_equal(th, theta) + assert np.array_equal(p, np.arange(6) * (k + 1) * 1e-6) + assert np.array_equal(m, np.arange(6) * (k + 1) * 2e-6) + + +def test_tau_roundtrip(tmp_path): + theta = _theta() + s = _base_sacc() + for k in (0, 2, 5): + sio.add_tau( + s, + (0, 0), + k, + theta, + np.arange(6) * (k + 1) * 1e-6, + np.arange(6) * (k + 1) * 2e-6, + ) + s2 = _roundtrip(s, tmp_path, "tau") + for k in (0, 2, 5): + th, p, m = sio.get_tau(s2, (0, 0), k) + assert np.array_equal(th, theta) + assert np.array_equal(p, np.arange(6) * (k + 1) * 1e-6) + assert np.array_equal(m, np.arange(6) * (k + 1) * 2e-6) + assert s2.data[s2.indices(sio.TAU_PLUS.format(k=k))[0]].tracers == ( + "source_0", + sio.PSF_TRACER, + ) + + +# --------------------------------------------------------------------------- # +# 2. Covariance alignment +# --------------------------------------------------------------------------- # +def _multi_statistic_sacc(): + """Sacc with ξ+/ξ−, Cℓ (ee/bb/eb) and COSEBIs, ready for a covariance.""" + s = _base_sacc() + _add_xi(s) + ell = np.array([30.0, 120.0, 210.0]) + W = np.random.default_rng(0).uniform(size=(20, 3)) + sio.add_pseudo_cl( + s, + (0, 0), + ell, + np.arange(3) * 1e-9, + np.arange(3) * 2e-9, + np.arange(3) * 3e-9, + window_ells=np.arange(2, 22).astype(float), + window_weights=W, + ) + sio.add_cosebis( + s, (0, 0), np.arange(1, 6) * 1e-6, (1.0, 100.0), Bn=np.arange(1, 6) * 1e-7 + ) + return s + + +def test_assemble_covariance_alignment(tmp_path): + s = _multi_statistic_sacc() + tr = ("source_0", "source_0") + # Block selectors in canonical (insertion) order. + xi, cl, co = _xi_block(s, tr), _cl_block(s, tr), _cosebi_block(s, tr) + cov_xi, cov_cl, cov_co = _spd(len(xi), 1), _spd(len(cl), 2), _spd(len(co), 3) + sio.assemble_covariance(s, [(xi, cov_xi), (cl, cov_cl), (co, cov_co)]) + s2 = _roundtrip(s, tmp_path, "cov") + assert type(s2.covariance).__name__ == "BlockDiagonalCovariance" + dense = s2.covariance.dense + # each block's sub-covariance is exactly what went in + assert np.array_equal(dense[np.ix_(xi, xi)], cov_xi) + assert np.array_equal(dense[np.ix_(cl, cl)], cov_cl) + assert np.array_equal(dense[np.ix_(co, co)], cov_co) + # zero cross-blocks + assert np.array_equal(dense[np.ix_(xi, cl)], np.zeros((len(xi), len(cl)))) + assert np.array_equal(dense[np.ix_(xi, co)], np.zeros((len(xi), len(co)))) + assert np.array_equal(dense[np.ix_(cl, co)], np.zeros((len(cl), len(co)))) + + +def test_assemble_covariance_selector_tuples(): + """Blocks addressed by (data_type, tracers) tuples, not raw indices.""" + s = _multi_statistic_sacc() + tr = ("source_0", "source_0") + xi, cl = _xi_block(s, tr), _cl_block(s, tr) + sio.assemble_covariance( + s, + [ + (xi, _spd(len(xi), 1)), + (cl, _spd(len(cl), 2)), + ((sio.COSEBI_EE, tr), _spd(len(s.indices(sio.COSEBI_EE, tr)), 3)), + ((sio.COSEBI_BB, tr), _spd(len(s.indices(sio.COSEBI_BB, tr)), 4)), + ], + ) + assert type(s.covariance).__name__ == "BlockDiagonalCovariance" + assert s.covariance.dense.shape == (len(s.mean), len(s.mean)) + + +# --------------------------------------------------------------------------- # +# 3. assemble_covariance failure modes +# --------------------------------------------------------------------------- # +def test_assemble_covariance_wrong_dimension(): + s = _base_sacc() + _add_xi(s) + idx = np.arange(len(s.mean)) + with pytest.raises(ValueError, match="span"): + sio.assemble_covariance(s, [(idx, _spd(len(idx) - 1, 1))]) + + +def test_assemble_covariance_non_contiguous(): + s = _base_sacc() + _add_xi(s) + idx = np.array([0, 2, 4, 6, 8, 10, 1, 3]) # not contiguous/ascending + with pytest.raises(ValueError, match="non-contiguous"): + sio.assemble_covariance(s, [(idx, _spd(len(idx), 1))]) + + +def test_assemble_covariance_missing_coverage(): + s = _multi_statistic_sacc() + tr = ("source_0", "source_0") + xi = _xi_block(s, tr) + with pytest.raises(ValueError, match="tile"): + sio.assemble_covariance( + s, [(xi, _spd(len(xi), 1))] + ) # leaves cl+cosebi uncovered + + +def test_assemble_covariance_non_square(): + s = _base_sacc() + _add_xi(s) + idx = np.arange(len(s.mean)) + with pytest.raises(ValueError, match="square"): + sio.assemble_covariance(s, [(idx, np.ones((len(idx), len(idx) - 1)))]) + + +def test_assemble_covariance_overlap(): + s = _multi_statistic_sacc() + tr = ("source_0", "source_0") + xi = _xi_block(s, tr) + # second block starts before the first ended -> gap/overlap error + with pytest.raises(ValueError, match="tile|gap|overlap"): + sio.assemble_covariance(s, [(xi, _spd(len(xi), 1)), (xi, _spd(len(xi), 2))]) + + +# --------------------------------------------------------------------------- # +# 4. Fine file: DiagonalCovariance from 1-D variances +# --------------------------------------------------------------------------- # +def test_diagonal_covariance_roundtrip(tmp_path): + theta = np.geomspace(0.1, 250.0, 40) + n = len(theta) + xip, xim = np.arange(n) * 1e-5, np.arange(n) * 2e-5 + varxip, varxim = np.arange(1, n + 1) * 1e-12, np.arange(1, n + 1) * 2e-12 + s = _base_sacc() + sio.add_xi(s, (0, 0), theta, xip, xim, grid="integration") + variances = np.concatenate([varxip, varxim]) # [xip; xim] order + sio.add_diagonal_covariance(s, variances) + assert type(s.covariance).__name__ == "DiagonalCovariance" + s2 = _roundtrip(s, tmp_path, "integration") + assert type(s2.covariance).__name__ == "DiagonalCovariance" + assert np.array_equal(np.diag(s2.covariance.dense), variances) + + +# --------------------------------------------------------------------------- # +# 5. extract(): sub-covariance alignment; original untouched; tag filter +# --------------------------------------------------------------------------- # +def test_extract_subblock_and_original_untouched(): + s = _multi_statistic_sacc() + tr = ("source_0", "source_0") + xi, cl, co = _xi_block(s, tr), _cl_block(s, tr), _cosebi_block(s, tr) + cov_co = _spd(len(co), 3) + sio.assemble_covariance( + s, [(xi, _spd(len(xi), 1)), (cl, _spd(len(cl), 2)), (co, cov_co)] + ) + n_before = len(s.mean) + sub = sio.extract(s, data_type=sio.COSEBI_EE, tracers=tr) + # original untouched + assert len(s.mean) == n_before + # subset covariance equals the COSEBI-EE diagonal sub-block + ee_local = np.arange(len(s.indices(sio.COSEBI_EE, tr))) + assert np.array_equal(sub.covariance.dense, cov_co[np.ix_(ee_local, ee_local)]) + + +def test_extract_tag_filter(): + theta = _theta() + s = _base_sacc() + _add_xi(s) + _add_xi(s, scale=3.0, grid="integration") + sub = sio.extract( + s, data_type=sio.XI_PLUS, tracers=("source_0", "source_0"), grid="integration" + ) + assert len(sub.mean) == len(theta) + assert set(sub.get_tag("grid", sio.XI_PLUS)) == {"integration"} + + +# --------------------------------------------------------------------------- # +# 6. Tomographic case: >=2 bins, >=3 pairs, per-pair selection +# --------------------------------------------------------------------------- # +def test_tomographic_per_pair_selection(tmp_path): + theta = _theta() + s = _base_sacc(nbins=2) + pairs = [(0, 0), (0, 1), (1, 1)] + for k, (i, j) in enumerate(pairs): + sio.add_xi( + s, + (i, j), + theta, + np.arange(6) * (k + 1) * 1e-5, + np.arange(6) * (k + 1) * 2e-5, + grid="reporting", + ) + s2 = _roundtrip(s, tmp_path, "tomo") + for k, (i, j) in enumerate(pairs): + th, p, m = sio.get_xi(s2, (i, j), grid="reporting") + assert np.array_equal(th, theta) + assert np.array_equal(p, np.arange(6) * (k + 1) * 1e-5) + assert np.array_equal(m, np.arange(6) * (k + 1) * 2e-5) + # selecting one pair does not bleed into another + assert len(s2.indices(sio.XI_PLUS, ("source_0", "source_1"))) == len(theta) + assert len(s2.indices(sio.XI_PLUS, ("source_1", "source_1"))) == len(theta) + + +# --------------------------------------------------------------------------- # +# 7. Readers mirror writers on a mixed file +# --------------------------------------------------------------------------- # +def test_readers_on_mixed_file(tmp_path): + theta = _theta() + # ξ+Cℓ+COSEBIs from the shared builder; the assertions below restate its + # values (ell, W match _multi_statistic_sacc). + s = _multi_statistic_sacc() + ell = np.array([30.0, 120.0, 210.0]) + W = np.random.default_rng(0).uniform(size=(20, 3)) + pure = {key: np.arange(6) * (i + 1) * 1e-6 for i, key in enumerate(sio.PURE_KEYS)} + sio.add_pure_eb(s, (0, 0), theta, **pure) + for k in range(6): + sio.add_rho( + s, k, theta, np.arange(6) * (k + 1) * 1e-7, np.arange(6) * (k + 1) * 2e-7 + ) + for k in (0, 2, 5): + sio.add_tau( + s, + (0, 0), + k, + theta, + np.arange(6) * (k + 1) * 1e-8, + np.arange(6) * (k + 1) * 2e-8, + ) + s2 = _roundtrip(s, tmp_path, "mixed") + + _, p, m = sio.get_xi(s2, (0, 0), grid="reporting") + assert np.array_equal(p, np.arange(6) * 1e-5) and np.array_equal( + m, np.arange(6) * 2e-5 + ) + ell_r, ee, bb, eb, win = sio.get_pseudo_cl(s2, (0, 0)) + assert np.array_equal(ell_r, ell) and np.array_equal(win.weight, W) + n, E, B = sio.get_cosebis(s2, (0, 0)) + assert np.array_equal(E, np.arange(1, 6) * 1e-6) + _, back = sio.get_pure_eb(s2, (0, 0)) + for key in sio.PURE_KEYS: + assert np.array_equal(back[key], pure[key]) + for k in range(6): + _, rp, rm = sio.get_rho(s2, k) + assert np.array_equal(rp, np.arange(6) * (k + 1) * 1e-7) + for k in (0, 2, 5): + _, tp, tm = sio.get_tau(s2, (0, 0), k) + assert np.array_equal(tp, np.arange(6) * (k + 1) * 1e-8) + + +# --------------------------------------------------------------------------- # +# 8. End-to-end one-file layout for a synthetic catalogue version +# --------------------------------------------------------------------------- # +def test_end_to_end_one_file_layout(tmp_path): + version = "vSYNTH" + theta_c = _theta(20) + theta_f = np.geomspace(0.1, 250.0, 200) + + # one file: analysis products first, fine-grid integration input last + s = _base_sacc() + sio.add_xi( + s, (0, 0), theta_c, np.arange(20) * 1e-5, np.arange(20) * 2e-5, grid="reporting" + ) + sio.add_cosebis( + s, (0, 0), np.arange(1, 11) * 1e-6, (1.0, 100.0), Bn=np.arange(1, 11) * 1e-7 + ) + sio.add_xi( + s, + (0, 0), + theta_f, + np.arange(200) * 1e-5, + np.arange(200) * 2e-5, + grid="integration", + ) + tr = ("source_0", "source_0") + xi_c = _xi_block(s, tr, grid="reporting") + co = _cosebi_block(s, tr) + xi_f = _xi_block(s, tr, grid="integration") + # dense fine block (CosmoCov integration covariance in production) + fine_block = _spd(len(xi_f), 3) + sio.assemble_covariance( + s, + [(xi_c, _spd(len(xi_c), 1)), (co, _spd(len(co), 2)), (xi_f, fine_block)], + ) + sio.save(s, str(tmp_path / f"{version}.sacc"), type="mock") + + a = sio.load(str(tmp_path / f"{version}.sacc")) + + th_c, p_c, _ = sio.get_xi(a, (0, 0), grid="reporting") + assert np.array_equal(th_c, theta_c) and np.array_equal(p_c, np.arange(20) * 1e-5) + n, E, B = sio.get_cosebis(a, (0, 0)) + assert np.array_equal(n, np.arange(1, 11)) + assert a.covariance.dense.shape == (len(a.mean), len(a.mean)) + + th_f, p_f, _ = sio.get_xi(a, (0, 0), grid="integration") + assert np.array_equal(th_f, theta_f) and np.array_equal(p_f, np.arange(200) * 1e-5) + + # extract() of the fine selection pulls the aligned dense sub-covariance + fine = sio.extract(a, sio.XI_PLUS, tr, grid="integration") + idx_p = a.indices(sio.XI_PLUS, tr, grid="integration") + assert np.allclose(fine.covariance.dense, a.covariance.dense[np.ix_(idx_p, idx_p)]) + # zero cross-blocks between analysis and fine points + assert np.all(a.covariance.dense[np.ix_(xi_c, xi_f)] == 0) + + +def test_one_file_layout_diagonal_fine_fallback(tmp_path): + """No CosmoCov covariance: the fine block is np.diag(varxip/varxim).""" + s = _base_sacc() + theta_f = np.geomspace(0.1, 250.0, 50) + _add_xi(s) + sio.add_xi( + s, + (0, 0), + theta_f, + np.arange(50) * 1e-5, + np.arange(50) * 2e-5, + grid="integration", + ) + tr = ("source_0", "source_0") + xi_c = _xi_block(s, tr, grid="reporting") + xi_f = _xi_block(s, tr, grid="integration") + variances = np.concatenate([np.arange(1, 51) * 1e-12, np.arange(1, 51) * 2e-12]) + sio.assemble_covariance(s, [(xi_c, _spd(len(xi_c), 1)), (xi_f, np.diag(variances))]) + sio.save(s, str(tmp_path / "vDIAG.sacc"), type="mock") + a = sio.load(str(tmp_path / "vDIAG.sacc")) + assert np.array_equal(np.diag(a.covariance.dense[np.ix_(xi_f, xi_f)]), variances) + + +# --------------------------------------------------------------------------- # +# 9. Ascending-grid enforcement (writers reject out-of-order grids) +# --------------------------------------------------------------------------- # +def test_add_xi_rejects_non_ascending_theta(): + s = _base_sacc() + theta = _theta()[::-1] # descending + with pytest.raises(ValueError, match="theta must be strictly ascending"): + sio.add_xi( + s, (0, 0), theta, np.arange(6) * 1e-5, np.arange(6) * 2e-5, grid="reporting" + ) + + +def test_add_pseudo_cl_rejects_non_ascending_ell(): + s = _base_sacc() + ell = np.array([210.0, 30.0, 120.0]) # not ascending + W = np.random.default_rng(0).uniform(size=(20, 3)) + with pytest.raises(ValueError, match="ell_eff must be strictly ascending"): + sio.add_pseudo_cl( + s, + (0, 0), + ell, + np.arange(3) * 1e-9, + np.arange(3) * 2e-9, + np.arange(3) * 3e-9, + window_ells=np.arange(2, 22).astype(float), + window_weights=W, + ) + + +def test_add_pure_eb_rho_tau_reject_non_ascending_theta(): + s = _base_sacc() + theta = _theta()[::-1] + pure = {key: np.arange(6) * 1e-6 for key in sio.PURE_KEYS} + with pytest.raises(ValueError, match="theta must be strictly ascending"): + sio.add_pure_eb(s, (0, 0), theta, **pure) + with pytest.raises(ValueError, match="theta must be strictly ascending"): + sio.add_rho(s, 0, theta, np.arange(6) * 1e-6, np.arange(6) * 2e-6) + with pytest.raises(ValueError, match="theta must be strictly ascending"): + sio.add_tau(s, (0, 0), 0, theta, np.arange(6) * 1e-6, np.arange(6) * 2e-6) + + +# --------------------------------------------------------------------------- # +# 10. Bin-pair normalisation: (1, 0) addresses the same points as (0, 1) +# --------------------------------------------------------------------------- # +def test_bin_pair_normalisation(): + theta = _theta() + xip, xim = np.arange(6) * 1e-5, np.arange(6) * 2e-5 + s = _base_sacc(nbins=2) + sio.add_xi(s, (0, 1), theta, xip, xim, grid="reporting") + th01, p01, m01 = sio.get_xi(s, (0, 1), grid="reporting") + th10, p10, m10 = sio.get_xi(s, (1, 0), grid="reporting") # reversed order + assert np.array_equal(th01, th10) + assert np.array_equal(p01, p10) and np.array_equal(p10, xip) + assert np.array_equal(m01, m10) and np.array_equal(m10, xim) + # writing under (1, 0) lands in the same tracer pair, not a new one + s2 = _base_sacc(nbins=2) + sio.add_xi(s2, (1, 0), theta, xip, xim, grid="reporting") + assert len(s2.indices(sio.XI_PLUS, ("source_0", "source_1"))) == len(theta) + + +# --------------------------------------------------------------------------- # +# 11. Reader/covariance alignment holds for ANY insertion order (the core +# regression the review caught): a tomographic multi-pair ξ covariance +# assembled as ONE contiguous pair-major block, per-pair sub-blocks +# recovered via extract(). +# --------------------------------------------------------------------------- # +def test_tomographic_xi_covariance_one_contiguous_block(): + theta = _theta() + nth = len(theta) + pairs = [(0, 0), (0, 1), (1, 1)] + s = _base_sacc(nbins=2) + for k, (i, j) in enumerate(pairs): + sio.add_xi( + s, + (i, j), + theta, + np.arange(nth) * (k + 1) * 1e-5, + np.arange(nth) * (k + 1) * 2e-5, + grid="reporting", + ) + # All ξ points as one contiguous block in insertion (pair-major) order. + xi_idx = np.arange(len(s.mean)) + assert np.array_equal(xi_idx, np.arange(3 * 2 * nth)) # 3 pairs x [xip; xim] + cov = _spd(len(xi_idx), 11) # dense, cross-pair correlations + sio.assemble_covariance(s, [(xi_idx, cov)]) + + # Per-pair xip sub-block: resolve indices, extract, compare to input. + for i, j in pairs: + idx_p = s.indices( + sio.XI_PLUS, ("source_" + str(min(i, j)), "source_" + str(max(i, j))) + ) + sub = sio.extract( + s, + data_type=sio.XI_PLUS, + tracers=("source_" + str(min(i, j)), "source_" + str(max(i, j))), + ) + assert np.array_equal(sub.covariance.dense, cov[np.ix_(idx_p, idx_p)]) + # readers stay covariance-aligned: get_xi returns in the same order + th, xip, _ = sio.get_xi(s, (i, j), grid="reporting") + assert np.array_equal(th, theta) + assert np.array_equal(s.mean[idx_p], xip) + + +# --------------------------------------------------------------------------- # +# 12. type stamping + fail-closed load (data/mock x concealed/not; escape +# hatch for the blinding/unblinding tooling) +# --------------------------------------------------------------------------- # +def _saved(tmp_path, name, *, type, concealed=None): + s = _base_sacc() + _add_xi(s) + if concealed is not None: + s.metadata["concealed"] = concealed + path = str(tmp_path / f"{name}.sacc") + sio.save(s, path, type=type) + return path + + +def test_load_mock_unconcealed(tmp_path): + s = sio.load(_saved(tmp_path, "m0", type="mock")) + assert s.metadata["type"] == "mock" + + +def test_load_mock_concealed(tmp_path): + s = sio.load(_saved(tmp_path, "m1", type="mock", concealed=True)) + assert s.metadata["concealed"] + + +def test_load_data_concealed(tmp_path): + s = sio.load(_saved(tmp_path, "d1", type="data", concealed=True)) + assert s.metadata["type"] == "data" + + +def test_load_data_unconcealed_fails_closed(tmp_path): + path = _saved(tmp_path, "d0", type="data") + with pytest.raises(ValueError, match="unblinded"): + sio.load(path) + # concealed=False is as unblinded as no stamp at all + path_f = _saved(tmp_path, "d0f", type="data", concealed=False) + with pytest.raises(ValueError, match="unblinded"): + sio.load(path_f) + + +def test_load_data_unconcealed_escape_hatch(tmp_path): + s = sio.load(_saved(tmp_path, "d0h", type="data"), allow_unblinded=True) + assert s.metadata["type"] == "data" + + +def test_save_requires_valid_type(tmp_path): + s = _base_sacc() + with pytest.raises(TypeError): + sio.save(s, str(tmp_path / "x.sacc")) # type is required + with pytest.raises(ValueError, match="'data' or 'mock'"): + sio.save(s, str(tmp_path / "x.sacc"), type="simulation") + + +def test_save_refuses_type_restamp(tmp_path): + s = _base_sacc() + sio.save(s, str(tmp_path / "x.sacc"), type="mock") + with pytest.raises(ValueError, match="re-stamp"): + sio.save(s, str(tmp_path / "x.sacc"), type="data") + + +def test_load_requires_type_tag(tmp_path): + import sacc as sacc_lib + + s = _base_sacc() # never stamped + path = str(tmp_path / "untyped.sacc") + s.save_fits(path, overwrite=True) + with pytest.raises(KeyError): + sio.load(path) + assert sacc_lib.Sacc.load_fits(path) is not None # raw loader still works + + +# --------------------------------------------------------------------------- # +# 13. merge(): per-statistic files combine into one; shared tracers stored +# once; covariance block-diagonal (all-or-none); metadata union with +# loud conflicts. update_statistic(): value-only merge-back. +# --------------------------------------------------------------------------- # +def _xi_sacc(metadata=None): + s = sio.new_sacc({0: _nz(0)}, metadata=metadata) + _add_xi(s) + return s + + +def _cosebi_sacc(metadata=None): + s = sio.new_sacc({0: _nz(0)}, metadata=metadata) + sio.add_cosebis( + s, (0, 0), np.arange(1, 6) * 1e-6, (1.0, 100.0), Bn=np.arange(1, 6) * 1e-7 + ) + return s + + +def test_merge_per_statistic_files(tmp_path): + meta = {"version": "vM", "type": "mock"} + s_xi, s_co = _xi_sacc(meta), _cosebi_sacc(meta) + merged = sio.merge([s_xi, s_co]) + # shared tracers stored once; all points present, xi first + assert set(merged.tracers) == {"source_0", sio.PSF_TRACER} + assert len(merged.mean) == len(s_xi.mean) + len(s_co.mean) + assert np.array_equal(merged.mean, np.concatenate([s_xi.mean, s_co.mean])) + assert merged.metadata["version"] == "vM" and merged.metadata["type"] == "mock" + # inputs untouched + assert s_xi.metadata["version"] == "vM" + # readers work on the merged file after a round-trip + sio.save(merged, str(tmp_path / "vM.sacc"), type="mock") + merged_rt = sio.load(str(tmp_path / "vM.sacc")) + _, p, _ = sio.get_xi(merged_rt, (0, 0), grid="reporting") + assert np.array_equal(p, np.arange(6) * 1e-5) + _, E, _ = sio.get_cosebis(merged_rt, (0, 0)) + assert np.array_equal(E, np.arange(1, 6) * 1e-6) + + +def test_merge_covariance_block_diagonal(): + s_xi, s_co = _xi_sacc(), _cosebi_sacc() + cov_xi, cov_co = _spd(len(s_xi.mean), 1), _spd(len(s_co.mean), 2) + s_xi.add_covariance(cov_xi) + s_co.add_covariance(cov_co) + merged = sio.merge([s_xi, s_co]) + dense = merged.covariance.dense + n_xi = len(s_xi.mean) + assert np.array_equal(dense[:n_xi, :n_xi], cov_xi) + assert np.array_equal(dense[n_xi:, n_xi:], cov_co) + assert np.all(dense[:n_xi, n_xi:] == 0) + + +def test_merge_block_diagonal_covariance_stays_block_diagonal(tmp_path): + """Merging two files that already carry a BlockDiagonalCovariance (e.g. + each assembled via ``assemble_covariance``) must not densify — the + result stays a ``BlockDiagonalCovariance``, on disk too.""" + s_xi, s_co = _xi_sacc(), _cosebi_sacc() + sio.assemble_covariance( + s_xi, [(np.arange(len(s_xi.mean)), _spd(len(s_xi.mean), 1))] + ) + sio.assemble_covariance( + s_co, [(np.arange(len(s_co.mean)), _spd(len(s_co.mean), 2))] + ) + assert type(s_xi.covariance).__name__ == "BlockDiagonalCovariance" + merged = sio.merge([s_xi, s_co]) + assert type(merged.covariance).__name__ == "BlockDiagonalCovariance" + sio.save(merged, str(tmp_path / "vBLK.sacc"), type="mock") + merged_rt = sio.load(str(tmp_path / "vBLK.sacc")) + assert type(merged_rt.covariance).__name__ == "BlockDiagonalCovariance" + n_xi = len(s_xi.mean) + assert np.array_equal( + merged_rt.covariance.dense[:n_xi, :n_xi], s_xi.covariance.dense + ) + assert np.array_equal( + merged_rt.covariance.dense[n_xi:, n_xi:], s_co.covariance.dense + ) + + +def test_merge_mixed_covariance_fails(): + s_xi, s_co = _xi_sacc(), _cosebi_sacc() + s_xi.add_covariance(_spd(len(s_xi.mean), 1)) # s_co has none + with pytest.raises(Exception): + sio.merge([s_xi, s_co]) + + +def test_merge_conflicting_metadata_fails(): + s_xi = _xi_sacc({"type": "data"}) + s_co = _cosebi_sacc({"type": "mock"}) + with pytest.raises(ValueError, match="conflicting metadata"): + sio.merge([s_xi, s_co]) + + +def test_update_statistic_values_only(): + s = _multi_statistic_sacc() + tr = ("source_0", "source_0") + xi, cl, co = _xi_block(s, tr), _cl_block(s, tr), _cosebi_block(s, tr) + cov = [(xi, _spd(len(xi), 1)), (cl, _spd(len(cl), 2)), (co, _spd(len(co), 3))] + sio.assemble_covariance(s, cov) + dense_before = s.covariance.dense.copy() + mean_before = s.mean.copy() + + # extract -> shift (a stand-in for conceal) -> merge back + sub = sio.extract(s, data_type=sio.XI_PLUS, tracers=tr) + for point in sub.data: + point.value += 1e-4 + sio.update_statistic(s, sub) + + idx_p = s.indices(sio.XI_PLUS, tr) + assert np.allclose(s.mean[idx_p], mean_before[idx_p] + 1e-4) + untouched = np.setdiff1d(np.arange(len(s.mean)), idx_p) + assert np.array_equal(s.mean[untouched], mean_before[untouched]) + # covariance and insertion order untouched + assert np.array_equal(s.covariance.dense, dense_before) + + +def test_update_statistic_requires_unique_match(): + s = _xi_sacc() + sub = sio.extract(s, data_type=sio.XI_PLUS, tracers=("source_0", "source_0")) + missing = sub.copy() + missing.data[0].tags["theta"] = 999.0 # matches nothing in s + with pytest.raises(ValueError, match="0 points"): + sio.update_statistic(s, missing) + + +def test_update_statistic_rejects_duplicate_sub_points(): + s = _xi_sacc() + sub = sio.extract(s, data_type=sio.XI_PLUS, tracers=("source_0", "source_0")) + sub.data[1].tags = dict(sub.data[0].tags) # two sub points -> one target + with pytest.raises(ValueError, match="same target"): + sio.update_statistic(s, sub) + + +def test_merge_rejects_divergent_shared_tracer(): + s_xi, s_co = _xi_sacc(), _cosebi_sacc() + s_co.tracers["source_0"].nz = s_co.tracers["source_0"].nz * 2.0 + with pytest.raises(ValueError, match="source_0.*differs"): + sio.merge([s_xi, s_co]) + + +# --------------------------------------------------------------------------- # +# 13b. merge(): theta-consistency guard across groups. +# --------------------------------------------------------------------------- # +def _rho_sacc(k=0, theta=None, scale=1.0, metadata=None): + s = sio.new_sacc({0: _nz(0)}, metadata=metadata) + theta = _theta() if theta is None else theta + sio.add_rho( + s, + k, + theta, + np.arange(len(theta)) * scale * 1e-6, + np.arange(len(theta)) * scale * 2e-6, + ) + return s + + +def test_merge_identical_theta_grids_passes(): + # xi and rho both default to grid='reporting' and share the grid bitwise + s_xi, s_rho = _xi_sacc(), _rho_sacc(theta=_theta()) + s_rho2 = _rho_sacc(k=1, theta=_theta()) + merged = sio.merge([s_xi, s_rho, s_rho2]) + assert len(merged.mean) == sum(len(s.mean) for s in (s_xi, s_rho, s_rho2)) + + +def test_merge_nearly_identical_theta_grids_raises(): + theta = _theta() + s_rho = _rho_sacc(theta=theta) + # same 'reporting' grid group, same length, ~1e-9 relative perturbation + s_rho2 = _rho_sacc(k=1, theta=theta * (1 + 1e-9)) + with pytest.raises(ValueError, match="theta grids under the same grid tag"): + sio.merge([s_rho, s_rho2]) + + +def test_merge_rho_off_xi_reporting_grid_raises(): + s_xi = _xi_sacc() # grid='reporting' + # rho defaults to 'reporting' too: a slightly-off grid must fail loud + s_rho = _rho_sacc(theta=_theta() * (1 + 1e-9)) + with pytest.raises(ValueError, match="theta grids under the same grid tag"): + sio.merge([s_xi, s_rho]) + + +def test_merge_clearly_different_theta_grids_same_tag_raises(): + s_rho = _rho_sacc(theta=_theta()) # theta in [1, 100] + # same 'reporting' group, same length, entirely different binning + s_rho2 = _rho_sacc(k=1, theta=np.geomspace(200.0, 400.0, 6)) + with pytest.raises(ValueError, match="theta grids under the same grid tag"): + sio.merge([s_rho, s_rho2]) + + +def test_merge_different_length_theta_grids_passes(): + s_rho = _rho_sacc(theta=_theta()) # 6-point theta + s_rho2 = _rho_sacc(k=1, theta=_theta(nbins=10)) # same group, subset OK + merged = sio.merge([s_rho, s_rho2]) + assert len(merged.mean) == len(s_rho.mean) + len(s_rho2.mean) + + +def test_merge_same_length_grids_under_different_tags_pass(): + # reporting vs integration differ by design — no cross-tag constraint + s = sio.new_sacc({0: _nz(0)}) + _add_xi(s, grid="reporting") # theta in [1, 100], 6 points + sio.add_xi( + s, + (0, 0), + np.geomspace(200.0, 400.0, 6), # same length, different values + np.arange(6) * 1e-5, + np.arange(6) * 2e-5, + grid="integration", + ) + merged = sio.merge([s, _rho_sacc(theta=_theta())]) + assert len(merged.mean) == len(s.mean) + 12 + + +def _cl_sacc(bin=0, ell=None, W=None, grid="reporting"): + ell = np.array([30.0, 120.0, 210.0, 300.0]) if ell is None else ell + nell, nbp = 50, len(ell) + W = np.random.default_rng(5).uniform(size=(nell, nbp)) if W is None else W + s = sio.new_sacc({bin: _nz(bin)}) + sio.add_pseudo_cl( + s, + (bin, bin), + ell, + np.arange(nbp) * 1e-9, + np.arange(nbp) * 2e-9, + np.arange(nbp) * 3e-9, + window_ells=np.arange(2, 2 + nell).astype(float), + window_weights=W, + grid=grid, + ) + return s + + +def test_merge_identical_ell_grids_and_windows_pass(): + s_a, s_b = _cl_sacc(bin=0), _cl_sacc(bin=1) # same ell, same window + merged = sio.merge([s_a, s_b]) + assert len(merged.mean) == len(s_a.mean) + len(s_b.mean) + + +def test_merge_nearly_identical_ell_grids_raises(): + ell = np.array([30.0, 120.0, 210.0, 300.0]) + s_a = _cl_sacc(bin=0, ell=ell) + s_b = _cl_sacc(bin=1, ell=ell * (1 + 1e-9)) # same 'reporting' group + with pytest.raises(ValueError, match="ell grids under the same grid tag"): + sio.merge([s_a, s_b]) + + +def test_merge_shared_ell_grid_different_windows_raises(): + W = np.random.default_rng(5).uniform(size=(50, 4)) + s_a = _cl_sacc(bin=0, W=W) + s_b = _cl_sacc(bin=1, W=W * (1 + 1e-6)) # same ell, perturbed window + with pytest.raises(ValueError, match="bandpower windows differ"): + sio.merge([s_a, s_b]) + + +def test_merge_different_ell_grids_across_tags_pass(): + s_a = _cl_sacc(bin=0) # grid='reporting' + s_b = _cl_sacc(bin=1, ell=np.array([40.0, 130.0, 220.0, 310.0]), grid="finer") + merged = sio.merge([s_a, s_b]) # different tag values: unconstrained + assert len(merged.mean) == len(s_a.mean) + len(s_b.mean) + + +# --------------------------------------------------------------------------- # +# 15. Unmatched selections fail loud — no silent-empty arrays anywhere. +# --------------------------------------------------------------------------- # +def test_readers_raise_on_unmatched_selection(): + s = _base_sacc() + _add_xi(s, grid="reporting") + with pytest.raises(ValueError, match="matched no points"): + sio.get_xi(s, (0, 0), grid="integration") # only 'reporting' exists + with pytest.raises(ValueError, match="matched no points"): + sio.get_xi(s, (0, 1), grid="reporting") # no such pair + + +def test_get_cosebis_raises_on_unmatched_scale_cut(): + s = _cosebi_sacc() # written with scale cut (1.0, 100.0) + with pytest.raises(ValueError, match="matched no points"): + sio.get_cosebis(s, (0, 0), scale_cut=(2.0, 50.0)) + + +def test_get_cosebis_rejects_ambiguous_multi_cut_file(): + s = _cosebi_sacc() + sio.add_cosebis( + s, (0, 0), np.arange(1, 6) * 1e-6, (2.0, 50.0), Bn=np.arange(1, 6) * 1e-7 + ) + with pytest.raises(ValueError, match="several COSEBIs scale cuts"): + sio.get_cosebis(s, (0, 0)) + n, En, _ = sio.get_cosebis(s, (0, 0), scale_cut=(2.0, 50.0)) + assert len(En) == 5 # explicit cut disambiguates + + +def test_extract_raises_on_empty_selection(): + s = _xi_sacc() + with pytest.raises(ValueError, match="no points"): + sio.extract(s, grid="integration") # only 'reporting' exists + + +def test_assemble_covariance_rejects_empty_selector(): + s = _xi_sacc() + with pytest.raises(ValueError, match="matched no points"): + sio.assemble_covariance( + s, [((sio.XI_PLUS, ("source_9", "source_9")), np.eye(6))] + ) + + +# --------------------------------------------------------------------------- # +# 14. get_pseudo_cl window/cl column correspondence: window column j maps to +# the returned ell_eff[j] (verified through window_ind tags). +# --------------------------------------------------------------------------- # +def test_pseudo_cl_window_column_correspondence(tmp_path): + ell_eff = np.array([30.0, 120.0, 210.0, 300.0]) + nell, nbp = 40, len(ell_eff) + window_ells = np.arange(2, 2 + nell).astype(float) + # Distinct columns so a permutation would be detectable. + W = np.zeros((nell, nbp)) + for b in range(nbp): + W[b * 5 : (b + 1) * 5, b] = 1.0 + ee, bb, eb = np.arange(nbp) * 1e-9, np.arange(nbp) * 2e-9, np.arange(nbp) * 3e-9 + s = _base_sacc() + sio.add_pseudo_cl( + s, + (0, 0), + ell_eff, + ee, + bb, + eb, + window_ells=window_ells, + window_weights=W, + ) + s2 = _roundtrip(s, tmp_path, "clwin") + ell, cl_ee, _, _, window = sio.get_pseudo_cl(s2, (0, 0)) + # returned ell array is in insertion order + assert np.array_equal(ell, ell_eff) + assert np.array_equal(cl_ee, ee) + # window_ind tag on each EE point indexes the matching window column + idx = s2.indices(sio.CL_EE, ("source_0", "source_0")) + for pos, i in enumerate(idx): + col = s2.data[i].tags["window_ind"] + assert col == pos # insertion order preserved => column j <-> ell[j] + assert np.array_equal(window.weight[:, col], W[:, pos]) + + +# --------------------------------------------------------------------------- # +# 16. Optionality: writers omit optional components; readers tolerate their +# absence in composite reads but still fail loud on an explicit selection +# naming a component that isn't there. +# --------------------------------------------------------------------------- # +def test_pseudo_cl_ee_bb_only_no_eb(tmp_path): + """EB is often not even computed — add_pseudo_cl must not require it.""" + ell_eff = np.array([30.0, 120.0, 210.0]) + nell, nbp = 20, len(ell_eff) + window_ells = np.arange(2, 2 + nell).astype(float) + W = np.random.default_rng(1).uniform(size=(nell, nbp)) + ee, bb = np.arange(nbp) * 1e-9, np.arange(nbp) * 2e-9 + s = _base_sacc() + sio.add_pseudo_cl( + s, (0, 0), ell_eff, ee, bb, window_ells=window_ells, window_weights=W + ) + tr = ("source_0", "source_0") + assert len(s.indices(sio.CL_EB, tr)) == 0 + s2 = _roundtrip(s, tmp_path, "cl_ee_bb") + ell, cl_ee, cl_bb, cl_eb, window = sio.get_pseudo_cl(s2, (0, 0)) + assert np.array_equal(ell, ell_eff) + assert np.array_equal(cl_ee, ee) + assert np.array_equal(cl_bb, bb) + assert cl_eb is None + + +def test_pseudo_cl_ee_only(tmp_path): + ell_eff = np.array([30.0, 120.0, 210.0]) + nell, nbp = 20, len(ell_eff) + W = np.random.default_rng(2).uniform(size=(nell, nbp)) + ee = np.arange(nbp) * 1e-9 + s = _base_sacc() + sio.add_pseudo_cl( + s, + (0, 0), + ell_eff, + ee, + window_ells=np.arange(2, 2 + nell).astype(float), + window_weights=W, + ) + s2 = _roundtrip(s, tmp_path, "cl_ee_only") + ell, cl_ee, cl_bb, cl_eb, window = sio.get_pseudo_cl(s2, (0, 0)) + assert np.array_equal(cl_ee, ee) + assert cl_bb is None + assert cl_eb is None + + +def test_cosebis_en_only_no_bn(tmp_path): + En = np.arange(1, 8) * 1e-6 + s = _base_sacc() + sio.add_cosebis(s, (0, 0), En, (1.0, 100.0)) + tr = ("source_0", "source_0") + assert len(s.indices(sio.COSEBI_BB, tr)) == 0 + s2 = _roundtrip(s, tmp_path, "cosebi_e_only") + n, E, B = sio.get_cosebis(s2, (0, 0)) + assert np.array_equal(E, En) + assert B is None + + +def test_pure_eb_e_only_no_b_no_amb(tmp_path): + theta = _theta() + xip_E, xim_E = np.arange(6) * 1e-6, np.arange(6) * 2e-6 + s = _base_sacc() + sio.add_pure_eb(s, (0, 0), theta, xip_E, xim_E) + s2 = _roundtrip(s, tmp_path, "pureeb_e_only") + th, back = sio.get_pure_eb(s2, (0, 0)) + assert np.array_equal(th, theta) + assert set(back) == {"xip_E", "xim_E"} + assert np.array_equal(back["xip_E"], xip_E) + assert np.array_equal(back["xim_E"], xim_E) + + +def test_pure_eb_e_and_b_no_amb(tmp_path): + theta = _theta() + arrays = { + key: np.arange(6) * (i + 1) * 1e-6 + for i, key in enumerate(("xip_E", "xim_E", "xip_B", "xim_B")) + } + s = _base_sacc() + sio.add_pure_eb(s, (0, 0), theta, **arrays) + s2 = _roundtrip(s, tmp_path, "pureeb_e_b") + th, back = sio.get_pure_eb(s2, (0, 0)) + assert set(back) == {"xip_E", "xim_E", "xip_B", "xim_B"} + + +def test_pure_eb_rejects_half_a_pair(): + theta = _theta() + xip_E, xim_E = np.arange(6) * 1e-6, np.arange(6) * 2e-6 + s = _base_sacc() + with pytest.raises(ValueError, match="xip_B and xim_B"): + sio.add_pure_eb(s, (0, 0), theta, xip_E, xim_E, xip_B=np.arange(6) * 1e-6) + with pytest.raises(ValueError, match="xip_amb and xim_amb"): + sio.add_pure_eb(s, (0, 0), theta, xip_E, xim_E, xim_amb=np.arange(6) * 1e-6) + + +def test_xi_only_plus_covariance_file(tmp_path): + """A file with only xi +/- and a covariance — no Cl/COSEBIs at all.""" + theta = _theta() + s = _base_sacc() + xip, xim = _add_xi(s) + tr = ("source_0", "source_0") + xi_idx = _xi_block(s, tr) + sio.assemble_covariance(s, [(xi_idx, _spd(len(xi_idx), 7))]) + s2 = _roundtrip(s, tmp_path, "xi_only") + th, p, m = sio.get_xi(s2, (0, 0), grid="reporting") + assert np.array_equal(th, theta) + assert np.array_equal(p, xip) + assert np.array_equal(m, xim) + assert s2.covariance is not None + with pytest.raises(ValueError, match="matched no points"): + sio._indices(s2, sio.CL_EE, tr) + + +def test_reader_explicit_selection_fails_loud_on_missing_optional_component(): + """Absence is silent for composite readers, but a targeted selection + naming a missing optional component (e.g. CL_EB) still fails loud.""" + ell_eff = np.array([30.0, 120.0, 210.0]) + nell, nbp = 20, len(ell_eff) + W = np.random.default_rng(3).uniform(size=(nell, nbp)) + s = _base_sacc() + sio.add_pseudo_cl( + s, + (0, 0), + ell_eff, + np.arange(nbp) * 1e-9, + window_ells=np.arange(2, 2 + nell).astype(float), + window_weights=W, + ) + tr = ("source_0", "source_0") + with pytest.raises(ValueError, match="matched no points"): + sio._mean(s, sio.CL_EB, tr) + with pytest.raises(ValueError, match="no points"): + sio.extract(s, data_type=sio.CL_EB, tracers=tr) + + +def test_merge_partial_files(): + """merge() combines a Cl file missing EB with a COSEBIs file missing Bn.""" + s_cl = sio.new_sacc({0: _nz(0)}) + ell_eff = np.array([30.0, 120.0, 210.0]) + nell, nbp = 20, len(ell_eff) + W = np.random.default_rng(4).uniform(size=(nell, nbp)) + sio.add_pseudo_cl( + s_cl, + (0, 0), + ell_eff, + np.arange(nbp) * 1e-9, + np.arange(nbp) * 2e-9, # BB only, no EB + window_ells=np.arange(2, 2 + nell).astype(float), + window_weights=W, + ) + s_co = sio.new_sacc({0: _nz(0)}) + sio.add_cosebis(s_co, (0, 0), np.arange(1, 6) * 1e-6, (1.0, 100.0)) # En only + + merged = sio.merge([s_cl, s_co]) + tr = ("source_0", "source_0") + assert len(merged.indices(sio.CL_EB, tr)) == 0 + assert len(merged.indices(sio.COSEBI_BB, tr)) == 0 + ell, cl_ee, cl_bb, cl_eb, window = sio.get_pseudo_cl(merged, (0, 0)) + assert cl_bb is not None and cl_eb is None + n, E, B = sio.get_cosebis(merged, (0, 0)) + assert B is None