diff --git a/rabbit/auxiliary.py b/rabbit/auxiliary.py new file mode 100644 index 0000000..6a54b20 --- /dev/null +++ b/rabbit/auxiliary.py @@ -0,0 +1,109 @@ +"""Generic ``auxiliary`` array bundles carried through the fit HDF5. + +An auxiliary entry is a named bundle of arbitrary arrays — numeric ndarrays +and/or 1-D string lists — stashed in the input HDF5 under a top-level +``auxiliary`` group. It is not used by the fit itself. + +This mirrors the ``external_terms`` mechanism (see +:mod:`rabbit.external_likelihood`): :class:`rabbit.tensorwriter.TensorWriter` +collects bundles via ``add_auxiliary`` and serializes them in ``write()``; +:class:`rabbit.inputdata.FitInputData` exposes them as ``self.auxiliary``. + +Round-trip guarantees: + +* numeric arrays survive bit-for-bit including dtype and shape, via + ``writeFlatInChunks`` (which stamps an ``original_shape`` attr) and + ``maketensor`` (which restores it); +* 1-D string lists (e.g. axis names) survive as a Python ``list[str]``, stored + as a vlen-str dataset like ``external_terms``' ``params``. +""" + +import h5py +import numpy as np + +from rabbit import h5pyutils_write +from rabbit.h5pyutils_read import maketensor + + +def _is_string_array(arr): + """True if ``arr`` should be stored as strings rather than numbers.""" + return arr.dtype.kind in ("U", "S", "O") + + +def write_auxiliary_group(parent, auxiliary, maxChunkBytes=1024**2): + """Serialize auxiliary bundles under ``parent``'s ``auxiliary`` group. + + Parameters + ---------- + parent : h5py.Group + Open HDF5 group/file to create the ``auxiliary`` subgroup in. + auxiliary : list[dict] + Bundles ``{"name": str, "datasets": {key: ndarray | list[str]}}`` as + collected by ``TensorWriter.add_auxiliary``. + maxChunkBytes : int + Chunk size passed through to ``writeFlatInChunks`` for numeric arrays. + + Returns + ------- + int + Number of raw array bytes written (0 if ``auxiliary`` is empty). + """ + if not auxiliary: + return 0 + + nbytes = 0 + aux_group = parent.create_group("auxiliary") + for aux in auxiliary: + g = aux_group.create_group(aux["name"]) + for key, val in aux["datasets"].items(): + arr = np.asarray(val) + if _is_string_array(arr): + # 1-D list of strings (e.g. axis names) -> vlen-str dataset, + # mirroring external_terms' "params". + flat = arr.reshape(-1) + ds = g.create_dataset( + key, + [flat.size], + dtype=h5py.special_dtype(vlen=str), + compression="gzip", + ) + ds[...] = [str(s) for s in flat] + else: + # numeric array -> flat chunked; shape recovered by maketensor. + nbytes += h5pyutils_write.writeFlatInChunks( + arr, g, key, maxChunkBytes=maxChunkBytes + ) + return nbytes + + +def read_auxiliary_from_h5(aux_group): + """Decode an HDF5 ``auxiliary`` group. + + Parameters + ---------- + aux_group : h5py.Group or None + The ``auxiliary`` group in the input HDF5 file, or ``None``. + + Returns + ------- + dict + ``{name: {key: ndarray | list[str]}}``. Numeric datasets are decoded via + ``maketensor`` (shape restored from ``original_shape``); string datasets + are decoded to a Python ``list[str]``. Empty dict if ``aux_group`` is + ``None``. + """ + if aux_group is None: + return {} + + out = {} + for name, g in aux_group.items(): + bundle = {} + for key, ds in g.items(): + if h5py.check_string_dtype(ds.dtype): + bundle[key] = [ + s.decode() if isinstance(s, bytes) else str(s) for s in ds[...] + ] + else: + bundle[key] = np.asarray(maketensor(ds)) + out[name] = bundle + return out diff --git a/rabbit/inputdata.py b/rabbit/inputdata.py index 250a43f..7cf1385 100644 --- a/rabbit/inputdata.py +++ b/rabbit/inputdata.py @@ -199,6 +199,11 @@ def __init__(self, filename, pseudodata=None): self.external_terms = read_external_terms_from_h5(f.get("external_terms")) + # Load generic auxiliary array bundles (optional). See rabbit.auxiliary. + from rabbit.auxiliary import read_auxiliary_from_h5 + + self.auxiliary = read_auxiliary_from_h5(f.get("auxiliary")) + @tf.function def expected_events_nominal(self): rnorm = tf.ones(self.nproc, dtype=self.dtype) diff --git a/rabbit/tensorwriter.py b/rabbit/tensorwriter.py index 1ba08e3..44a98e9 100644 --- a/rabbit/tensorwriter.py +++ b/rabbit/tensorwriter.py @@ -6,7 +6,7 @@ import numpy as np from wums.sparse_hist import SparseHist # noqa: F401 re-exported for convenience -from rabbit import common, h5pyutils_write +from rabbit import auxiliary, common, h5pyutils_write from wums import ioutils, logging # isort: skip @@ -73,6 +73,11 @@ def __init__( # add_external_likelihood_term for details. self.external_terms = [] + # Generic auxiliary array bundles (not used by the fit). Each entry is + # {"name": str, "datasets": {key: ndarray | list[str]}}; see + # add_auxiliary and rabbit.auxiliary. + self.auxiliary = [] + self.clipSystVariations = False if self.clipSystVariations > 0.0: self.clip = np.abs(np.log(self.clipSystVariations)) @@ -1442,6 +1447,28 @@ def add_external_likelihood_term(self, grad=None, hess=None, name=None): } ) + def add_auxiliary(self, name, datasets): + """Store a named bundle of arbitrary arrays in the output. + + The bundle is written under the top-level ``auxiliary`` HDF5 group and + exposed on the read side as ``FitInputData.auxiliary[name]``. It is not + used by the fit itself; it is a side channel for ParamModels to carry + pre-computed inputs (e.g. a response matrix) that must stay consistent + with the datacard. See :mod:`rabbit.auxiliary`. + + Parameters + ---------- + name : str + Bundle identifier (the auxiliary subgroup name). Must be unique. + datasets : dict + ``{key: np.ndarray | list[str]}``. Numeric arrays round-trip + bit-for-bit (dtype + shape); 1-D string lists round-trip as + ``list[str]``. + """ + if any(a["name"] == name for a in self.auxiliary): + raise RuntimeError(f"auxiliary '{name}' already added") + self.auxiliary.append({"name": name, "datasets": dict(datasets)}) + @staticmethod def _sparse_values_at(sparse_csr, indices): """Extract values from a flat CSR array at the given flat indices. @@ -2212,6 +2239,11 @@ def create_dataset( maxChunkBytes=self.chunkSize, ) + # Write generic auxiliary array bundles (not used by the fit). See rabbit.auxiliary. + nbytes += auxiliary.write_auxiliary_group( + f, self.auxiliary, maxChunkBytes=self.chunkSize + ) + logger.info(f"Total raw bytes in arrays = {nbytes}") def get_systsstandard(self): diff --git a/tests/test_auxiliary.py b/tests/test_auxiliary.py new file mode 100644 index 0000000..4b309a0 --- /dev/null +++ b/tests/test_auxiliary.py @@ -0,0 +1,181 @@ +"""Test the generic ``auxiliary`` array bundles carried through the fit HDF5. + +Covers: + * round-trip through the real path: TensorWriter.add_auxiliary -> write() -> + FitInputData.auxiliary, asserting numeric arrays survive bit-for-bit + (dtype + shape, incl. >2D) and string lists survive as list[str]; + * a datacard with no auxiliary reads back as an empty dict; + * add_auxiliary rejects a duplicate bundle name; + * a direct write_auxiliary_group/read_auxiliary_from_h5 round-trip with + multiple bundles and mixed dtypes (no full datacard needed). + +This mirrors the scetlib_np response-matrix use case (R reco x gen, N_gen, +axis names + edges) without depending on WRemnants. +""" + +import os +import tempfile + +import h5py +import hist +import numpy as np + +from rabbit import inputdata, tensorwriter +from rabbit.auxiliary import read_auxiliary_from_h5, write_auxiliary_group + + +def build_minimal_writer(): + """A minimal valid TensorWriter: one channel, data, one signal process, + one shape systematic (matches the proven test_external_term setup).""" + np.random.seed(0) + ax = hist.axis.Regular(20, -5, 5, name="x") + + h_data = hist.Hist(ax, storage=hist.storage.Double()) + h_bkg = hist.Hist(ax, storage=hist.storage.Weight()) + + x_bkg = np.random.uniform(-5, 5, 5000) + h_data.fill(x_bkg) + h_bkg.fill(x_bkg, weight=np.ones(len(x_bkg))) + + weights = 0.01 * (ax.centers - ax.centers[0]) - 0.05 + h_up = h_bkg.copy() + h_dn = h_bkg.copy() + h_up.values()[...] = h_bkg.values() * (1 + weights) + h_dn.values()[...] = h_bkg.values() * (1 - weights) + + writer = tensorwriter.TensorWriter() + writer.add_channel([ax], "ch0") + writer.add_data(h_data, "ch0") + writer.add_process(h_bkg, "bkg", "ch0", signal=True) + writer.add_systematic([h_up, h_dn], "shape", "bkg", "ch0", symmetrize="average") + return writer + + +def make_scetlib_np_bundle(): + """A scetlib_np-shaped bundle: multi-dim float64 R, 1-D float64 N_gen, + reco/gen axis name lists, and one edges array per axis.""" + np.random.seed(42) + reco_axes = ["ptll", "yll"] + gen_axes = ["ptVGen", "absYVGen"] + reco_shape = (4, 3) + gen_shape = (5, 2) + R = np.random.uniform(0.0, 1.0, size=reco_shape + gen_shape).astype(np.float64) + N_gen = np.random.uniform(1.0, 10.0, size=gen_shape).astype(np.float64) + edges = { + "edges__ptll": np.array([0.0, 5.0, 10.0, 20.0, 44.0], dtype=np.float64), + "edges__yll": np.array([0.0, 1.0, 2.0, 2.5], dtype=np.float64), + "edges__ptVGen": np.array([0.0, 4.0, 8.0, 16.0, 44.0, 100.0], dtype=np.float64), + "edges__absYVGen": np.array([0.0, 1.25, 2.5], dtype=np.float64), + } + datasets = { + "R": R, + "N_gen": N_gen, + "reco_axes": reco_axes, + "gen_axes": gen_axes, + **edges, + } + return datasets + + +def assert_bundle_equal(got, expected): + assert set(got.keys()) == set( + expected.keys() + ), f"keys differ: {sorted(got)} != {sorted(expected)}" + for key, exp in expected.items(): + val = got[key] + if isinstance(exp, list): # string list + assert val == exp, f"{key}: {val} != {exp}" + else: # numeric array, bit-for-bit incl dtype + shape + exp_arr = np.asarray(exp) + assert ( + val.shape == exp_arr.shape + ), f"{key}: shape {val.shape} != {exp_arr.shape}" + assert ( + val.dtype == exp_arr.dtype + ), f"{key}: dtype {val.dtype} != {exp_arr.dtype}" + assert np.array_equal(val, exp_arr), f"{key}: values differ" + + +def test_through_writer(tmpdir): + datasets = make_scetlib_np_bundle() + + writer = build_minimal_writer() + writer.add_auxiliary("scetlib_np", datasets) + writer.write(outfolder=tmpdir, outfilename="with_aux") + + indata_obj = inputdata.FitInputData(os.path.join(tmpdir, "with_aux.hdf5")) + assert "scetlib_np" in indata_obj.auxiliary, "scetlib_np bundle missing on read" + assert_bundle_equal(indata_obj.auxiliary["scetlib_np"], datasets) + print("PASS: add_auxiliary -> write -> FitInputData round-trip (scetlib_np)") + + +def test_no_auxiliary(tmpdir): + writer = build_minimal_writer() + writer.write(outfolder=tmpdir, outfilename="no_aux") + indata_obj = inputdata.FitInputData(os.path.join(tmpdir, "no_aux.hdf5")) + assert ( + indata_obj.auxiliary == {} + ), f"expected empty auxiliary, got {indata_obj.auxiliary}" + print("PASS: datacard with no auxiliary reads back as empty dict") + + +def test_duplicate_name_guard(): + writer = build_minimal_writer() + writer.add_auxiliary("dup", {"a": np.zeros(3)}) + try: + writer.add_auxiliary("dup", {"b": np.ones(3)}) + except RuntimeError as exc: + assert "already added" in str(exc) + print("PASS: add_auxiliary rejects a duplicate bundle name") + else: + raise AssertionError("expected RuntimeError on duplicate auxiliary name") + + +def test_direct_roundtrip(tmpdir): + """write_auxiliary_group / read_auxiliary_from_h5 directly, no datacard. + Exercises multiple bundles and mixed dtypes (float32, int, >2D).""" + bundles = [ + { + "name": "b0", + "datasets": { + "arr3d": np.arange(24, dtype=np.float64).reshape(2, 3, 4), + "f32": np.array([1.5, 2.5], dtype=np.float32), + "ints": np.array([[1, 2], [3, 4]], dtype=np.int64), + "names": ["alpha", "beta", "gamma"], + }, + }, + { + "name": "b1", + "datasets": {"x": np.linspace(0, 1, 7, dtype=np.float64)}, + }, + ] + path = os.path.join(tmpdir, "direct.hdf5") + with h5py.File(path, "w") as f: + write_auxiliary_group(f, bundles) + with h5py.File(path, "r") as f: + out = read_auxiliary_from_h5(f.get("auxiliary")) + + assert set(out.keys()) == {"b0", "b1"} + for b in bundles: + assert_bundle_equal(out[b["name"]], b["datasets"]) + # explicit dtype checks (assert_bundle_equal already covers, but be loud) + assert out["b0"]["f32"].dtype == np.float32 + assert out["b0"]["ints"].dtype == np.int64 + assert out["b0"]["arr3d"].shape == (2, 3, 4) + # empty group -> {} + assert read_auxiliary_from_h5(None) == {} + print("PASS: direct write_auxiliary_group/read_auxiliary_from_h5 round-trip") + + +def main(): + with tempfile.TemporaryDirectory() as tmpdir: + test_through_writer(tmpdir) + test_no_auxiliary(tmpdir) + test_duplicate_name_guard() + test_direct_roundtrip(tmpdir) + print() + print("ALL CHECKS PASSED") + + +if __name__ == "__main__": + main()