atlas-ftag-tools 0.2.11__py3-none-any.whl → 0.2.13__py3-none-any.whl

{atlas_ftag_tools-0.2.11.dist-info → atlas_ftag_tools-0.2.13.dist-info}/METADATA

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: atlas-ftag-tools
-Version: 0.2.11
+Version: 0.2.13
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT
 Project-URL: Homepage, https://github.com/umami-hep/atlas-ftag-tools/
-Requires-Python: <3.12,>=3.8
+Requires-Python: <3.12,>=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: h5py>=3.0

{atlas_ftag_tools-0.2.11.dist-info → atlas_ftag_tools-0.2.13.dist-info}/RECORD

@@ -1,32 +1,32 @@
-atlas_ftag_tools-0.2.11.dist-info/licenses/LICENSE,sha256=R4o6bZfajQ1KxwcIeavTC00qYTdL33YGNe1hzfV53gM,11349
-ftag/__init__.py,sha256=BGQ1MtuhqCHFXRAh9S9f_ZnOCLWB5RA0ZtL9lW2tofs,748
+atlas_ftag_tools-0.2.13.dist-info/licenses/LICENSE,sha256=R4o6bZfajQ1KxwcIeavTC00qYTdL33YGNe1hzfV53gM,11349
+ftag/__init__.py,sha256=UdYmO_mROM7jvqpPUMbnaQxdCrlR8O0KLlhatyMnapw,748
 ftag/cli_utils.py,sha256=w3TtQmUHSyAKChS3ewvOtcSDAUJAZGIIomaNi8f446U,298
 ftag/cuts.py,sha256=9_ooLZHaO3SnIQBNxwbaPZn-qptGdKnB27FdKQGTiTY,2933
 ftag/flavours.py,sha256=ShH4M2UjQZpZ_NlCctTm2q1tJbzYxjmGteioQ2GcqEU,114
-ftag/flavours.yaml,sha256=CrVTJKndHeL15LT2nkjPodi6Ck9mk_oUtdRby6X_Rcc,9921
+ftag/flavours.yaml,sha256=b86gXX_FMIewLK7_pr0bNgz7RJ84fDf9nAvmjB4J-Ks,9920
 ftag/fraction_optimization.py,sha256=IlMEJe5fD0soX40f-LO4dYAYld2gMqgZRuBLctoPn9A,5566
 ftag/git_check.py,sha256=Y-XqM80CVXZ5ZKrDdZcYOJt3X64uU6W3OP6Z0D7AZU0,1663
-ftag/labeller.py,sha256=IXUgU9UBir39PxVWRKs5r5fqI66Tv0x7nJD3-RYpbrg,2780
+ftag/labeller.py,sha256=6tKLG0SrBijMIZdzWjGdQU9qN_dkRV4eELLX_8YQvTQ,2772
 ftag/labels.py,sha256=2nmcmrZD8mWQPxJsGiOgcLDhSVgWfS_cEzqsBV-Qy8o,4198
-ftag/mock.py,sha256=P2D7nNKAz2jRBbmfpHTDj9sBVU9r7HGd0rpWZOJYZ90,5980
+ftag/mock.py,sha256=syysvzLsBHU8aw7Uy5g3G4HB6LnSmHlGa2BfeXv5mQ4,5970
 ftag/region.py,sha256=ANv0dGI2W6NJqD9fp7EfqAUReH4FOjc1gwl_Qn8llcM,360
 ftag/sample.py,sha256=3N0FrRcu9l1sX8ohuGOHuMYGD0See6gMO4--7NzR2tE,2538
 ftag/track_selector.py,sha256=fJNk_kIBQriBqV4CPT_3ReJbOUnavDDzO-u3EQlRuyk,2654
 ftag/transform.py,sha256=uEGGJSnqoKOzLYQv650XdK_kDNw4Aw-5dc60z9Dp_y0,3963
