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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atlas-ftag-tools
-Version: 0.2.11
+Version: 0.2.12
 Summary: ATLAS Flavour Tagging Tools
 Author: Sam Van Stroud, Philipp Gadow
 License: MIT

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

@@ -1,14 +1,14 @@
-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.12.dist-info/licenses/LICENSE,sha256=R4o6bZfajQ1KxwcIeavTC00qYTdL33YGNe1hzfV53gM,11349
+ftag/__init__.py,sha256=CU1RjEu6pHq11LQ2kAy9YDittMHXB51fNWvuy1NFr7o,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
@@ -18,15 +18,15 @@ 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.12.dist-info/METADATA,sha256=bGfabVRARSL6PZTsDqen30IkQVqVGw8Tg9lMCnzY-5w,2152
+atlas_ftag_tools-0.2.12.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+atlas_ftag_tools-0.2.12.dist-info/entry_points.txt,sha256=acr7WwxMIJ3x2I7AheNxNnpWE7sS8XE9MA1eUJGcU5A,169
+atlas_ftag_tools-0.2.12.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.2.12.dist-info/RECORD,,

ftag/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-__version__ = "v0.2.11"
+__version__ = "v0.2.12"
 
 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))