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

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

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

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

@@ -1,7 +1,7 @@
-atlas_ftag_tools-0.2.13.dist-info/licenses/LICENSE,sha256=R4o6bZfajQ1KxwcIeavTC00qYTdL33YGNe1hzfV53gM,11349
-ftag/__init__.py,sha256=UdYmO_mROM7jvqpPUMbnaQxdCrlR8O0KLlhatyMnapw,748
+atlas_ftag_tools-0.2.14.dist-info/licenses/LICENSE,sha256=R4o6bZfajQ1KxwcIeavTC00qYTdL33YGNe1hzfV53gM,11349
+ftag/__init__.py,sha256=N224CzevGVd0J9m-ak4UuIuCUXc5Cf6Op9O993TqX80,748
 ftag/cli_utils.py,sha256=w3TtQmUHSyAKChS3ewvOtcSDAUJAZGIIomaNi8f446U,298
-ftag/cuts.py,sha256=9_ooLZHaO3SnIQBNxwbaPZn-qptGdKnB27FdKQGTiTY,2933
+ftag/cuts.py,sha256=r1g0vHJtafwEsCPlss685v9a5YDCSSMf-AXsqHrlxV0,3511
 ftag/flavours.py,sha256=ShH4M2UjQZpZ_NlCctTm2q1tJbzYxjmGteioQ2GcqEU,114
 ftag/flavours.yaml,sha256=b86gXX_FMIewLK7_pr0bNgz7RJ84fDf9nAvmjB4J-Ks,9920
 ftag/fraction_optimization.py,sha256=IlMEJe5fD0soX40f-LO4dYAYld2gMqgZRuBLctoPn9A,5566
@@ -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=NbHohY3RSicM3qnX_0Y1TfGAaDg3wgfEjYlGaWDJmug,14268
+ftag/hdf5/h5reader.py,sha256=kDZykPPGS2ABixp3rBPhG-6TTRTN3VKKBKzFV_sv-Eg,18542
 ftag/hdf5/h5split.py,sha256=4Wy6Xc3J58MdD9aBaSZHf5ZcVFnJSkWsm42R5Pgo-R4,2448
 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.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,,
+atlas_ftag_tools-0.2.14.dist-info/METADATA,sha256=ZUQJ5F0_oCFY-MtQ3-jChGevsFcIui11nQYJxF0tWjw,2153
+atlas_ftag_tools-0.2.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+atlas_ftag_tools-0.2.14.dist-info/entry_points.txt,sha256=acr7WwxMIJ3x2I7AheNxNnpWE7sS8XE9MA1eUJGcU5A,169
+atlas_ftag_tools-0.2.14.dist-info/top_level.txt,sha256=qiYQuKcAvMim-31FwkT3MTQu7WQm0s58tPAia5KKWqs,5
+atlas_ftag_tools-0.2.14.dist-info/RECORD,,

ftag/__init__.py

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

ftag/cuts.py

@@ -42,11 +42,28 @@ class Cut:
     @property
     def value(self) -> int | float:
         if isinstance(self._value, str):
+            txt = self._value.strip().lower()
+            if txt == "nan":
+                return float("nan")
+            if txt in {"+inf", "inf"}:
+                return float("inf")
+            if txt == "-inf":
+                return float("-inf")
             return literal_eval(self._value)
         return self._value
 
     def __call__(self, array):
-        return OPERATORS[self.operator](array[self.variable], self.value)
+        col = array[self.variable]
+        val = self.value
+
+        if isinstance(val, float) and np.isnan(val):
+            if self.operator == "!=":
+                return ~np.isnan(col)
+            if self.operator == "==":
+                return np.isnan(col)
+            raise ValueError("'nan' only makes sense with '==' or '!=' operators")
+
+        return OPERATORS[self.operator](col, val)
 
     def __str__(self) -> str:
         return f"{self.variable} {self.operator} {self.value}"

ftag/hdf5/h5reader.py

@@ -68,6 +68,37 @@ class H5SingleReader:
                     )
         return {name: array[keep_idx] for name, array in data.items()}
 