-ftag/vds.py,sha256=wqj1cA6mIJ4enk8inkearo7ccTw5KCbvuNo2oon51fc,4565
+ftag/vds.py,sha256=l6b54naOK7z0gZjvvtIAQv2Ky4X1w1yLrisZZZYqvbY,11259
 ftag/working_points.py,sha256=RJws2jPMEDQDspCbXUZBifS1CCBmlMJ5ax0eMyDzCRA,15949
 ftag/hdf5/__init__.py,sha256=8yzVQITge-HKkBQQ60eJwWmWDycYZjgVs-qVg4ShVr0,385
 ftag/hdf5/h5add_col.py,sha256=htS5wn4Tm4S3U6mrJ8s24VUnbI7o28Z6Ll-J_V68xTA,12558
 ftag/hdf5/h5move.py,sha256=oYpRu0IDCIJIQ2ML52HBAdoyDxmKkHTeM9JdbPEgKfI,947
-ftag/hdf5/h5reader.py,sha256=i31pDAqmOSaxdeRhc4iSBlld8xJ0pmp4rNd7CugNzw0,13706
+ftag/hdf5/h5reader.py,sha256=NbHohY3RSicM3qnX_0Y1TfGAaDg3wgfEjYlGaWDJmug,14268
 ftag/hdf5/h5split.py,sha256=4Wy6Xc3J58MdD9aBaSZHf5ZcVFnJSkWsm42R5Pgo-R4,2448
-ftag/hdf5/h5utils.py,sha256=-4zKTMtNCrDZr_9Ww7uzfsB7M7muBKpmm_1IkKJnHOI,3222
-ftag/hdf5/h5writer.py,sha256=2gBztierWdwZIqcFItoYz8oua_7hphOI8mbDg7xBdPs,5784
+ftag/hdf5/h5utils.py,sha256=EbCLOF_j1EBFwD95Z3QvlNpBpNkaZoqySZybhVja67U,3542
+ftag/hdf5/h5writer.py,sha256=SMurvZ8FPvqieZUaYRX2SBu-jIyZ6Fx8IasUrEOxIvM,7185
 ftag/utils/__init__.py,sha256=U3YyLY77-FzxRUbudxciieDoy_mnLlY3OfBquA3PnTE,524
 ftag/utils/logging.py,sha256=54NaQiC9Bh4vSznSqzoPfR-7tj1PXfmoH7yKgv_ZHZk,3192
 ftag/utils/metrics.py,sha256=zQI4nPeRDSyzqKpdOPmu0GU560xSWoW1wgL13rrja-I,12664
-atlas_ftag_tools-0.2.11.dist-info/METADATA,sha256=DVmllPN7YQNNmyDcTs3hEGo8mX8ogSReXq9gs6MwUR0,2152
-atlas_ftag_tools-0.2.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-atlas_ftag_tools-0.2.11.dist-info/entry_points.txt,sha256=acr7WwxMIJ3x2I7AheNxNnpWE7sS8XE9MA1eUJGcU5A,169
-atlas_ftag_tools-0.2.11.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
-atlas_ftag_tools-0.2.11.dist-info/RECORD,,
+atlas_ftag_tools-0.2.13.dist-info/METADATA,sha256=ZpQ5GggkLyizsv9uHEOvIlzRqPmC-4tNaoaMgV6unF4,2153
+atlas_ftag_tools-0.2.13.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+atlas_ftag_tools-0.2.13.dist-info/entry_points.txt,sha256=acr7WwxMIJ3x2I7AheNxNnpWE7sS8XE9MA1eUJGcU5A,169
+atlas_ftag_tools-0.2.13.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.2.13.dist-info/RECORD,,

ftag/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.11"
+__version__ = "v0.2.13"
 
 from . import hdf5, utils
 from .cuts import Cuts

ftag/flavours.yaml

@@ -113,7 +113,7 @@
   category: xbb
 - name: qcdnonbb
   label: $\mathrm{QCD} \rightarrow \mathrm{non-} b \bar{b}$
-  cuts: ["R10TruthLabel_R22v1 == 10", "GhostBHadronsFinalCount != 2"]
+  cuts: ["R10TruthLabel_R22v1 == 10", "GhostBHadronsFinalCount < 2"]
   colour: "silver"
   category: xbb
 - name: qcdbx

ftag/hdf5/h5reader.py

@@ -74,10 +74,12 @@ class H5SingleReader:
         num_jets: int | None = None,
         cuts: Cuts | None = None,
         start: int = 0,
+        skip_batches: int = 0,
     ) -> Generator:
         if num_jets is None:
             num_jets = self.num_jets
