theory_fn protocol, default CCL backend, pure order-independent shift draw

src/smokescreen/__main__.py

@@ -3,103 +3,84 @@
 from typing import Union, Dict, Tuple
 from jsonargparse import CLI
 from jsonargparse.typing import Path_drw, Path_fr
-from pyccl import Cosmology as CosmologyType
-import pyccl as ccl
 # warnings related to sacc files
 import warnings
 from smokescreen import ConcealDataVector
 from smokescreen.encryption import encrypt_file, decrypt_file
-from smokescreen.utils import load_cosmology_from_partial_dict, load_sacc_file
+from smokescreen.utils import load_sacc_file
 from . import __version__
 warnings.filterwarnings("ignore")
 
 # banner to be printed in the terminal
 banner = rf"""
 
- (                                                          
- )\ )                 )                                     
-(()/(    )         ( /(    (          (      (    (         
- /(_))  (      (   )\())  ))\ (    (  )(    ))\  ))\  (     
-(_))    )\  '  )\ ((_)\  /((_))\   )\(()\  /((_)/((_) )\ )  
-/ __| _((_))  ((_)| |(_)(_)) ((_) ((_)((_)(_)) (_))  _(_/(  
-\__ \| '  \()/ _ \| / / / -_)(_-</ _|| '_|/ -_)/ -_)| ' \)) 
-|___/|_|_|_| \___/|_\_\ \___|/__/\__||_|  \___|\___||_||_|  
-                                                            
+ (
+ )\ )                 )
+(()/(    )         ( /(    (          (      (    (
+ /(_))  (      (   )\())  ))\ (    (  )(    ))\  ))\  (
+(_))    )\  '  )\ ((_)\  /((_))\   )\(()\  /((_)/((_) )\ )
+/ __| _((_))  ((_)| |(_)(_)) ((_) ((_)((_)(_)) (_))  _(_/(
+\__ \| '  \()/ _ \| / / / -_)(_-</ _|| '_|/ -_)/ -_)| ' \))
+|___/|_|_|_| \___/|_\_\ \___|/__/\__||_|  \___|\___||_||_|
+
  - DESC Pipeline for Concealing your Cosmology Results -
                  Version {__version__}
 """
 
 
 def datavector_main(path_to_sacc: Path_fr,
-                    likelihood_path: str,
+                    fiducial_params: Dict[str, float],
                     shifts_dict: Dict[str, Union[float, Tuple[float, float]]],
-                    systematics: dict = None,
+                    seed: Union[int, str],
                     shift_type: str = 'add',
                     shift_distribution: str = 'flat',
-                    seed: Union[int, str] = 2112,
-                    # Note: dict must come first in Union for jsonargparse to correctly
-                    # parse partial cosmology dictionaries without trying to instantiate Cosmology
-                    reference_cosmology: Union[dict, CosmologyType] = ccl.CosmologyVanillaLCDM(),
                     path_to_output: Path_drw = None,
                     keep_original_sacc: bool = False,
                     output_suffix: str = None,
                     ) -> None:
-    r"""Main function to conceal a SACC file using a firecrown likelihood.
+    r"""Conceal a cosmic-shear SACC file with the default CCL theory backend.
 
     Args:
-        path_to_sacc (str): Path to the sacc file to blind.
-        likelihood_path (str): Path to the firecrown likelihood module file.
-        shifts_dict (dict): Dictionary with fixed values for the firecrown shifts parameters.
-            Example: {"Omega_c": (0.20, 0.39), "sigma8": (0.70, 0.90)}
-        shift_type (str): Type of shift to apply to the data vector.
-            Options are 'add' and 'mult'. Defaults to 'add'.
-        systematics (dict): Dictionary with fixed values for the firecrown systematics parameters.
-        shift_distribution (str): Distribution type for the parameter shifts.
-            Options are 'flat' and 'gaussian'. Defaults to 'flat'.
-        seed (int, str): Seed for the blinding process. Defaults to 2112.
-        reference_cosmology (Union[CosmologyType, dict]):
-            Cosmology object or dictionary with cosmological
-            parameters you want different than the VanillaLCDM as reference cosmology.
-            Defaults to ccl.CosmologyVanillaLCDM().
-        path_to_output (str): Path to save the blinded sacc file. Defaults to None.
-        keep_original_sacc (bool): If True, keeps the original sacc file.
-            Defaults to False [keeps only the encrypted file].
+        path_to_sacc (str): Path to the SACC file to blind. It must contain
+            exactly the cosmic-shear rows the default CCL backend models
+            (galaxy_shear_cl_ee and/or galaxy_shear_xi_plus/minus), with
+            weak-lensing tracers carrying n(z).
+        fiducial_params (dict): Fiducial cosmological parameters (CCL-native
+            names, e.g. {"sigma8": 0.8, "Omega_c": 0.25, "Omega_b": 0.05,
+            "h": 0.67, "n_s": 0.96}).
+        shifts_dict (dict): Shift envelopes, interpreted as deltas about zero.
+            Example: {"Omega_c": (-0.05, 0.05), "sigma8": 0.05}
+        seed (int, str): Seed for the blinding process (no default; must be a
+            deliberate, secret choice).
+        shift_type (str): Concealing factor type, 'add' or 'mult'. Default 'add'.
+        shift_distribution (str): 'flat' or 'gaussian'. Default 'flat'.
+        path_to_output (str): Directory to save the blinded SACC. Default None
+            (uses the input file's directory).
+        keep_original_sacc (bool): If True, keeps the original SACC file.
+            Default False (keeps only the encrypted file).
         output_suffix (str): Custom suffix for the output file name.
-            Defaults to None (uses 'concealed_data_vector').
     """
     print(banner)
-    if isinstance(reference_cosmology, dict):
-        cosmo = load_cosmology_from_partial_dict(reference_cosmology)
-    else:
-        cosmo = reference_cosmology
-    # tests if the sacc file exists
     assert os.path.exists(path_to_sacc), f"File {path_to_sacc} does not exist."
-    assert os.path.exists(likelihood_path), f"File {likelihood_path} does not exist."
     # reads the sacc file (returns sacc object and detected format)
     sacc_data, input_format = load_sacc_file(path_to_sacc)
-    # creates the smokescreen object
-    smoke = ConcealDataVector(cosmo,  likelihood_path, shifts_dict, sacc_data, systematics, seed,
-                              shift_distr=shift_distribution, input_format=input_format)
+    # creates the smokescreen object with the default CCL backend
+    smoke = ConcealDataVector(fiducial_params, shifts_dict, sacc_data,
+                              seed=seed, shift_distr=shift_distribution,
+                              input_format=input_format)
     # blinds the sacc file
     smoke.calculate_concealing_factor(factor_type=shift_type)
-    # applies the blinding factor to the sacc file
     smoke.apply_concealing_to_likelihood_datavec()
     print(f">> User {getpass.getuser()}",
           f"used Smokescreen on {path_to_sacc} ... it is super effective!")
     # get root name of the input file
     root_name = os.path.splitext(os.path.basename(path_to_sacc))[0]
     # saves the blinded sacc file
-    if path_to_output is not None:
-        smoke.save_concealed_datavector(path_to_output, root_name,
-                                        output_format=input_format,
-                                        suffix=output_suffix)
-    else:
-        # get the input file directory
+    if path_to_output is None:
         path_to_output = os.path.dirname(path_to_sacc)
-        smoke.save_concealed_datavector(path_to_output, root_name,
-                                        output_format=input_format,
-                                        suffix=output_suffix)
-    # Determine extension based on format
+    smoke.save_concealed_datavector(path_to_output, root_name,
+                                    output_format=input_format,
+                                    suffix=output_suffix)
     ext = '.hdf5' if input_format == 'hdf5' else '.fits'
     _suffix = output_suffix if output_suffix is not None else "concealed_data_vector"
     outprintfile = f"{path_to_output}/{root_name}_{_suffix}{ext}"

src/smokescreen/backends/__init__.py

@@ -0,0 +1,10 @@
+# author: Arthur Loureiro <arthur.loureiro@fysik.su.se>
+# license: BSD 3-Clause
+'''
+Theory backends (:mod:`smokescreen.backends`)
+=============================================
+
+Built-in ``theory_fn`` implementations. The default CCL cosmic-shear backend
+lives in :mod:`smokescreen.backends.ccl`; it is imported lazily (importing this
+package does not import ``pyccl``).
+'''

src/smokescreen/backends/ccl.py

@@ -0,0 +1,146 @@
+# author: Arthur Loureiro <arthur.loureiro@fysik.su.se>
+# license: BSD 3-Clause
+'''
+Default CCL cosmic-shear backend (:mod:`smokescreen.backends.ccl`)
+=================================================================
+
+.. currentmodule:: smokescreen.backends.ccl
+
+Builds a default ``theory_fn`` that computes a cosmic-shear theory vector from
+CCL for a standard weak-lensing SACC file. ``pyccl`` is imported here, so it
+enters ``sys.modules`` only once this backend is constructed --- importing
+``smokescreen`` does not import it.
+
+The backend accepts CCL-native cosmological-parameter names
+(``sigma8``, ``Omega_c``, ``Omega_b``, ``h``, ``n_s``, ...). It reads the
+weak-lensing tracers' n(z) and the SACC row layout, and returns, for a given
+parameter mapping, the theory vector aligned element-for-element to
+``sacc_data.mean``.
+
+Two conventions of this convenience backend, part of its row-to-theory
+contract (Smokescreen does not verify them):
+
+- **Transfer function / power spectrum.** It defaults to CCL's native
+  ``transfer_function="eisenstein_hu"`` and ``matter_power_spectrum="halofit"``
+  so the shipped backend runs against a bare ``pyccl`` install with no
+  Boltzmann code (CAMB/CLASS). Both are overridable through ``cosmo_params``;
+  a caller who wants a Boltzmann-code power spectrum can supply their own
+  ``theory_fn`` instead.
+- **ξ± angle units.** The SACC ``theta`` tag of ``galaxy_shear_xi_plus`` /
+  ``galaxy_shear_xi_minus`` rows is interpreted as **arcminutes** and converted
+  to degrees for :func:`pyccl.correlation`. ``galaxy_shear_cl_ee`` rows use
+  their ``ell`` tag directly.
+
+.. autofunction:: build_ccl_theory_fn
+'''
+import numpy as np
+import pyccl as ccl
+
+
+# SACC data types the default backend understands, and the CCL correlation
+# 'type' string (Fourier -> real space) each ξ± measurement maps to.
+_XI_TYPES = {
+    "galaxy_shear_xi_plus": "GG+",
+    "galaxy_shear_xi_minus": "GG-",
+}
+_CL_TYPES = {"galaxy_shear_cl_ee"}
+
+
+def _tracer_nz(sacc_data, name):
+    """Return (z, nz) for a SACC tracer carrying a redshift distribution."""
+    tracer = sacc_data.tracers[name]
+    if not (hasattr(tracer, "z") and hasattr(tracer, "nz")):
+        raise ValueError(
+            f"Default CCL backend requires n(z) tracers; tracer '{name}' "
+            f"({tracer.tracer_type}) carries no redshift distribution."
+        )
+    return np.asarray(tracer.z), np.asarray(tracer.nz)
+
+
+def _plan_blocks(sacc_data):
+    """
+    Decompose the SACC into cosmic-shear blocks aligned to ``sacc_data.mean``.
+
+    Returns a list of ``(data_type, tracer_pair, angle_values, row_indices)``,
+    one per (data_type, tracer-pair) group the default backend supports. Raises
+    if the SACC carries a data type this backend does not model.
+    """
+    supported = set(_XI_TYPES) | _CL_TYPES
+    present = set(sacc_data.get_data_types())
+    unsupported = present - supported
+    if unsupported:
+        raise ValueError(
+            f"Default CCL backend only models cosmic-shear data types "
+            f"{sorted(supported)}; SACC carries unsupported {sorted(unsupported)}. "
+            f"Supply an explicit theory_fn for this SACC."
+        )
+
+    blocks = []
+    for data_type in present:
+        angle_key = "theta" if data_type in _XI_TYPES else "ell"
+        for pair in sacc_data.get_tracer_combinations(data_type):
+            idx = sacc_data.indices(data_type, pair)
+            angles = np.array(
+                [sacc_data.data[i].get_tag(angle_key) for i in idx]
+            )
+            blocks.append((data_type, pair, angles, np.asarray(idx)))
+    return blocks
+
+
+def build_ccl_theory_fn(sacc_data):
+    """
+    Build a default cosmic-shear ``theory_fn`` from a SACC file.
+
+    Parameters
+    ----------
+    sacc_data : sacc.sacc.Sacc
+        SACC carrying weak-lensing tracers with n(z) and cosmic-shear rows
+        (``galaxy_shear_cl_ee`` and/or ``galaxy_shear_xi_plus`` /
+        ``galaxy_shear_xi_minus``). The row layout it exposes is the layout the
+        returned callable aligns to.
+
+    Returns
+    -------
+    Callable[[Mapping[str, float]], np.ndarray]
+        A ``theory_fn``: given a mapping of CCL-native cosmological-parameter
+        names to values, it constructs a :class:`pyccl.Cosmology`, builds a
+        :class:`pyccl.WeakLensingTracer` per bin, and returns the cosmic-shear
+        theory vector aligned to ``sacc_data.mean``.
+    """
+    blocks = _plan_blocks(sacc_data)
+    n_rows = len(sacc_data.mean)
+    nz_cache = {
+        name: _tracer_nz(sacc_data, name)
+        for block in blocks
+        for name in block[1]
+    }
+    # ℓ grid for the Cℓ used by the real-space correlation transform.
+    ell_grid = np.unique(np.geomspace(2, 30000, 512).astype(int))
+
+    def theory_fn(cosmo_params):
+        # Default to pyccl-native transfer/power (no CAMB/CLASS dependency), so
+        # the shipped backend runs against a bare pyccl install. A caller-supplied
+        # theory_fn is free to route through a Boltzmann code.
+        params = {"transfer_function": "eisenstein_hu",
+                  "matter_power_spectrum": "halofit"}
+        params.update(cosmo_params)
+        cosmo = ccl.Cosmology(**params)
+        wl = {
+            name: ccl.WeakLensingTracer(cosmo, dndz=nz_cache[name])
+            for name in nz_cache
+        }
+        vec = np.empty(n_rows)
+        for data_type, (t1, t2), angles, idx in blocks:
+            cl = ccl.angular_cl(cosmo, wl[t1], wl[t2], ell_grid)
+            if data_type in _CL_TYPES:
+                values = np.interp(angles, ell_grid, cl)
+            else:
+                theta_deg = angles / 60.0  # arcmin -> degrees
+                values = ccl.correlation(
+                    cosmo, ell=ell_grid, C_ell=cl, theta=theta_deg,
+                    type=_XI_TYPES[data_type],
+                )
+            vec[idx] = values
+        return vec
+
+    return theory_fn