f9columnar 0.2.2__tar.gz → 0.2.4__tar.gz

{f9columnar-0.2.2 → f9columnar-0.2.4}/PKG-INFO

@@ -1,8 +1,7 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: f9columnar
-Version: 0.2.2
+Version: 0.2.4
 Summary: Columnar analysis utils.
-Home-page: https://gitlab.cern.ch/ijs-f9-ljubljana/f9columnar
 License: MIT
 Author: Jan Gavranovic
 Author-email: jan.gavranovic@cern.ch
@@ -43,82 +42,38 @@ Description-Content-Type: text/markdown
 
 # F9 Columnar
 
-### 🚧 **Work in Progress** ⚠️
+A lightweight Python package for batch processing of ROOT and HDF5 event data.
 
-##  Setup
+### Project description
+
+This package is designed for efficient handling of large datasets. Built on PyTorch, Awkward Array, and Uproot, it utilizes PyTorch's DataLoader with an IterableDataset to enable parallel processing. It implements a columnar event loop, returning batches of events in a format compatible with standard PyTorch training loops over multiple epochs.
+
+Optimized for machine learning applications, the package provides `RootDataLoader` and `Hdf5DataLoader` classes for data loading from ROOT and HDF5 files. Additionally, it supports parallel data processing through a modular pipeline of processor classes, allowing users to chain operations for complex computations and histogramming.
 
-### User install
+##  Setup
 
-#### With PyTorch GPU
+### Install with PyTorch GPU
 
 ```shell
 pip install f9columnar[torch]
 ```
 
-#### With PyTorch CPU (recommended)
+### Install with PyTorch CPU (recommended)
 
 ```shell
 pip install f9columnar
 pip install torch --index-url https://download.pytorch.org/whl/cpu
 ```
 
-#### Without PyTorch
+### Install without PyTorch
 
 ```shell
 pip install f9columnar
 ```
 