-
+        if skip_batches > 0:
+            assert not self.shuffle, "Cannot skip batches if shuffle is True"
         if num_jets > self.num_jets:
             log.warning(
                 f"{num_jets:,} jets requested but only {self.num_jets:,} available in {self.fname}."
@@ -97,7 +99,8 @@ class H5SingleReader:
             indices = list(range(start, self.num_jets + start, self.batch_size))
             if self.shuffle:
                 self.rng.shuffle(indices)
-
+            if skip_batches > 0:
+                indices = indices[skip_batches:]
             # loop over batches and read file
             for low in indices:
                 for name in variables:
@@ -176,7 +179,12 @@ class H5Reader:
 
         # calculate batch sizes
         if self.weights is None:
-            self.weights = [1 / len(self.fname)] * len(self.fname)
+            rows_per_file = [
+                H5SingleReader(f, jets_name=self.jets_name).num_jets for f in self.fname
+            ]
+            num_total = sum(rows_per_file)
+            self.weights = [num / num_total for num in rows_per_file]
+
         self.batch_sizes = [int(w * self.batch_size) for w in self.weights]
 
         # create readers
@@ -233,6 +241,7 @@ class H5Reader:
         num_jets: int | None = None,
         cuts: Cuts | None = None,
         start: int = 0,
+        skip_batches: int = 0,
     ) -> Generator:
         """Generate batches of selected jets.
 
@@ -246,6 +255,8 @@ class H5Reader:
             Selection cuts to apply, by default None
         start : int, optional
             Starting index of the first jet to read, by default 0
+        skip_batches : int, optional
+            Number of batches to skip, by default 0
 
         Yields
         ------
@@ -266,7 +277,9 @@ class H5Reader:
 
         # get streams for selected jets from each reader
         streams = [
-            r.stream(variables, int(r.num_jets / self.num_jets * num_jets), cuts, start)
+            r.stream(
+                variables, int(r.num_jets / self.num_jets * num_jets), cuts, start, skip_batches
+            )
             for r in self.readers
         ]
 

ftag/hdf5/h5utils.py

@@ -13,6 +13,7 @@ def get_dtype(
     variables: list[str] | None = None,
     precision: str | None = None,
     transform: Transform | None = None,
+    full_precision_vars: list[str] | None = None,
 ) -> np.dtype:
     """Return a dtype based on an existing dataset and requested variables.
 