+    def _process_batch(self, data: dict, cuts: Cuts | None = None) -> dict:
+        """Apply cuts and transformations to the batch.
+
+        Parameters
+        ----------
+        data : dict
+            Dictionary of arrays for each group.
+        cuts : Cuts | None, optional
+            Selection cuts to apply, by default None
+
+        Returns
+        -------
+        dict
+            Processed data dictionary with arrays for each group. After applying cuts,
+            (optional) removal of infs, and (optional) transformation.
+        """
+        # apply selections
+        if cuts:
+            idx = cuts(data[self.jets_name]).idx
+            data = {name: array[idx] for name, array in data.items()}
+
+        # check for inf and remove
+        if self.do_remove_inf:
+            data = self.remove_inf(data)
+
+        # apply transform
+        if self.transform:
+            data = self.transform(data)
+
+        return data
+
     def stream(
         self,
         variables: dict | None = None,
@@ -106,18 +137,8 @@ class H5SingleReader:
                 for name in variables:
                     data[name] = self.read_chunk(f[name], arrays[name], low)
 
-                # apply selections
-                if cuts:
-                    idx = cuts(data[self.jets_name]).idx
-                    data = {name: array[idx] for name, array in data.items()}
-
-                # check for inf and remove
-                if self.do_remove_inf:
-                    data = self.remove_inf(data)
-
-                # apply transform
-                if self.transform:
-                    data = self.transform(data)
+                # Apply cuts and transformations
+                data = self._process_batch(data, cuts)
 
                 # check for completion
                 total += len(data[self.jets_name])
@@ -129,6 +150,58 @@ class H5SingleReader:
 
                 yield data
 
+    def get_batch_reader(
+        self,
+        variables: dict | None = None,
+        cuts: Cuts | None = None,
+    ):
+        """Get a function to read batches of selected jets.
+
+        Parameters
+        ----------
+        variables : dict | None, optional
+            Dictionary of variables to for each group, by default use all jet variables.
+        cuts : Cuts | None, optional
+            Selection cuts to apply, by default None
+
+        Returns
+        -------
+        function
+            Function that takes an index and returns a batch of selected jets.
+        """
+        if variables is None:
+            variables = {self.jets_name: None}
+        h5 = h5py.File(self.fname, "r")
+        arrays = {name: self.empty(h5[name], var) for name, var in variables.items()}
+        # nonlocal data
+        data = {name: self.empty(h5[name], var) for name, var in variables.items()}
+
+        def get_batch(idx: int) -> dict | None:
+            """Get a batch of data from the HDF5 file.
+
+            Parameters
+            ----------
+            idx : int
+                Index of the batch to read.
+
+            Returns
+            -------
+            dict | None
+            Dictionary of arrays for each group, or None if no more batches are available.
+            """
+            low = idx * self.batch_size
+            if low >= self.num_jets:
+                return None
+
+            for name in variables:
+                data[name] = self.read_chunk(h5[name], arrays[name], low)
+
+            data_out = {name: array.copy() for name, array in data.items()}
+
+            return self._process_batch(data_out, cuts)
+
+        return get_batch
+
 
 @dataclass
 class H5Reader:
@@ -315,6 +388,62 @@ class H5Reader:
             # yield batch
             yield data
 
+    def get_batch_reader(
+        self, variables: dict | None = None, cuts: Cuts | None = None, shuffle: bool = True
+    ):
+        """Get a function to read batches of selected jets.
+
+        Parameters
+        ----------
+        variables : dict | None, optional
+            Dictionary of variables to for each group, by default use all jet variables.
+        cuts : Cuts | None, optional
+            Selection cuts to apply, by default None
+        shuffle : bool, optional
+            Read batches in a shuffled order, by default True
+
+        Returns
+        -------
+        function
+            Function that takes an index and returns a batch of selected jets.
+        """
+        if variables is None:
+            variables = {self.jets_name: None}
+
+        # create batch readers for each sample
+        batch_readers = [r.get_batch_reader(variables, cuts) for r in self.readers]
+
+        def get_batch(idx: int) -> dict | None:
+            """Get a batch of data from the HDF5 files.
+
+            Parameters
+            ----------
+            idx : int
+                Index of the batch to read.
+
+            Returns
+            -------
+            dict | None
+                Dictionary of arrays for each group, or None if no more batches are available.
+            """
+            assert idx >= 0, "Index must be non-negative"
+            if idx * self.batch_size >= self.num_jets:
+                return None
+            # get a batch from each sample
+            samples = [br(idx) for br in batch_readers]
+            samples = [s for s in samples if s is not None]
+            if len(samples) == 0:
+                return None
+            # combine samples and shuffle
+            data = {name: np.concatenate([s[name] for s in samples]) for name in variables}
+            if shuffle:
+                idx = np.arange(len(data[self.jets_name]))
+                self.rng.shuffle(idx)
+                data = {name: array[idx] for name, array in data.items()}
+            return data
+
+        return get_batch
+
     def load(
         self, variables: dict | None = None, num_jets: int | None = None, cuts: Cuts | None = None
     ) -> dict: