Add generic auxiliary array bundles to the fit input HDF5

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

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)

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):

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()