@@ -26,6 +27,8 @@ def get_dtype(
         Precision to cast floats to, "half" or "full", by default None
     transform : Transform | None, optional
         Transform to apply to variables names, by default None
+    full_precision_vars : list[str] | None, optional
+        List of variables to keep in full precision, by default None
 
     Returns
     -------
@@ -39,6 +42,8 @@ def get_dtype(
     """
     if variables is None:
         variables = ds.dtype.names
+    if full_precision_vars is None:
+        full_precision_vars = []
 
     if (missing := set(variables) - set(ds.dtype.names)) and transform is not None:
         variables = transform.map_variable_names(ds.name, variables, inverse=True)
@@ -50,7 +55,10 @@ def get_dtype(
 
     dtype = [(n, x) for n, x in ds.dtype.descr if n in variables]
     if precision:
-        dtype = [(n, cast_dtype(x, precision)) for n, x in dtype]
+        dtype = [
+            (n, cast_dtype(x, precision)) if n not in full_precision_vars else (n, x)
+            for n, x in dtype
+        ]
 
     return np.dtype(dtype)
 
@@ -78,6 +86,7 @@ def cast_dtype(typestr: str, precision: str) -> np.dtype:
     t = np.dtype(typestr)
     if t.kind != "f":
         return t
+
     if precision == "half":
         return np.dtype("f2")
     if precision == "full":

ftag/hdf5/h5writer.py

@@ -47,18 +47,28 @@ class H5Writer:
     precision: str = "full"
     full_precision_vars: list[str] | None = None
     shuffle: bool = True
+    num_jets: int | None = None  # Allow dynamic mode by defaulting to None
 
     def __post_init__(self):
         self.num_written = 0
         self.rng = np.random.default_rng(42)
-        self.num_jets = [shape[0] for shape in self.shapes.values()]
-        assert len(set(self.num_jets)) == 1, "Must have same number of jets per group"
-        self.num_jets = self.num_jets[0]
+
+        # Infer number of jets from shapes if not explicitly passed
+        inferred_num_jets = [shape[0] for shape in self.shapes.values()]
+        if self.num_jets is None:
+            assert len(set(inferred_num_jets)) == 1, "Shapes must agree in first dimension"
+            self.fixed_mode = False
+        else:
+            self.fixed_mode = True
+            for name in self.shapes:
+                self.shapes[name] = (self.num_jets,) + self.shapes[name][1:]
 
         if self.precision == "full":
             self.fp_dtype = np.float32
         elif self.precision == "half":
             self.fp_dtype = np.float16
+        elif self.precision is None:
+            self.fp_dtype = None
         else:
             raise ValueError(f"Invalid precision: {self.precision}")
 
@@ -71,16 +81,34 @@ class H5Writer:
             self.create_ds(name, dtype)
 
     @classmethod
-    def from_file(cls, source: Path, num_jets: int | None = None, **kwargs) -> H5Writer:
+    def from_file(
+        cls, source: Path, num_jets: int | None = 0, variables=None, **kwargs
+    ) -> H5Writer:
         with h5py.File(source, "r") as f:
             dtypes = {name: ds.dtype for name, ds in f.items()}
             shapes = {name: ds.shape for name, ds in f.items()}
-            if num_jets is not None:
+
+            if variables:
+                new_dtye = {}
+                new_shape = {}
+                for name, ds in f.items():
+                    if name not in variables:
+                        continue
+                    new_dtye[name] = ftag.hdf5.get_dtype(
+                        ds,
+                        variables=variables[name],
+                        precision=kwargs.get("precision"),
+                        full_precision_vars=kwargs.get("full_precision_vars"),
+                    )
+                    new_shape[name] = ds.shape
+                dtypes = new_dtye
+                shapes = new_shape
+            if num_jets != 0:
                 shapes = {name: (num_jets,) + shape[1:] for name, shape in shapes.items()}
             compression = [ds.compression for ds in f.values()]
             assert len(set(compression)) == 1, "Must have same compression for all groups"
             compression = compression[0]
-            if compression not in kwargs:
+            if "compression" not in kwargs:
                 kwargs["compression"] = compression
         return cls(dtypes=dtypes, shapes=shapes, **kwargs)
 
@@ -88,36 +116,47 @@ class H5Writer:
         if name == self.jets_name and self.add_flavour_label and "flavour_label" not in dtype.names:
             dtype = np.dtype([*dtype.descr, ("flavour_label", "i4")])
 
-        # adjust dtype based on specified precision
-        full_precision_vars = [] if self.full_precision_vars is None else self.full_precision_vars
-        # If the field is in full_precision_vars, use the full precision dtype
+        fp_vars = self.full_precision_vars or []
+        # If no precision is defined, or the field is in full_precision_vars, or its non-float,
+        # keep it at the original dtype
         dtype = np.dtype([
             (
                 field,
-                self.fp_dtype
-                if field not in full_precision_vars and np.issubdtype(dt, np.floating)
-                else dt,
+                (
+                    self.fp_dtype
+                    if (self.fp_dtype and field not in fp_vars and np.issubdtype(dt, np.floating))
+                    else dt
+                ),
             )
             for field, dt in dtype.descr
         ])
 
-        # optimal chunking is around 100 jets, only aply for track groups
         shape = self.shapes[name]
         chunks = (100,) + shape[1:] if shape[1:] else None
 
-        # note: enabling the hd5 shuffle filter doesn't improve write performance
-        self.file.create_dataset(
-            name, dtype=dtype, shape=shape, compression=self.compression, chunks=chunks
-        )
+        if self.fixed_mode:
+            self.file.create_dataset(
+                name, dtype=dtype, shape=shape, compression=self.compression, chunks=chunks
+            )
+        else:
+            maxshape = (None,) + shape[1:]
+            self.file.create_dataset(
+                name,
+                dtype=dtype,
+                shape=(0,) + shape[1:],
+                maxshape=maxshape,
+                compression=self.compression,
+                chunks=chunks,
+            )
 
     def close(self) -> None:
-        with h5py.File(self.dst) as f:
-            written = len(f[self.jets_name])
-        if self.num_written != written:
-            raise ValueError(
-                f"Attemped to close file {self.dst} when only {self.num_written:,} out of"
-                f" {written:,} jets have been written"
-            )
+        if self.fixed_mode:
+            written = len(self.file[self.jets_name])
+            if self.num_written != written:
+                raise ValueError(
+                    f"Attempted to close file {self.dst} when only {self.num_written:,} out of"
+                    f" {written:,} jets have been written"
+                )
         self.file.close()
 
     def get_attr(self, name, group=None):
@@ -137,18 +176,25 @@ class H5Writer:
                 for attr_name, value in ds.attrs.items():
                     self.add_attr(attr_name, value, group=name)
 
-    def write(self, data: dict[str, np.array]) -> None:
-        if (total := self.num_written + len(data[self.jets_name])) > self.num_jets:
-            raise ValueError(
-                f"Attempted to write more jets than expected: {total:,} > {self.num_jets:,}"
-            )
-        idx = np.arange(len(data[self.jets_name]))
+    def write(self, data: dict[str, np.ndarray]) -> None:
+        batch_size = len(data[self.jets_name])
+        idx = np.arange(batch_size)
         if self.shuffle:
             self.rng.shuffle(idx)
             data = {name: array[idx] for name, array in data.items()}
 
         low = self.num_written
-        high = low + len(idx)
+        high = low + batch_size
+
+        if self.fixed_mode and high > self.num_jets:
+            raise ValueError(
+                f"Attempted to write more jets than expected: {high:,} > {self.num_jets:,}"
+            )
+
         for group in self.dtypes:
-            self.file[group][low:high] = data[group]
-        self.num_written += len(idx)
+            ds = self.file[group]
+            if not self.fixed_mode:
+                ds.resize((high,) + ds.shape[1:])
+            ds[low:high] = data[group]
+
+        self.num_written += batch_size

ftag/labeller.py

@@ -30,7 +30,7 @@ class Labeller:
     def __post_init__(self) -> None:
         if isinstance(self.labels, LabelContainer):
             self.labels = list(self.labels)
-        self.labels = sorted([Flavours[label] for label in self.labels])
+        self.labels = [Flavours[label] for label in self.labels]
 
     @property
     def variables(self) -> list[str]:

ftag/mock.py

@@ -106,11 +106,11 @@ def get_mock_scores(labels: np.ndarray, is_xbb: bool = False) -> np.ndarray:
     for i in range(n_classes):
         tmp_means = []
         tmp_means = [
-            0 if j != i else mean_scale_list[np.random.randint(0, len(mean_scale_list))]
+            0 if j != i else mean_scale_list[rng.integers(0, len(mean_scale_list))]
             for j in range(n_classes)
         ]
         means.append(tmp_means)
-        scales.append(mean_scale_list[np.random.randint(0, len(mean_scale_list))])
+        scales.append(mean_scale_list[rng.integers(0, len(mean_scale_list))])
 
     # Map the labels to the means
     label_mapping = dict(zip(label_dict.values(), means))

ftag/vds.py

@@ -8,114 +8,334 @@ import sys
 from pathlib import Path
 
 import h5py
+import numpy as np
 
 
-def parse_args(args):
+def parse_args(args=None):
     parser = argparse.ArgumentParser(
-        description="Create a lightweight wrapper around a set of h5 files"
+        description="Create a lightweight HDF5 wrapper (virtual datasets + "
+        "summed cutBookkeeper counts) around a set of .h5 files"
+    )
+    parser.add_argument(
+        "pattern",
+        type=Path,
+        help="quotes-enclosed glob pattern of files to merge, "
+        "or a regex if --use_regex is given",
     )
-    parser.add_argument("pattern", type=Path, help="quotes-enclosed glob pattern of files to merge")
     parser.add_argument("output", type=Path, help="path to output virtual file")
-    parser.add_argument("--use_regex", help="if provided pattern is a regex", action="store_true")
-    parser.add_argument("--regex_path", type=str, required="--regex" in sys.argv, default=None)
+    parser.add_argument(
+        "--use_regex",
+        action="store_true",
+        help="treat PATTERN as a regular expression instead of a glob",
+    )
+    parser.add_argument(
+        "--regex_path",
+        type=str,
+        required="--use_regex" in (args or sys.argv),
+        default=None,
+        help="directory whose entries the regex is applied to "
+        "(defaults to the current working directory)",
+    )
     return parser.parse_args(args)
 
 
-def get_virtual_layout(fnames: list[str], group: str):
-    # get sources
+def get_virtual_layout(fnames: list[str], group: str) -> h5py.VirtualLayout:
+    """Concatenate group from multiple files into a single VirtualDataset.
+
+    Parameters
+    ----------
+    fnames : list[str]
+        List with the file names
+    group : str
+        Name of the group that is concatenated
+
+    Returns
+    -------
+    h5py.VirtualLayout
+        Virtual layout of the new virtual dataset
+    """
     sources = []
     total = 0
+
+    # Loop over the input files
     for fname in fnames:
-        with h5py.File(fname) as f:
-            vsource = h5py.VirtualSource(f[group])
-            total += vsource.shape[0]
-            sources.append(vsource)
+        with h5py.File(fname, "r") as f:
+            # Get the file and append its length
+            vsrc = h5py.VirtualSource(f[group])
+            total += vsrc.shape[0]
+            sources.append(vsrc)
 
-    # define layout of the vds
-    with h5py.File(fnames[0]) as f:
+    # Define the layout of the output vds
+    with h5py.File(fnames[0], "r") as f:
         dtype = f[group].dtype
         shape = f[group].shape
+
+    # Update the shape finalize the output layout
     shape = (total,) + shape[1:]
     layout = h5py.VirtualLayout(shape=shape, dtype=dtype)
 
-    # fill the vds
+    # Fill the vds
     idx = 0
-    for source in sources:
-        length = source.shape[0]
-        layout[idx : idx + length] = source
+    for vsrc in sources:
+        length = vsrc.shape[0]
+        layout[idx : idx + length] = vsrc
         idx += length
 
     return layout
 
 
-def glob_re(pattern, regex_path):
+def glob_re(pattern: str | None, regex_path: str | None) -> list[str] | None:
+    """Return list of filenames that match REGEX pattern inside regex_path.
+
+    Parameters
+    ----------
+    pattern : str
+        Pattern for the input files
+    regex_path : str
+        Regex path for the input files
+
+    Returns
+    -------
+    list[str]
+        List of the file basenames that matched the regex pattern
+    """
+    if pattern is None or regex_path is None:
+        return None
+
     return list(filter(re.compile(pattern).match, os.listdir(regex_path)))
 
 
-def regex_files_from_dir(reg_matched_fnames, regex_path):
+def regex_files_from_dir(
+    reg_matched_fnames: list[str] | None,
+    regex_path: str | None,
+) -> list[str] | None:
+    """Turn a list of basenames into full paths; dive into sub-dirs if needed.
+
+    Parameters
+    ----------
+    reg_matched_fnames : list[str]
+        List of the regex matched file names
+    regex_path : str
+        Regex path for the input files
+
+    Returns
+    -------
+    list[str]
+        List of file paths (as strings) that matched the regex and any subsequent
+        globbing inside matched directories.
+    """
+    if reg_matched_fnames is None or regex_path is None:
+        return None
+
     parent_dir = regex_path or str(Path.cwd())
-    full_paths = [parent_dir + "/" + fname for fname in reg_matched_fnames]
-    paths_to_glob = [fname + "/*.h5" if Path(fname).is_dir() else fname for fname in full_paths]
-    nested_fnames = [glob.glob(fname) for fname in paths_to_glob]
+    full_paths = [Path(parent_dir) / fname for fname in reg_matched_fnames]
+    paths_to_glob = [str(fp / "*.h5") if fp.is_dir() else str(fp) for fp in full_paths]
+    nested_fnames = [glob.glob(p) for p in paths_to_glob]
     return sum(nested_fnames, [])
 
 
+def sum_counts_once(counts: np.ndarray) -> np.ndarray:
+    """Reduce the arrays in the counts dataset for one file to a scalar via summation.
+
+    Parameters
+    ----------
+    counts : np.ndarray
+        Array from the h5py dataset (counts) from the cutBookkeeper groups
+
+    Returns
+    -------
+    np.ndarray
+        Array with the summed variables for the file
+    """
+    dtype = counts.dtype
+    summed = np.zeros((), dtype=dtype)
+    for field in dtype.names:
+        summed[field] = counts[field].sum()
+    return summed
+
+
+def check_subgroups(fnames: list[str], group_name: str = "cutBookkeeper") -> list[str]:
+    """Check which subgroups are available for the bookkeeper.
+
+    Find the intersection of sub-group names that have a 'counts' dataset
+    in every input file. (Using the intersection makes the script robust
+    even if a few files are missing a variation.)
+
+    Parameters
+    ----------
+    fnames : list[str]
+        List of the input files
+    group_name : str, optional
+        Group name in the h5 files of the bookkeeper, by default "cutBookkeeper"
+
+    Returns
+    -------
+    set[str]
+        Returns the files with common sub-groups
+
+    Raises
+    ------
+    KeyError
+        When a file does not have a bookkeeper
+    ValueError
+        When no common bookkeeper sub-groups were found
+    """
+    common: set[str] | None = None
+    for fname in fnames:
+        with h5py.File(fname, "r") as f:
+            if group_name not in f:
+                raise KeyError(f"{fname} has no '{group_name}' group")
+            these = {
+                name
+                for name, item in f[group_name].items()
+                if isinstance(item, h5py.Group) and "counts" in item
+            }
+            common = these if common is None else common & these
+    if not common:
+        raise ValueError("No common cutBookkeeper sub-groups with 'counts' found")
+    return sorted(common)
+
+
+def aggregate_cutbookkeeper(
+    fnames: list[str],
+    group_name: str = "cutBookkeeper",
+) -> dict[str, np.ndarray] | None:
+    """Aggregate the cutBookkeeper in the input files.
+
+    For every input file:
+    For every sub-group (nominal, sysUp, sysDown, …):
+    1. Sum the 4-entry record array inside each file into 1 record
+    1. Add those records from all files together into grand total
+    Returns a dict  {subgroup_name: scalar-record-array}
+
+    Parameters
+    ----------
+    fnames : list[str]
+        List of the input files
+
+    Returns
+    -------
+    dict[str, np.ndarray] | None
+        Dict with the accumulated cutBookkeeper groups. If the cut bookkeeper
+        is not in the files, return None.
+    """
+    if any(group_name not in h5py.File(f, "r") for f in fnames):
+        return None
+
+    subgroups = check_subgroups(fnames, group_name=group_name)
+
+    # initialise an accumulator per subgroup (dtype taken from 1st file)
+    accum: dict[str, np.ndarray] = {}
+    with h5py.File(fnames[0], "r") as f0:
+        for sg in subgroups:
+            dtype = f0[f"{group_name}/{sg}/counts"].dtype
+            accum[sg] = np.zeros((), dtype=dtype)
+
+    # add each files contribution field-wise
+    for fname in fnames:
+        with h5py.File(fname, "r") as f:
+            for sg in subgroups:
+                per_file = sum_counts_once(f[f"{group_name}/{sg}/counts"][()])
+                for fld in accum[sg].dtype.names:
+                    accum[sg][fld] += per_file[fld]
+
+    return accum
+
+
 def create_virtual_file(
     pattern: Path | str,
-    out_fname: Path | None = None,
+    out_fname: Path | str | None = None,
     use_regex: bool = False,
     regex_path: str | None = None,
     overwrite: bool = False,
-):
-    # get list of filenames
+    bookkeeper_name: str = "cutBookkeeper",
+) -> Path:
+    """Create the virtual dataset file for the given inputs.
+
+    Parameters
+    ----------
+    pattern : Path | str
+        Pattern of the input files used. Wildcard is supported
+    out_fname : Path | str | None, optional
+        Output path to which the virtual dataset file is written. By default None
+    use_regex : bool, optional
+        If you want to use regex instead of glob, by default False
+    regex_path : str | None, optional
+        Regex logic used to define the input files, by default None
+    overwrite : bool, optional
+        Decide, if an existing output file is overwritten, by default False
+    bookkeeper_name : str, optional
+        Name of the cut bookkeeper in the h5 files.
+
+    Returns
+    -------
+    Path
+        Path object of the path to which the output file is written
+
+    Raises
+    ------
+    FileNotFoundError
+        If not input files were found for the given pattern
+    ValueError
+        If no output file is given and the input comes from multiple directories
+    """
+    # Get list of filenames
     pattern_str = str(pattern)
-    if use_regex:
-        reg_matched_fnames = glob_re(pattern_str, regex_path)
-        print("reg matched fnames: ", reg_matched_fnames)
-        fnames = regex_files_from_dir(reg_matched_fnames, regex_path)
+
+    # Use regex to find input files else use glob
+    if use_regex is True:
+        matched = glob_re(pattern_str, regex_path)
+        fnames = regex_files_from_dir(matched, regex_path)
     else:
         fnames = glob.glob(pattern_str)
+
+    # Throw error if no input files were found
     if not fnames:
-        raise FileNotFoundError(f"No files matched pattern {pattern}")
-    print("Files to merge to vds: ", fnames)
+        raise FileNotFoundError(f"No files matched pattern {pattern!r}")
 
-    # infer output path if not given
+    # Infer output path if not given
     if out_fname is None:
-        assert len({Path(fname).parent for fname in fnames}) == 1
+        if len({Path(f).parent for f in fnames}) != 1:
+            raise ValueError("Give --output when files reside in multiple dirs")
         out_fname = Path(fnames[0]).parent / "vds" / "vds.h5"
     else:
         out_fname = Path(out_fname)
 
-    # check if file already exists
+    # If overwrite is not active and a file exists, stop here
     if not overwrite and out_fname.is_file():
         return out_fname
 
-    # identify common groups across all files
+    # Identify common groups across all files
     common_groups: set[str] = set()
     for fname in fnames:
-        with h5py.File(fname) as f:
+        with h5py.File(fname, "r") as f:
             groups = set(f.keys())
-            common_groups = groups if not common_groups else common_groups.intersection(groups)
-
-    if not common_groups:
-        raise ValueError("No common groups found across files")
-
-    # create virtual file
-    out_fname.parent.mkdir(exist_ok=True)
-    with h5py.File(out_fname, "w") as f:
-        for group in common_groups:
-            layout = get_virtual_layout(fnames, group)
-            f.create_virtual_dataset(group, layout)
-            attrs_dict: dict = {}
-            for fname in fnames:
-                with h5py.File(fname) as g:
-                    for name, value in g[group].attrs.items():
-                        if name not in attrs_dict:
-                            attrs_dict[name] = []
-                        attrs_dict[name].append(value)
-            for name, value in attrs_dict.items():
-                if len(value) > 0:
-                    f[group].attrs[name] = value[0]
+            common_groups = groups if not common_groups else common_groups & groups
+
+    # Ditch the bookkeeper. We will process it separately
+    common_groups.discard("cutBookkeeper")
+
+    # Check that the directory of the output file exists
+    out_fname.parent.mkdir(parents=True, exist_ok=True)
+
+    # Build the output file
+    with h5py.File(out_fname, "w") as fout:
+        # Build "standard" groups
+        for gname in sorted(common_groups):
+            layout = get_virtual_layout(fnames, gname)
+            fout.create_virtual_dataset(gname, layout)
+
+            # Copy first-file attributes to VDS root object
+            with h5py.File(fnames[0], "r") as f0:
+                for k, v in f0[gname].attrs.items():
+                    fout[gname].attrs[k] = v
+
+        # Build the cutBookkeeper
+        counts_total = aggregate_cutbookkeeper(fnames=fnames, group_name=bookkeeper_name)
+        if counts_total is not None:
+            for sg, record in counts_total.items():
+                grp = fout.require_group(f"{bookkeeper_name}/{sg}")
+                grp.create_dataset("counts", data=record, shape=(), dtype=record.dtype)
 
     return out_fname
 
@@ -123,19 +343,20 @@ def create_virtual_file(
 def main(args=None) -> None:
     args = parse_args(args)
     matching_mode = "Applying regex to" if args.use_regex else "Globbing"
-    print(f"{matching_mode} {args.pattern}...")
-    create_virtual_file(
-        args.pattern,
-        args.output,
+    print(f"{matching_mode} {args.pattern} ...")
+    out_path = create_virtual_file(
+        pattern=args.pattern,
+        out_fname=args.output,
         use_regex=args.use_regex,
         regex_path=args.regex_path,
         overwrite=True,
     )
-    with h5py.File(args.output) as f:
+
+    with h5py.File(out_path, "r") as f:
         key = next(iter(f.keys()))
-        num = len(f[key])
-    print(f"Virtual dataset '{key}' has {num:,} entries")
-    print(f"Saved virtual file to {args.output.resolve()}")
+        print(f"Virtual dataset '{key}' has {len(f[key]):,} entries")
+
+    print(f"Saved virtual file to {out_path.resolve()}")
 
 
 if __name__ == "__main__":