-### Development install
-
-Use [poetry](https://python-poetry.org/) to install the required packages:
-
-```shell
-poetry config cache-dir $PWD
-poetry config virtualenvs.in-project true
-poetry install -E torch
-```
-
-This environment is duplicated for batch processing on dCache.
-
-## aCT
-
-[ARC](https://doc.vega.izum.si/arc) Control Tower (aCT) is a system for submitting and managing payloads on ARC (and other) Computing Elements. It is used to submit jobs on sites in Slovenia.
-
-### Installation
-
-Install aCT client from the repository with the following command to the virtual environment (or with poetry):
-
-```shell
-pip install "git+https://github.com/ARCControlTower/aCT.git@test#subdirectory=src/act/client/aCT-client"
-```
-
-The command `act` is available in `PATH` as the virtual environment is activated. See the scripts in the [`submit`](f9columnar/submit/) directory for further details.
-
-### Voms proxy setup
-
-Note that it is recommended to be in `/atlas/si` group and make the proxy with it. Active it using (in a separate terminal):
-
-```shell
-setupATLAS
-lsetup emi
-voms-proxy-init --valid 96:0 --voms atlas:/atlas/si
-```
-
-To propagate the proxy to the system use
-
-```shell
-act proxy
-```
-
-At this point you are ready to use aCT.
-
 ## Getting started example
 
-The code is written in a way that it can be used standalone or as a part of an analysis containing data and MC samples. The following example shows the standalone usage.
-
-### Basic example
-
-The main idea is to have a columnar event loop that returns arrays of events. The code and usage is the same as in a standard torch training loop over epochs, but instead of having epochs we iterate over batches of events.
+The following example demonstrates how to load data from multiple ROOT files, apply a simple filter to select two branches, define variables, apply a cut, and create a histogram.
 
 ```python
 from f9columnar.root_dataloader import get_root_dataloader
@@ -144,7 +99,9 @@ for events in root_dataloader:
     # ... do something with the arrays
 ```
 
-Doing calculations on arrays inside of workers can be done using a `Processor`. Many processors can be chained together into a `ProcessorsGraph` (DAG) to perform more complex calculations. Processors are applied to the arrays in the order given by the topological sort of the DAG. Note that each worker runs the same processor graph on batches of array events and returns the result to the event loop when done. So in the above example there would be 12 (`num_workers`) processor graphs running in parallel on small batches of events. An example of calculating tau visible mass and then applying a cut on this variable is shown below.
+Calculations on arrays within worker processes can be performed using a `Processor`. Multiple processors can be linked together in a `ProcessorsGraph`, forming a directed acyclic graph (DAG). These processors are applied to arrays in the sequence determined by the DAG’s topological order.
+
+Each worker executes the same processor graph on batches of event data and returns the results to the event loop once processing is complete. In the example above, 12 (`num_workers`) processor graphs would be running in parallel, each handling small batches of events. Below is an example demonstrating how to calculate the tau visible mass and apply a cut to this variable.
 
 ```python
 from f9columnar.processors import ProcessorsGraph, CheckpointProcessor
@@ -228,8 +185,61 @@ for processed_graph in dataloader:
     # ... do something with the histograms and arrays
 ```
 
+A higher level of abstraction is also possible using the [`ColumnarEventLoop`](f9columnar/run.py) class. See benchmark [examples](benchmark/f9columnar_benchmark.py) for some more details.
+
+## aCT
+
+Basic job submitting to Slovenian grid is also possible using aCT. Currently it only supports Ntuple analysis data format from rucio (for [R21](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/MultiLeptonAnalysis) and [R25](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/EnhancedCPToolkit)).
+
+### Installation
+
+[ARC](https://doc.vega.izum.si/arc) Control Tower (aCT) is a system for submitting and managing payloads on ARC (and other) Computing Elements. It is used to submit jobs on sites in Slovenia. Install aCT client from the repository with the following command to the virtual environment (or with poetry):
+
+```shell
+pip install "git+https://github.com/ARCControlTower/aCT.git@test#subdirectory=src/act/client/aCT-client"
+```
+
+The command `act` is available in `PATH` as the virtual environment is activated. See the scripts in the [`submit`](f9columnar/submit/) directory for further details.
+
+### Voms proxy setup
+
+Note that it is recommended to be in `/atlas/si` group and make the proxy with it. Active it using (in a separate terminal):
+
+```shell
+setupATLAS
+lsetup emi
+voms-proxy-init --valid 96:0 --voms atlas:/atlas/si
+```
+
+To propagate the proxy to the system use
+
+```shell
+act proxy
+```
+
+At this point you are ready to use aCT.
+
+## Examples
+
+- [Mini-analysis](https://gitlab.cern.ch/jgavrano/columnar-seesaw)
+- [Converting ROOT to HDF5](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/SeeSawML/-/blob/main/seesaw/fakes/hdf5_converter.py?ref_type=heads)
+- [HDF5 Dataloader for ML](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/SeeSawML/-/blob/main/seesaw/fakes/hdf5_dataloader.py?ref_type=heads)
+- [Histogramming](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/SeeSawML/-/blob/main/seesaw/fakes/el_fake_hists.py?ref_type=heads)
+
 ## Development
 
+### Development install
+
+Use [poetry](https://python-poetry.org/) to install the required packages:
+
+```shell
+poetry config cache-dir $PWD
+poetry config virtualenvs.in-project true
+poetry install -E torch
+```
+
+Note: this environment should be duplicated for batch processing on dCache.
+
 ### Making a portable venv with conda
 
 Make sure you have [Miniconda](https://docs.anaconda.com/miniconda/) installed:

{f9columnar-0.2.2 → f9columnar-0.2.4}/README.md

@@ -1,81 +1,37 @@
 # F9 Columnar
 
-### 🚧 **Work in Progress** ⚠️
+A lightweight Python package for batch processing of ROOT and HDF5 event data.
 
-##  Setup
+### Project description
+
+This package is designed for efficient handling of large datasets. Built on PyTorch, Awkward Array, and Uproot, it utilizes PyTorch's DataLoader with an IterableDataset to enable parallel processing. It implements a columnar event loop, returning batches of events in a format compatible with standard PyTorch training loops over multiple epochs.
+
+Optimized for machine learning applications, the package provides `RootDataLoader` and `Hdf5DataLoader` classes for data loading from ROOT and HDF5 files. Additionally, it supports parallel data processing through a modular pipeline of processor classes, allowing users to chain operations for complex computations and histogramming.
 
-### User install
+##  Setup
 
-#### With PyTorch GPU
+### Install with PyTorch GPU
 
 ```shell
 pip install f9columnar[torch]
 ```
 
-#### With PyTorch CPU (recommended)
+### Install with PyTorch CPU (recommended)
 
 ```shell
 pip install f9columnar
 pip install torch --index-url https://download.pytorch.org/whl/cpu
 ```
 
-#### Without PyTorch
+### Install without PyTorch
 
 ```shell
 pip install f9columnar
 ```
 
-### Development install
-
-Use [poetry](https://python-poetry.org/) to install the required packages:
-
-```shell
-poetry config cache-dir $PWD
-poetry config virtualenvs.in-project true
-poetry install -E torch
-```
-
-This environment is duplicated for batch processing on dCache.
-
-## aCT
-
-[ARC](https://doc.vega.izum.si/arc) Control Tower (aCT) is a system for submitting and managing payloads on ARC (and other) Computing Elements. It is used to submit jobs on sites in Slovenia.
-
-### Installation
-
-Install aCT client from the repository with the following command to the virtual environment (or with poetry):
-
-```shell
-pip install "git+https://github.com/ARCControlTower/aCT.git@test#subdirectory=src/act/client/aCT-client"
-```
-
-The command `act` is available in `PATH` as the virtual environment is activated. See the scripts in the [`submit`](f9columnar/submit/) directory for further details.
-
-### Voms proxy setup
-
-Note that it is recommended to be in `/atlas/si` group and make the proxy with it. Active it using (in a separate terminal):
-
-```shell
-setupATLAS
-lsetup emi
-voms-proxy-init --valid 96:0 --voms atlas:/atlas/si
-```
-
-To propagate the proxy to the system use
-
-```shell
-act proxy
-```
-
-At this point you are ready to use aCT.
-
 ## Getting started example
 
-The code is written in a way that it can be used standalone or as a part of an analysis containing data and MC samples. The following example shows the standalone usage.
-
-### Basic example
-
-The main idea is to have a columnar event loop that returns arrays of events. The code and usage is the same as in a standard torch training loop over epochs, but instead of having epochs we iterate over batches of events.
+The following example demonstrates how to load data from multiple ROOT files, apply a simple filter to select two branches, define variables, apply a cut, and create a histogram.
 
 ```python
 from f9columnar.root_dataloader import get_root_dataloader
@@ -101,7 +57,9 @@ for events in root_dataloader:
     # ... do something with the arrays
 ```
 
-Doing calculations on arrays inside of workers can be done using a `Processor`. Many processors can be chained together into a `ProcessorsGraph` (DAG) to perform more complex calculations. Processors are applied to the arrays in the order given by the topological sort of the DAG. Note that each worker runs the same processor graph on batches of array events and returns the result to the event loop when done. So in the above example there would be 12 (`num_workers`) processor graphs running in parallel on small batches of events. An example of calculating tau visible mass and then applying a cut on this variable is shown below.
+Calculations on arrays within worker processes can be performed using a `Processor`. Multiple processors can be linked together in a `ProcessorsGraph`, forming a directed acyclic graph (DAG). These processors are applied to arrays in the sequence determined by the DAG’s topological order.
+
+Each worker executes the same processor graph on batches of event data and returns the results to the event loop once processing is complete. In the example above, 12 (`num_workers`) processor graphs would be running in parallel, each handling small batches of events. Below is an example demonstrating how to calculate the tau visible mass and apply a cut to this variable.
 
 ```python
 from f9columnar.processors import ProcessorsGraph, CheckpointProcessor
@@ -185,8 +143,61 @@ for processed_graph in dataloader:
     # ... do something with the histograms and arrays
 ```
 
+A higher level of abstraction is also possible using the [`ColumnarEventLoop`](f9columnar/run.py) class. See benchmark [examples](benchmark/f9columnar_benchmark.py) for some more details.
+
+## aCT
+
+Basic job submitting to Slovenian grid is also possible using aCT. Currently it only supports Ntuple analysis data format from rucio (for [R21](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/MultiLeptonAnalysis) and [R25](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/EnhancedCPToolkit)).
+
+### Installation
+
+[ARC](https://doc.vega.izum.si/arc) Control Tower (aCT) is a system for submitting and managing payloads on ARC (and other) Computing Elements. It is used to submit jobs on sites in Slovenia. Install aCT client from the repository with the following command to the virtual environment (or with poetry):
+
+```shell
+pip install "git+https://github.com/ARCControlTower/aCT.git@test#subdirectory=src/act/client/aCT-client"
+```
+
+The command `act` is available in `PATH` as the virtual environment is activated. See the scripts in the [`submit`](f9columnar/submit/) directory for further details.
+
+### Voms proxy setup
+
+Note that it is recommended to be in `/atlas/si` group and make the proxy with it. Active it using (in a separate terminal):
+
+```shell
+setupATLAS
+lsetup emi
+voms-proxy-init --valid 96:0 --voms atlas:/atlas/si
+```
+
+To propagate the proxy to the system use
+
+```shell
+act proxy
+```
+
+At this point you are ready to use aCT.
+
+## Examples
+
+- [Mini-analysis](https://gitlab.cern.ch/jgavrano/columnar-seesaw)
+- [Converting ROOT to HDF5](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/SeeSawML/-/blob/main/seesaw/fakes/hdf5_converter.py?ref_type=heads)
+- [HDF5 Dataloader for ML](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/SeeSawML/-/blob/main/seesaw/fakes/hdf5_dataloader.py?ref_type=heads)
+- [Histogramming](https://gitlab.cern.ch/atlas-dch-seesaw-analyses/SeeSawML/-/blob/main/seesaw/fakes/el_fake_hists.py?ref_type=heads)
+
 ## Development
 
+### Development install
+
+Use [poetry](https://python-poetry.org/) to install the required packages:
+
+```shell
+poetry config cache-dir $PWD
+poetry config virtualenvs.in-project true
+poetry install -E torch
+```
+
+Note: this environment should be duplicated for batch processing on dCache.
+
 ### Making a portable venv with conda
 
 Make sure you have [Miniconda](https://docs.anaconda.com/miniconda/) installed:

{f9columnar-0.2.2 → f9columnar-0.2.4}/f9columnar/hdf5_dataloader.py

@@ -229,7 +229,7 @@ class Hdf5LoaderIterator:
         name: str,
         iterators_df: pd.DataFrame,
         worker_id: int,
-        processors: list[Callable[[ArrayLike, dict], tuple[ArrayLike, dict]]] | ProcessorsGraph | None,
+        processors: list[Callable[[ArrayLike, dict], tuple[ArrayLike, dict]]] | ProcessorsGraph | None = None,
         hdf5_files_desc_dct: dict[str, dict[str, Any]] | None = None,
     ) -> None:
         self.name = name
@@ -372,13 +372,14 @@ class Hdf5IterableDataset(IterableDataset):
                 logging.warning(f"Dataset {dataset_name} not found in {file}. Skipping!")
                 continue
 
-            piles_metadata[file] = {"piles_lst": None, "piles_shapes": None}
+            if file not in piles_metadata:
+                piles_metadata[file] = {"piles_lst": [], "piles_shapes": []}
 
             piles_lst = self._get_piles_metadata(file)["piles"][dataset_name]
-            piles_metadata[file]["piles_lst"] = piles_lst
+            piles_metadata[file]["piles_lst"] += piles_lst
 
             piles_shapes = self._get_piles_shape(file, piles_lst)
-            piles_metadata[file]["piles_shapes"] = piles_shapes
+            piles_metadata[file]["piles_shapes"] += piles_shapes
 
         for file, metadata in piles_metadata.items():
             for pile, shape in zip(metadata["piles_lst"], metadata["piles_shapes"]):

{f9columnar-0.2.2 → f9columnar-0.2.4}/f9columnar/hdf5_writer.py

@@ -443,9 +443,17 @@ class NtupleHdf5Writer(ArraysHdf5Writer):
         write_mc: bool = True,
         write_data: bool = True,
         name: str = "datasetHDF5Writer",
+        dataset_names: list[str] | None = None,
         **hdf5_kwargs: Any,
     ) -> None:
-        super().__init__(file_path, dataset_names=None, name=name)
+
+        if dataset_names is None:
+            dataset_names = ["mc", "data"]
+
+        if len(dataset_names) != 2:
+            raise ValueError("Dataset names must be a list of two strings!")
+
+        super().__init__(file_path, dataset_names=dataset_names, name=name)
         self.save_node = save_node
         self.chunk_shape = chunk_shape
 
@@ -463,25 +471,26 @@ class NtupleHdf5Writer(ArraysHdf5Writer):
             self.data_column_names = data_column_names
 
         metadata = {}
+        mc_group, data_group = self.dataset_names[0], self.dataset_names[1]
 
         if write_mc:
             self.create_datasets(
-                dataset_names=["mc"],
+                dataset_names=[mc_group],
                 shape=(self.chunk_shape, len(self.mc_column_names)),
                 maxshape=(None, len(self.mc_column_names)),
                 **hdf5_kwargs,
             )
-            metadata["mc_columns"] = self.mc_column_names
+            metadata[f"{mc_group}_columns"] = self.mc_column_names
 
         if write_data:
             self.create_datasets(
                 mode="a" if self.write_mc else "w",
-                dataset_names=["data"],
+                dataset_names=[data_group],
                 shape=(self.chunk_shape, len(self.data_column_names)),
                 maxshape=(None, len(self.data_column_names)),
                 **hdf5_kwargs,
             )
-            metadata["data_columns"] = self.data_column_names
+            metadata[f"{data_group}_columns"] = self.data_column_names
 
         self.add_metadata(metadata)
 
@@ -498,12 +507,12 @@ class NtupleHdf5Writer(ArraysHdf5Writer):
 
         if self.is_data and self.write_data:
             self._current_idx, self._current_shape = self._current_mc_idx, self._current_mc_shape
-            self.write_arrays(arrays, "data", self.data_column_names, self.chunk_shape)
+            self.write_arrays(arrays, self.dataset_names[1], self.data_column_names, self.chunk_shape)
             self._current_mc_idx, self._current_mc_shape = self._current_idx, self._current_shape
 
         if not self.is_data and self.write_mc:
             self._current_idx, self._current_shape = self._current_data_idx, self._current_data_shape
-            self.write_arrays(arrays, "mc", self.mc_column_names, self.chunk_shape)
+            self.write_arrays(arrays, self.dataset_names[0], self.mc_column_names, self.chunk_shape)
             self._current_data_idx, self._current_data_shape = self._current_idx, self._current_shape
 
         return {"processors": processors}
@@ -522,9 +531,19 @@ class NtupleHdf5PileWriter(ArraysHdf5PileWriter):
         write_mc: bool = True,
         write_data: bool = True,
         name: str = "datasetHDF5Writer",
+        dataset_names_map: dict[str, str] | None = None,
         **hdf5_kwargs: Any,
     ) -> None:
-        super().__init__(file_path, n_piles, pile_assignment, dataset_names=["mc", "data"], name=name)
+
+        if dataset_names_map is None:
+            dataset_names = ["mc", "data"]
+        else:
+            dataset_names = [dataset_names_map["mc"], dataset_names_map["data"]]
+
+        if len(dataset_names) != 2:
+            raise ValueError("Dataset names must be a list of two strings!")
+
+        super().__init__(file_path, n_piles, pile_assignment, dataset_names=dataset_names, name=name)
         self.save_node = save_node
         self.chunk_shape = chunk_shape
 
@@ -542,25 +561,26 @@ class NtupleHdf5PileWriter(ArraysHdf5PileWriter):
             self.data_column_names = data_column_names
 
         metadata: dict[str, Any] = {}
+        mc_group, data_group = self.dataset_names[0], self.dataset_names[1]
 
         if write_mc:
             self.create_datasets(
-                dataset_names=self.pile_datasets_dct["mc"],
+                dataset_names=self.pile_datasets_dct[mc_group],
                 shape=(self.chunk_shape, len(self.mc_column_names)),
                 maxshape=(None, len(self.mc_column_names)),
                 **hdf5_kwargs,
             )
-            metadata["mc_columns"] = self.mc_column_names
+            metadata[f"{mc_group}_columns"] = self.mc_column_names
 
         if write_data:
             self.create_datasets(
                 mode="a" if self.write_mc else "w",
-                dataset_names=self.pile_datasets_dct["data"],
+                dataset_names=self.pile_datasets_dct[data_group],
                 shape=(self.chunk_shape, len(self.data_column_names)),
                 maxshape=(None, len(self.data_column_names)),
                 **hdf5_kwargs,
             )
-            metadata["data_columns"] = self.data_column_names
+            metadata[f"{data_group}_columns"] = self.data_column_names
 
         metadata["piles"] = self.pile_datasets_dct
 
@@ -575,10 +595,10 @@ class NtupleHdf5PileWriter(ArraysHdf5PileWriter):
             raise AttributeError("Arrays attribute not found in the processor!")
 
         if self.is_data and self.write_data:
-            self.write_arrays(arrays, "data", self.data_column_names, self.chunk_shape)
+            self.write_arrays(arrays, self.dataset_names[1], self.data_column_names, self.chunk_shape)
 
         if not self.is_data and self.write_mc:
-            self.write_arrays(arrays, "mc", self.mc_column_names, self.chunk_shape)
+            self.write_arrays(arrays, self.dataset_names[0], self.mc_column_names, self.chunk_shape)
 
         return {"processors": processors}
 

{f9columnar-0.2.2 → f9columnar-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "f9columnar"
-version = "0.2.2"
+version = "0.2.4"
 description = "Columnar analysis utils."
 authors = ["Jan Gavranovic <jan.gavranovic@cern.ch>"]
 readme = "README.md"