diff --git a/.github/actions/install/action.yml b/.github/actions/install/action.yml index df9348490..a21ccdef9 100644 --- a/.github/actions/install/action.yml +++ b/.github/actions/install/action.yml @@ -68,5 +68,5 @@ runs: echo "Installing graphnet with flags: ${{ steps.flags.outputs.user_flag }} ${{ steps.flags.outputs.editable_flag }}" echo "pip install ${{ steps.flags.outputs.user_flag }} ${{ steps.flags.outputs.editable_flag }} .[torch-${{ steps.flags.outputs.torch_flag }},${{ inputs.extras }}] -f https://data.pyg.org/whl/torch-${{ inputs.torch_version }}+${{ inputs.hardware }}.html" pip install ${{ steps.flags.outputs.user_flag }} ${{ steps.flags.outputs.editable_flag }} .[torch-${{ steps.flags.outputs.torch_flag }},${{ inputs.extras }}] -f https://data.pyg.org/whl/torch-${{ inputs.torch_version }}+${{ inputs.hardware }}.html - + pip install git+https://github.com/thoglu/jammy_flows.git diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e0146e760..c2adcc966 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -32,7 +32,7 @@ jobs: fi build-icetray: - name: Unit tests - IceTray (v1.13.0 - 3.10) + name: Unit tests - IceTray (v1.13.0 - 3.10) needs: [ check-Qlty-credentials ] runs-on: ubuntu-22.04 strategy: @@ -43,7 +43,7 @@ jobs: # GitHub Actions overwrite the docker container entrypoint # inspect --format '{{.Config.Entrypoint}}' image: icecube/icetray:icetray-devel-v1.13.0-ubuntu22.04-2025-02-12 - + steps: - name: Mimmick Docker Entrypoint # The entrypoint of the container sets python paths @@ -62,7 +62,7 @@ jobs: - name: Verify IceCube is importable run: python3 -c "import icecube; print('Import Successful')" - name: Show python version - run: | + run: | python3 --version pip --version pip3 list @@ -122,9 +122,9 @@ jobs: shell: bash run: | source ~/venv/bin/activate - coverage run --source=graphnet -m pytest tests/ --ignore=tests/examples/04_training --ignore=tests/utilities + coverage run --source=graphnet -m pytest tests/ --ignore=tests/examples/04_training --ignore=tests/utilities coverage run -a --source=graphnet -m pytest tests/examples/04_training - coverage run -a --source=graphnet -m pytest tests/utilities + coverage run -a --source=graphnet -m pytest tests/utilities coverage xml -o coverage.xml coverage report -m - name: Work around permission issue @@ -136,7 +136,7 @@ jobs: with: token: ${{secrets.QLTY_TOKEN}} files: coverage.xml - + build-matrix-examples: name: Examples - Ubuntu 22.04 diff --git a/.github/workflows/code-quality.yml b/.github/workflows/code-quality.yml index 04c5230d6..200f74af6 100644 --- a/.github/workflows/code-quality.yml +++ b/.github/workflows/code-quality.yml @@ -23,14 +23,13 @@ jobs: hook: ["black", "flake8", "docformatter", "pydocstyle", "mypy"] steps: - uses: actions/checkout@v4 - - name: Set up Python 3.10 - uses: actions/setup-python@v3 + - name: Set up Python 3.11 + uses: actions/setup-python@v5 with: - python-version: "3.10" + python-version: "3.11" - name: Install package uses: ./.github/actions/install - with: - + - name: Run ${{ matrix.hook }} run: | pip show ${{ matrix.hook }} diff --git a/.github/workflows/docker_action.yml b/.github/workflows/docker_action.yml index 00569ec45..dbe0a2542 100644 --- a/.github/workflows/docker_action.yml +++ b/.github/workflows/docker_action.yml @@ -4,14 +4,14 @@ on: workflow_call: inputs: DOCKERHUB_NAME: { type: string, required: true } - DOCKERFILE_PATH: { type: string, required: true} + DOCKERFILE_PATH: { type: string, required: true} GRAPHNET_VERSION: { type: string, default: "main" } BASE_IMAGE: { type: string, default: "ubuntu:22.04" } TORCH_CHOICE: { type: string, default: "torch-2.6.0" } # torch-2.7.0|2.6.0|2.5.1|no_torch HARDWARE: { type: string, default: "cpu" } # cpu|cu118|cu121|cu124|cu126|cu128 PYTHON_VERSION: { type: string, default: "3.11" } PLATFORMS: { type: string, default: "linux/amd64" } - CONTEXT_PATH: { type: string, default: "." } + CONTEXT_PATH: { type: string, default: "." } secrets: DOCKERHUB_USERNAME: required: true @@ -102,7 +102,7 @@ jobs: uses: docker/build-push-action@v6 with: context: ${{ inputs.CONTEXT_PATH }} - file: ${{ inputs.DOCKERFILE_PATH }} + file: ${{ inputs.DOCKERFILE_PATH }} platforms: ${{ env.PLATFORMS }} push: true tags: ${{ steps.tag.outputs.full }} diff --git a/.github/workflows/paper.yml b/.github/workflows/paper.yml index 4ee36dba8..f8c0828b9 100644 --- a/.github/workflows/paper.yml +++ b/.github/workflows/paper.yml @@ -28,4 +28,4 @@ jobs: # This is the output path where Pandoc will write the compiled # PDF. Note, this should be the same directory as the input # paper.md - path: paper/paper.pdf \ No newline at end of file + path: paper/paper.pdf diff --git a/.gitignore b/.gitignore index d8d702060..8b679248a 100644 --- a/.gitignore +++ b/.gitignore @@ -150,9 +150,9 @@ data/examples/output/ **.pth **.db **.parquet -# Exception to pre-trained folders; here we want all .pth files +# Exception to pre-trained folders; here we want all .pth files !/graphnet/src/graphnet/models/pretrained/**/**/**/**/**.pth -# Exception to geometry tables +# Exception to geometry tables !/data/geometry_tables/**/**.parquet !/data/tests/sqlite/upgrade_genie_step4_140028_000998_first_5_frames/upgrade_genie_step4_140028_000998_first_5_frames.db !/data/tests/parquet/oscNext_genie_level7_v02/merged/** diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index a29be6d3e..450ab6556 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,8 +9,8 @@ repos: - repo: https://github.com/pycqa/flake8 rev: 7.1.1 hooks: - - id: flake8 - language_version: python3 + - id: flake8 + args: ["--config=.flake8", "--show-source", "--statistics"] - repo: https://github.com/pycqa/docformatter rev: v1.7.7 hooks: @@ -26,4 +26,9 @@ repos: hooks: - id: mypy args: [--follow-imports=silent, --disallow-untyped-defs, --disallow-incomplete-defs, --disallow-untyped-calls] - language_version: python3 \ No newline at end of file + language_version: python3 + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.6.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer diff --git a/CITATION.cff b/CITATION.cff index 16b4ebbaa..7ac855912 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -54,4 +54,4 @@ version: 1.0.0 date-released: 2023-05-12 url: "https://github.com/graphnet-team/graphnet" license: Apache-2.0 -type: software \ No newline at end of file +type: software diff --git a/black.toml b/black.toml index 1e39d837a..246a44a58 100644 --- a/black.toml +++ b/black.toml @@ -7,4 +7,4 @@ extend-exclude = ''' versioneer.py | src/graphnet/_version.py ) -''' \ No newline at end of file +''' diff --git a/configs/datasets/training_classification_example_data_sqlite.yml b/configs/datasets/training_classification_example_data_sqlite.yml index 3a13f8749..4defb9372 100644 --- a/configs/datasets/training_classification_example_data_sqlite.yml +++ b/configs/datasets/training_classification_example_data_sqlite.yml @@ -67,4 +67,4 @@ selection: train_nu_e: 60 random events ~ event_no % 5 > 1 & (dummy_pid == 12 | dummy_pid == -12) train_nu_mu: 60 random events ~ event_no % 5 > 1 & (dummy_pid == 14 | dummy_pid == -14) train_nu_tau: 60 random events ~ event_no % 5 > 1 & (dummy_pid == 16 | dummy_pid == -16) - train_mu: 60 random events ~ event_no % 5 > 1 & (dummy_pid == 13 | dummy_pid == -13) \ No newline at end of file + train_mu: 60 random events ~ event_no % 5 > 1 & (dummy_pid == 13 | dummy_pid == -13) diff --git a/configs/datasets/training_example_data_parquet.yml b/configs/datasets/training_example_data_parquet.yml index 67abca0c4..16ad02e46 100644 --- a/configs/datasets/training_example_data_parquet.yml +++ b/configs/datasets/training_example_data_parquet.yml @@ -53,4 +53,4 @@ seed: 21 selection: test: event_no % 5 == 0 validation: event_no % 5 == 1 - train: event_no % 5 > 1 \ No newline at end of file + train: event_no % 5 > 1 diff --git a/configs/datasets/training_example_data_sqlite.yml b/configs/datasets/training_example_data_sqlite.yml index 20c4aa8c0..dd6901123 100644 --- a/configs/datasets/training_example_data_sqlite.yml +++ b/configs/datasets/training_example_data_sqlite.yml @@ -53,4 +53,4 @@ seed: 21 selection: test: event_no % 5 == 0 validation: event_no % 5 == 1 - train: event_no % 5 > 1 \ No newline at end of file + train: event_no % 5 > 1 diff --git a/data/tests/km3net/mcv6.example_offline_reco_format.root b/data/tests/km3net/mcv6.example_offline_reco_format.root new file mode 100644 index 000000000..28c6e3116 Binary files /dev/null and b/data/tests/km3net/mcv6.example_offline_reco_format.root differ diff --git a/docker/gnn-benchmarking/dockerfile b/docker/gnn-benchmarking/dockerfile index 26cd5a101..84a5ab8da 100644 --- a/docker/gnn-benchmarking/dockerfile +++ b/docker/gnn-benchmarking/dockerfile @@ -18,4 +18,4 @@ WORKDIR /root ADD module.py ./module.py ADD apply.py ./apply.py -CMD [ "/bin/bash" ] \ No newline at end of file +CMD [ "/bin/bash" ] diff --git a/dockerfile b/dockerfile index 8749b1695..99dda49e7 100644 --- a/dockerfile +++ b/dockerfile @@ -40,4 +40,4 @@ RUN echo 'PS1="🐳 \[\033[38;2;86;138;242m\]graphnet@\h \[\033[0m\]❯ \[\033[0 RUN echo 'PS2="\[\033[38;5;236m\]❯\[\033[38;5;239m\]❯\[\033[0m\]❯ "' >> ~/.bashrc # Default command for executing container. -CMD [ "/bin/bash" ] \ No newline at end of file +CMD [ "/bin/bash" ] diff --git a/docs/source/installation/install.rst b/docs/source/installation/install.rst index e47f7bab4..69df698fd 100644 --- a/docs/source/installation/install.rst +++ b/docs/source/installation/install.rst @@ -55,6 +55,52 @@ To achieve this, we recommend installing |graphnet|\ GraphNeT into a CVMFS with Once installed, |graphnet|\ GraphNeT is available whenever you open the CVMFS locally. +Installation with km3io (KM3NeT) +----------------------------------------------- + +This installation is only necessary if you want to process KM3NeT/ARCA or KM3NeT/ORCA files. Processing means converting them from a `.root` offline format into a suitable format for training using |graphnet|. If you already have your KM3NeT data in `SQLite` or `parquet` format and only want to train a model or perform inference on this database, this specific installation is not needed. + +Note that this installation will add `km3io` ensuring it is built with a compatible versions. The steps below are provided for a conda environment, with an enviroment created in the same way it is done above in this page, but feel free to choose a different enviroment setup. + +As mentioned, it is highly reommended to create a conda enviroment where your installation is done to do not mess up any dependecy. It can be done with the following commands: + +.. code-block:: bash + + # Create an environment with Python 3.10 + conda create -p --no-default-packages python=3.10 -y + # Activate the environment and move to the graphnet repository you just cloned. If using conda: + conda activate + +The isntallation of GraphNeT is then done by: + +.. code-block:: bash + + git clone https://github.com/graphnet-team/graphnet.git + cd graphnet + +Choose the appropriate requirements file based on your system. Here there is just an example of installation with PyTorch-2.5.1 but check the matrix above for a full idea of all the versions can be installed. + +For CPU-only enviroments: + +.. code-block:: bash + + pip3 install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cpu + pip3 install -e .[torch-25] -f https://data.pyg.org/whl/torch-2.5.1+cpu.html + +For GPU enviroments with, for instance, CUDA 11.8 drivers: + +.. code-block:: bash + + pip3 install torch==2.5.1 torchvision==0.20.1 torchaudio==2.5.1 --index-url https://download.pytorch.org/whl/cu118 + pip3 install -e .[torch-25] -f https://data.pyg.org/whl/torch-2.5.1+cu118.html + +Downgrade setuptools for compatibility between km3io and GraphNeT. + +.. code-block:: bash + + pip3 install --force-reinstall setuptools==70.3.0 + pip3 install km3io==1.2.0 + + .. note:: We recommend installing |graphnet|\ GraphNeT without GPU in clean metaprojects. - diff --git a/examples/04_training/README.md b/examples/04_training/README.md index e94830cdf..7536adb80 100644 --- a/examples/04_training/README.md +++ b/examples/04_training/README.md @@ -30,7 +30,7 @@ This subfolder contains two main training scripts: # Same as above, as this is the default model config. (graphnet) $ python examples/04_training/03_train_dynedge_from_config.py \ --model-config configs/models/example_energy_reconstruction_model.yml - + # Train a vertex position reconstruction model (graphnet) $ python examples/04_training/03_train_dynedge_from_config.py \ --model-config configs/models/example_vertex_position_reconstruction_model.yml diff --git a/examples/05_liquido/README.md b/examples/05_liquido/README.md index daf010218..2e025d6b8 100644 --- a/examples/05_liquido/README.md +++ b/examples/05_liquido/README.md @@ -1,3 +1,3 @@ # LiquidO data conversion example -This directory contain examples of converting data from LiqidO to intermediate data formats suitable for deep learning in GraphNeT +This directory contain examples of converting data from LiqidO to intermediate data formats suitable for deep learning in GraphNeT diff --git a/examples/06_prometheus/README.md b/examples/06_prometheus/README.md index 0a9e5bbd9..4445901e3 100644 --- a/examples/06_prometheus/README.md +++ b/examples/06_prometheus/README.md @@ -1,5 +1,5 @@ # Prometheus data conversion example -This directory contain examples of converting data from the open-source neutrino telescope Prometheus simulation to intermediate data formats suitable for deep learning in GraphNeT. +This directory contain examples of converting data from the open-source neutrino telescope Prometheus simulation to intermediate data formats suitable for deep learning in GraphNeT. More details on the Prometheus simulation project can be found [here](https://arxiv.org/pdf/2304.14526.pdf) diff --git a/examples/07_km3net/01_convert_km3net.py b/examples/07_km3net/01_convert_km3net.py new file mode 100644 index 000000000..a3a53b49d --- /dev/null +++ b/examples/07_km3net/01_convert_km3net.py @@ -0,0 +1,104 @@ +"""Code to run the extraction of km3net data.""" + +import os +import warnings + +from graphnet.constants import EXAMPLE_OUTPUT_DIR, TEST_DATA_DIR +from graphnet.data.readers import KM3NeTReader +from graphnet.data.writers import ParquetWriter, SQLiteWriter +from graphnet.data import DataConverter +from graphnet.data.extractors.km3net import ( + KM3NeTTruthExtractor, + KM3NeTFullPulseExtractor, + KM3NeTTriggPulseExtractor, + KM3NeTHNLTruthExtractor, + KM3NeTRegularRecoExtractor, + KM3NeTHNLRecoExtractor, +) + +from graphnet.utilities.argparse import ArgumentParser + + +def main(backend: str, triggered: str, HNL: str, OUTPUT_DIR: str) -> None: + """Convert ROOT files from KM3NeT to `backend` format.""" + warnings.simplefilter(action="ignore", category=FutureWarning) + + input_dir = [f"{TEST_DATA_DIR}/km3net"] + if OUTPUT_DIR != "None": + outdir = f"{OUTPUT_DIR}/{backend}" + else: + outdir = f"{EXAMPLE_OUTPUT_DIR}/{backend}" + os.makedirs(outdir, exist_ok=True) + print(60 * "*") + print(f"Saving to {outdir}") + print(60 * "*") + if backend == "parquet": + save_method = ParquetWriter(truth_table="truth") + elif backend == "sqlite": + save_method = SQLiteWriter() # type: ignore + else: + raise ValueError("Invalid backend choice") + + if HNL == "km3net-vars": + truth_extractor = KM3NeTTruthExtractor(name="truth") + reco_extractor = KM3NeTRegularRecoExtractor(name="reco") + elif HNL == "hnl-vars": + truth_extractor = KM3NeTHNLTruthExtractor(name="truth") # type: ignore + reco_extractor = KM3NeTHNLRecoExtractor(name="reco") # type: ignore + else: + raise ValueError("Invalid HNL choice") + + if triggered == "Triggered": + pulse_extractor = KM3NeTTriggPulseExtractor(name="pulse_map") + elif triggered == "Snapshot": + pulse_extractor = KM3NeTFullPulseExtractor( + name="pulse_map" + ) # type: ignore + else: + raise ValueError("Invalid triggered choice") + + converter = DataConverter( + file_reader=KM3NeTReader(), + save_method=save_method, + extractors=[truth_extractor, pulse_extractor, reco_extractor], + outdir=outdir, + num_workers=1, + ) + + converter(input_dir=input_dir) + + +if __name__ == "__main__": + + # Parse command-line arguments + parser = ArgumentParser( + description=""" + Convert root files from KM3NeT to an sqlite or parquet. + """ + ) + + parser.add_argument( + "backend", + choices=["sqlite", "parquet"], + help="Choose the backend format", + ) + parser.add_argument( + "triggered", + choices=["Triggered", "Snapshot"], + help="Choose between triggered or snapshot pulse maps", + ) + parser.add_argument( + "HNL", + choices=["km3net-vars", "hnl-vars"], + help="Km3net truth or adding Heavy Neutral Lepton info", + ) + parser.add_argument( + "OUTPUT_DIR", + default="None", + help="Output directory (optional)", + ) + + args, unknown = parser.parse_known_args() + + # Run example script + main(args.backend, args.triggered, args.HNL, args.OUTPUT_DIR) diff --git a/examples/07_km3net/README.md b/examples/07_km3net/README.md new file mode 100644 index 000000000..e7d963ecb --- /dev/null +++ b/examples/07_km3net/README.md @@ -0,0 +1,87 @@ +# KM3NeT Data Conversion + +This folder contains an example script for extracting information from ROOT files of KM3NeT offline data and converting it into intermediate file formats suitable for deep learning training or inference using GraphNeT. Supported output formats include SQLite and Parquet. After this conversion, training and inference on KM3NeT data can be performed efficiently. + +## Example Usage + +The following example demonstrates how to perform the conversion using a sample KM3NeT-like file containing a few events with random information: + +```bash +python 01_convert_km3net.py [OUTPUT_DIR] +``` + +### Arguments: +- ``: Specifies the output format, either `sqlite` or `parquet`. +- ``: Determines whether to extract all pulses (`Snapshot`) or only the triggered ones (`Triggered`). +- ``: Defines the variables to include, such as `km3net-vars` for standard neutrino-related data or `hnl-vars` for additional quantities related to Heavy Neutral Lepton searches. +- `[OUTPUT_DIR]` (optional): Specifies the output directory. If not provided, the output will be stored in GraphNeT's default example output directory, which can be found using: + +```python +from graphnet.constants import EXAMPLE_OUTPUT_DIR +print(EXAMPLE_OUTPUT_DIR) +``` + +The path to the ROOT file converted can be found by running: +```python +from graphnet.constants import TEST_DATA_DIR +print(TEST_DATA_DIR) +``` + +### Output Structure + +The generated SQLite or Parquet file contains: +- A **pulse table**, storing hit-by-hit information for each event, with a unique identifier linking pulses to their respective events. +- A **true Monte Carlo event table**, including ground-truth event information. If available and selected, it may also contain reconstructed information from likelihood-based methods. +- Unavailable variables (e.g., true Monte Carlo information in real data files) will be filled with unphysical placeholder values. + +### Reading the Output Files + +The output files can be read using Python. + +- **If you chose to create a Parquet output**: + You will find several `.parquet` files in the output folder, each corresponding to a different extracted table (e.g., a table with the true event information, a table with pulse information, etc.). + To read one of these tables: + + ```python + import pandas as pd + + df = pd.read_parquet("FILE_NAME.parquet") + print(df.head()) + +- **If you chose to create an SQLite output**: + In this case, you will find a single `.db` file per converted input, which contains all the tables inside. + To list the table names and preview their contents: + + ```python + import pandas as pd + import sqlite3 + + # Connect to the SQLite database + conn = sqlite3.connect("FILE_NAME.db") + cursor = conn.cursor() + + # Get the table names + cursor.execute("SELECT name FROM sqlite_master WHERE type='table';") + tables = [t[0] for t in cursor.fetchall()] + print("The following tables are stored inside the file:", tables) + + # Preview the first 5 rows of each table + for t in tables: + print(f"\nTable: {t}") + df = pd.read_sql_query(f"SELECT * FROM {t[0]} LIMIT 5;", conn) + print(df) + + +## Help + +For more information on available options, use the help flag: + +```bash +python 01_convert_km3net.py -h +``` + +or + +```bash +python 01_convert_km3net.py --help +``` diff --git a/examples/README.md b/examples/README.md index 74c85f3ca..e28f42ae5 100644 --- a/examples/README.md +++ b/examples/README.md @@ -7,7 +7,9 @@ Examples are grouped into five numbered subfolders, roughly in order of how you 2. **Data.** Reading in data in intermediate formats, plotting feature distributions, and converting data between intermediate file formats. These examples are entirely self-contained and can be run by anyone. 3. **Weights.** Fitting per-event weights. 4. **Training.** Training GNN models on various physics tasks. -5**LiquidO.** Converting h5 files from the LiquidO experiment into intermediate formats suitable for deep learning. +5. **LiquidO.** Converting h5 files from the LiquidO experiment into intermediate formats suitable for deep learning. +6. **Prometheus.** Converting parquet files from the Prometheus simulation software into intermediate formats suitable for deep learning. +7. **KM3NeT.** Converting root files from the KM3NeT experiment into intermediate formats suitable for deep learning. Each subfolder contains similarly numbered example scripts. Each example script comes with a simple command-line interface and help functionality, e.g. diff --git a/setup.cfg b/setup.cfg index 476acd91c..83f0219ff 100644 --- a/setup.cfg +++ b/setup.cfg @@ -28,4 +28,5 @@ wrap-summaries = 79 wrap-descriptions = 79 [mypy] -ignore_missing_imports = True \ No newline at end of file +ignore_missing_imports = True +python_version = 3.11 diff --git a/src/graphnet/data/constants.py b/src/graphnet/data/constants.py index 241b4527d..702974da9 100644 --- a/src/graphnet/data/constants.py +++ b/src/graphnet/data/constants.py @@ -52,6 +52,18 @@ class FEATURES: ] KAGGLE = ["x", "y", "z", "time", "charge", "auxiliary"] LIQUIDO = ["sipm_x", "sipm_y", "sipm_z", "t"] + KM3NET = [ + "t", + "pos_x", + "pos_y", + "pos_z", + "dir_x", + "dir_y", + "dir_z", + "tot", + "trig", + ] + KM3NET_HNL = KM3NET class TRUTH: @@ -71,8 +83,6 @@ class TRUTH: "interaction_type", "interaction_time", # Added for vertex reconstruction "inelasticity", - "visible_inelasticity", - "visible_energy", "stopped_muon", ] DEEPCORE = ICECUBE86 @@ -167,3 +177,49 @@ class TRUTH: "energy", "pid", ] + KM3NET = [ + "true_pdgid", + "true_E", + "true_pos_x", + "true_pos_y", + "true_pos_z", + "true_dir_x", + "true_dir_y", + "true_dir_z", + "true_zenith", + "true_azimuth", + "run_id", + "evt_id", + "frame_index", + "trigger_counter", + "n_hits", + "event_no", + "is_cc_flag", + "tau_topology", + ] + KM3NET_HNL = [ + "true_pdgid", + "true_E", + "true_pos_x", + "true_pos_y", + "true_pos_z", + "true_dir_x", + "true_dir_y", + "true_dir_z", + "true_zenith", + "run_id", + "evt_id", + "frame_index", + "trigger_counter", + "n_hits", + "event_no", + "is_cc_flag", + "tau_topology", + "zenith_hnl", + "azimuth_hnl", + "angle_between_showers", + "Energy_hnl", + "Energy_second_shower", + "Energy_imbalance", + "distance", + ] diff --git a/src/graphnet/data/dataconverter.py b/src/graphnet/data/dataconverter.py index 075bf03d3..44f530a38 100644 --- a/src/graphnet/data/dataconverter.py +++ b/src/graphnet/data/dataconverter.py @@ -21,6 +21,7 @@ from .extractors.liquido import H5Extractor from .extractors.internal import ParquetExtractor from .extractors.prometheus import PrometheusExtractor +from .extractors.km3net import KM3NeTExtractor from .dataclasses import I3FileSet @@ -51,6 +52,7 @@ def __init__( List[ParquetExtractor], List[H5Extractor], List[PrometheusExtractor], + List[KM3NeTExtractor], ], index_column: str = "event_no", num_workers: int = 1, diff --git a/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py b/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py index d491d938f..19228d21c 100644 --- a/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py +++ b/src/graphnet/data/extractors/icecube/i3highesteparticleextractor.py @@ -401,7 +401,7 @@ def highest_energy_track( return particle, EonEntrance, distance, visible_length, containment - def highest_energy_starting( + def highest_energy_starting( # noqa: C901 self, frame: "icetray.I3Frame", min_e: float = 0, diff --git a/src/graphnet/data/extractors/km3net/__init__.py b/src/graphnet/data/extractors/km3net/__init__.py new file mode 100644 index 000000000..bf72c954e --- /dev/null +++ b/src/graphnet/data/extractors/km3net/__init__.py @@ -0,0 +1,13 @@ +"""Extractors for extracting pure-python data from KM3NeT-Offline files.""" + +from .km3netextractor import KM3NeTExtractor +from .km3netpulseextractor import ( + KM3NeTTriggPulseExtractor, + KM3NeTFullPulseExtractor, +) +from .km3nettruthextractor import ( + KM3NeTTruthExtractor, + KM3NeTHNLTruthExtractor, + KM3NeTRegularRecoExtractor, + KM3NeTHNLRecoExtractor, +) diff --git a/src/graphnet/data/extractors/km3net/km3netextractor.py b/src/graphnet/data/extractors/km3net/km3netextractor.py new file mode 100644 index 000000000..9242a43f8 --- /dev/null +++ b/src/graphnet/data/extractors/km3net/km3netextractor.py @@ -0,0 +1,30 @@ +"""Base class for all KM3NeTROOT extractors.""" + +from typing import Any +from abc import abstractmethod + +from graphnet.data.extractors import Extractor + +# needs to be implemented at the end. It is a class that will kind of +# gather all the specific extractors for the different data types and +# help to call them all from the reader. Equivalent to the I3extractor (that +# I don't yet understand) in the IceCube example + + +class KM3NeTExtractor(Extractor): + """Base class for all KM3NeT extractors.""" + + def __init__(self, extractor_name: str): + """Initailize KM3NeTTExtractor. + + Args: + extractor_name: Name of the `KM3NeTExtractor` instance. + Used to keep track of the provenance of different data, + and to name tables to which this data is saved. + """ + super().__init__(extractor_name=extractor_name) + + @abstractmethod + def __call__(self, file: Any) -> dict: + """Extract information from file.""" + pass diff --git a/src/graphnet/data/extractors/km3net/km3netpulseextractor.py b/src/graphnet/data/extractors/km3net/km3netpulseextractor.py new file mode 100644 index 000000000..bbd1f896b --- /dev/null +++ b/src/graphnet/data/extractors/km3net/km3netpulseextractor.py @@ -0,0 +1,158 @@ +"""Module for extracting pulse information from a KM3NeT file.""" + +from typing import Any, Dict, TYPE_CHECKING +import numpy as np +import pandas as pd + +from .km3netextractor import KM3NeTExtractor +from graphnet.utilities.imports import has_km3net_package +from graphnet.data.extractors.km3net.utilities.km3net_utilities import ( + create_unique_id_run_by_run, + assert_no_uint_values, + creating_time_zero, +) + +if has_km3net_package() or TYPE_CHECKING: + import awkward as ak + + +class KM3NeTPulseExtractor(KM3NeTExtractor): + """Base class for extracting pulse information from a file.""" + + def __init__(self, name: str, filter_triggered_pulses: bool = False): + """Initialize the base extractor with optional pulse filtering.""" + super().__init__(name) + self.filter_triggered_pulses = filter_triggered_pulses + + def __call__(self, file: Any) -> Dict[str, Any]: + """Extract pulse map information and return a dataframe. + + Args: + file (Any): The file from which to extract pulse map information. + + Returns: + Dict[str, Any]: A dictionary containing pulse map information. + """ + pulsemap_df = self._extract_pulse_map(file) + pulsemap_df = assert_no_uint_values(pulsemap_df) + return pulsemap_df + + def _extract_pulse_map(self, file: Any) -> pd.DataFrame: + """Extract the pulse information and assign unique IDs. + + Args: + file (Any): The file from which to extract pulse information. + + Returns: + pd.DataFrame: A dataframe containing pulse information. + """ + # Process Monte Carlo or events/noise files + unique_id = self._determine_unique_id(file) + + # Extract hits data + hits = file.hits + keys_to_extract = [ + "t", + "pos_x", + "pos_y", + "pos_z", + "dir_x", + "dir_y", + "dir_z", + "tot", + "trig", + ] + hits_array = hits.arrays(keys_to_extract, library="ak") + hits_array["entry"] = ak.local_index(hits_array, axis=0) + df = ak.to_dataframe(hits_array).reset_index(drop=True) + + # Add unique event ID + unique_extended = [ + int(unique_id[index]) for index in df["entry"].values + ] + df["event_no"] = unique_extended + + # Optionally filter triggered pulses + if self.filter_triggered_pulses: + df = df[df["trig"] != 0] + + # Final processing + df = ( + df.drop(["entry", "subentry"], axis=1) + if "subentry" in df.columns + else df.drop(["entry"], axis=1) + ) + df = creating_time_zero(df) + df = df.reset_index(drop=True) + + return df + + def _determine_unique_id(self, file: Any) -> np.ndarray: + """Determine the unique ID for events.""" + if len(file.mc_trks.E[0]) > 0: # Monte Carlo + primaries = file.mc_trks[:, 0] + nus_flavor = [12, 14, 16] + + if (abs(np.array(primaries.pdgid)[0]) not in nus_flavor) and ( + abs(np.array(primaries.pdgid)[0]) > 0.1 + ): # Muon + return create_unique_id_run_by_run( + file_type="muon", + run_id=np.array(file.run_id), + evt_id=np.array(file.id), + hnl_model="none", + ) + elif (5914 in file.mc_trks.pdgid[0]) or ( + file.mc_trks.pdgid[0][0] == 0 + ): # HNL + try: + model_hnl = file.header.model.interaction + except Exception: + model_hnl = "none" + return create_unique_id_run_by_run( + file_type="hnl", + run_id=np.array(file.run_id), + evt_id=np.array(file.id), + hnl_model=model_hnl, + ) + elif (abs(np.array(primaries.pdgid)[0]) in nus_flavor) and ( + 5914 not in file.mc_trks.pdgid[0] + ): # Neutrino + return create_unique_id_run_by_run( + file_type="neutrino", + run_id=np.array(file.run_id), + evt_id=np.array(file.id), + hnl_model="none", + ) + elif len(file.mc_trks.E[0]) == 0: # Events or Noise + if file.header["calibration"] == "dynamical": # Data + return create_unique_id_run_by_run( + file_type="data", + run_id=np.array(file.run_id), + evt_id=np.array(file.id), + hnl_model="none", + ) + elif file.header["calibration"] == "statical": # Noise + return create_unique_id_run_by_run( + file_type="noise", + run_id=np.array(file.run_id), + evt_id=np.array(file.id), + hnl_model="none", + ) + raise ValueError("File type not recognized or corrupted.") + + +class KM3NeTTriggPulseExtractor(KM3NeTPulseExtractor): + """Extractor for triggered pulses.""" + + def __init__(self, name: str = "trigg_pulse_map"): + """Initialize the extractor.""" + super().__init__(name, filter_triggered_pulses=True) + + +class KM3NeTFullPulseExtractor(KM3NeTPulseExtractor): + """Extractor for all pulses (no filtering).""" + + def __init__(self, name: str = "full_pulse_map"): + """Initialize the extractor.""" + super().__init__(name, filter_triggered_pulses=False) diff --git a/src/graphnet/data/extractors/km3net/km3nettruthextractor.py b/src/graphnet/data/extractors/km3net/km3nettruthextractor.py new file mode 100644 index 000000000..06b12b8ec --- /dev/null +++ b/src/graphnet/data/extractors/km3net/km3nettruthextractor.py @@ -0,0 +1,487 @@ +"""Code to extract the truth event information from the KM3NeT ROOT file.""" + +from typing import Any, Dict, TYPE_CHECKING +import numpy as np +import pandas as pd +from graphnet.utilities.imports import has_km3net_package + +from .km3netextractor import KM3NeTExtractor +from graphnet.data.extractors.km3net.utilities.km3net_utilities import ( + create_unique_id_run_by_run, + xyz_dir_to_zen_az, + assert_no_uint_values, + filter_None_NaN, +) + +if has_km3net_package() or TYPE_CHECKING: + import km3io as ki + + +class KM3NeTTruthExtractor(KM3NeTExtractor): + """Class for extracting the truth information from a file.""" + + def __init__( + self, + name: str, + add_hnl_info: bool = False, + add_reco_info: bool = False, + ): + """Initialize the class to extract the truth information.""" + super().__init__(name) + self.add_hnl_info = add_hnl_info + self.add_reco_info = add_reco_info + + def __call__(self, file: Any) -> Dict[str, Any]: + """Extract truth event information as a dataframe.""" + truth_df = self._extract_truth_dataframe(file) + truth_df = assert_no_uint_values(truth_df) # asserts the data format + + return truth_df + + def _extract_particle_attributes( + self, + particle: Any, + attributes: Any, + padding_value: int, + prefix: str, + file_type: str, + ) -> Dict[str, Any]: + if (file_type != "data") and (file_type != "noise"): + dict_particle = {} + for attr in attributes: + dict_particle[prefix + "_" + attr] = filter_None_NaN( + getattr(particle, attr), padding_value + ) + + # also add the zenith and the azimuth computed from the direction + if ( + "dir_x" in attributes + and "dir_y" in attributes + and "dir_z" in attributes + ): + zen, az = xyz_dir_to_zen_az( + filter_None_NaN(getattr(particle, "dir_x"), padding_value), + filter_None_NaN(getattr(particle, "dir_y"), padding_value), + filter_None_NaN(getattr(particle, "dir_z"), padding_value), + padding_value, + ) + dict_particle[prefix + "_zenith"] = zen + dict_particle[prefix + "_azimuth"] = az + elif ( + (file_type != "hnl") + and (file_type != "neutrino") + and (file_type != "muon") + ): + dict_particle = {} + for attr in attributes: + dict_particle[prefix + "_" + attr] = padding_value * np.ones( + len(particle.pos_x) + ) + + # also add the zenith and the azimuth computed from the direction + if ( + "dir_x" in attributes + and "dir_y" in attributes + and "dir_z" in attributes + ): + dict_particle[prefix + "_zenith"] = padding_value * np.ones( + len(particle.pos_x) + ) + dict_particle[prefix + "_azimuth"] = padding_value * np.ones( + len(particle.pos_x) + ) + else: + raise ValueError("File type not recognized") + + return dict_particle + + def _extract_event_attributes( + self, + file: Any, + primaries: Any, + primaries_jshower: Any, + padding_value: int, + file_type: str, + ) -> Dict[str, Any]: + """Return dictionary with event information.""" + evt_id, run_id, frame_index, trigger_counter = ( + np.array(file.id), + np.array(file.run_id), + np.array(file.frame_index), + np.array(file.trigger_counter), + ) + + n_hits = np.array(file.n_hits) + + if file_type == "data": + is_cc_flag = padding_value * np.ones(len(primaries_jshower.E)) + tau_topologies = padding_value * np.ones(len(primaries_jshower.E)) + elif file_type == "noise": + is_cc_flag = padding_value * np.ones(len(primaries_jshower.E)) + tau_topologies = padding_value * np.ones(len(primaries_jshower.E)) + elif file_type == "hnl": + is_cc_flag = np.ones(len(primaries.pos_x)) + tau_topologies = padding_value * np.ones(len(primaries.pos_x)) + elif file_type == "neutrino": + is_cc_flag = np.array(file.w2list[:, 10] == 2) + tau_topologies = [ + ( + 2 + if 16 in np.abs(primaries.pdgid) + and 13 in np.abs(file.mc_trks.pdgid[i]) + else 1 if 16 in np.abs(primaries.pdgid) else 3 + ) + for i in range(len(primaries.pdgid)) + ] + elif file_type == "muon": + is_cc_flag = padding_value * np.ones(len(primaries.pos_x)) + tau_topologies = padding_value * np.ones(len(primaries.pos_x)) + if file_type == "hnl": + try: + model_hnl = file.header.model.interaction + except Exception: + model_hnl = "none" + elif ( + (file_type == "neutrino") + or (file_type == "muon") + or (file_type == "data") + or (file_type == "noise") + ): + model_hnl = "none" + else: + raise ValueError("File type not recognized") + + unique_id = create_unique_id_run_by_run( + file_type=file_type, + run_id=np.array(file.run_id), + evt_id=np.array(file.id), + hnl_model=model_hnl, + ) + + return { + "run_id": run_id, + "evt_id": evt_id, + "frame_index": frame_index, + "trigger_counter": trigger_counter, + "n_hits": n_hits, + "event_no": np.array(unique_id).astype(int), + "is_cc_flag": is_cc_flag, + "tau_topology": tau_topologies, + } + + def _extract_hnl_attributes( + self, + file: Any, + primaries: Any, + primaries_jshower: Any, + padding_value: int, + file_type: str, + ) -> Dict[str, Any]: + """Return dictionary with specific information of the HNL event.""" + if (file_type == "data") or (file_type == "noise"): + return { + "zenith_hnl": padding_value + * np.ones(len(primaries_jshower.E)), + "azimuth_hnl": padding_value + * np.ones(len(primaries_jshower.E)), + "angle_between_showers": padding_value + * np.ones(len(primaries_jshower.E)), + "Energy_hnl": padding_value + * np.ones(len(primaries_jshower.E)), + "Energy_second_shower": padding_value + * np.ones(len(primaries_jshower.E)), + "Energy_imbalance": padding_value + * np.ones(len(primaries_jshower.E)), + "distance": padding_value * np.ones(len(primaries_jshower.E)), + "is_hnl": padding_value * np.ones(len(primaries_jshower.E)), + } + elif file_type == "hnl": + if len(file.mc_trks.E[0]) > 3.5: + hnl = file.mc_trks[:, 1] + first_shower = file.mc_trks[:, 2] + second_shower = file.mc_trks[:, 4] + elif len(file.mc_trks.E[0]) < 3.5: + hnl = file.mc_trks[:, 0] + first_shower = file.mc_trks[:, 1] + second_shower = file.mc_trks[:, 2] + + energy_imbalance = ( + np.array(first_shower.E) - np.array(second_shower.E) + ) / (np.array(first_shower.E) + np.array(second_shower.E)) + distance = np.sqrt( + (np.array(primaries.pos_x) - np.array(second_shower.pos_x)) + ** 2 + + (np.array(primaries.pos_y) - np.array(second_shower.pos_y)) + ** 2 + + (np.array(primaries.pos_z) - np.array(second_shower.pos_z)) + ** 2 + ) + + # compute the angle between the two showers + angle = np.arccos( + np.array(first_shower.dir_x) * np.array(second_shower.dir_x) + + np.array(first_shower.dir_y) * np.array(second_shower.dir_y) + + np.array(first_shower.dir_z) * np.array(second_shower.dir_z) + ) + + # the zenith and azimuth of the heavy neutrino + zen_hnl, az_hnl = xyz_dir_to_zen_az( + np.array(hnl.dir_x), + np.array(hnl.dir_y), + np.array(hnl.dir_z), + padding_value, + ) + + return { + "zenith_hnl": zen_hnl, + "azimuth_hnl": az_hnl, + "angle_between_showers": angle, + "Energy_hnl": np.array(hnl.E), + "Energy_second_shower": np.array(second_shower.E), + "Energy_imbalance": energy_imbalance, + "distance": distance, + "is_hnl": np.ones(len(primaries.pos_x)), + } + elif (file_type == "neutrino") or (file_type == "muon"): + return { + "zenith_hnl": padding_value * np.ones(len(primaries.pos_x)), + "azimuth_hnl": padding_value * np.ones(len(primaries.pos_x)), + "angle_between_showers": np.zeros(len(primaries.pos_x)), + "Energy_hnl": padding_value * np.ones(len(primaries.pos_x)), + "Energy_second_shower": padding_value + * np.ones(len(primaries.pos_x)), + "Energy_imbalance": padding_value + * np.ones(len(primaries.pos_x)), + "distance": np.zeros(len(primaries.pos_x)), + "is_hnl": np.zeros(len(primaries.pos_x)), + } + else: + raise ValueError("File type not recognized") + + def _construct_truth_dictionary( + self, + primaries: Any, + primaries_jshower: Any, + primaries_jmuon: Any, + file: Any, + padding_value: int, + file_type: str, + ) -> Dict[str, Any]: + true_attrs = [ + "pdgid", + "E", + "pos_x", + "pos_y", + "pos_z", + "dir_x", + "dir_y", + "dir_z", + ] + dict_truth = self._extract_particle_attributes( + primaries, + true_attrs, + padding_value, + prefix="true", + file_type=file_type, + ) + + primaries_jshower, primaries_jmuon = ki.tools.best_jshower( + file.trks + ), ki.tools.best_jmuon(file.trks) + jshower_attrs = [ + "E", + "pos_x", + "pos_y", + "pos_z", + "dir_x", + "dir_y", + "dir_z", + ] + jmuon_attrs = [ + "E", + "pos_x", + "pos_y", + "pos_z", + "dir_x", + "dir_y", + "dir_z", + ] + if self.add_reco_info: + jshower_data = self._extract_particle_attributes( + primaries_jshower, + jshower_attrs, + padding_value, + prefix="jshower", + file_type=file_type, + ) + jmuon_data = self._extract_particle_attributes( + primaries_jmuon, + jmuon_attrs, + padding_value, + prefix="jmuon", + file_type=file_type, + ) + dict_truth.update(jshower_data) + dict_truth.update(jmuon_data) + + evt_data = self._extract_event_attributes( + file, + primaries, + primaries_jshower, + padding_value, + file_type=file_type, + ) + dict_truth.update(evt_data) + + if self.add_hnl_info: + hnl_data = self._extract_hnl_attributes( + file, + primaries, + primaries_jshower, + padding_value, + file_type=file_type, + ) + dict_truth.update(hnl_data) + + return dict_truth + + def _extract_truth_dataframe(self, file: Any) -> Any: + """Extract truth information from a file and returns a dataframe. + + Args: + file (Any): The file from which to extract truth information. + + Returns: + pd.DataFrame: A dataframe containing truth information. + """ + nus_flavor = [12, 14, 16] + padding_value = int(99999999) + if len(file.mc_trks.E[0]) > 0: + primaries = file.mc_trks[:, 0] + ############################################################## + # MUON-FILE#################################################### + ############################################################## + if (abs(np.array(primaries.pdgid)[0]) not in nus_flavor) and ( + abs(np.array(primaries.pdgid)[0]) > 0.1 + ): + primaries_jshower = ki.tools.best_jshower(file.trks) + primaries_jmuon = ki.tools.best_jmuon(file.trks) + dict_truth = self._construct_truth_dictionary( + primaries, + primaries_jshower, + primaries_jmuon, + file, + padding_value, + file_type="muon", + ) + + ############################################################### + # HNL-FILE###################################################### + ############################################################### + elif (5914 in file.mc_trks.pdgid[0]) or ( + file.mc_trks.pdgid[0][0] == 0 + ): + primaries_jshower = ki.tools.best_jshower(file.trks) + primaries_jmuon = ki.tools.best_jmuon(file.trks) + dict_truth = self._construct_truth_dictionary( + primaries, + primaries_jshower, + primaries_jmuon, + file, + padding_value, + file_type="hnl", + ) + + elif (abs(np.array(primaries.pdgid)[0]) in nus_flavor) and ( + 5914 not in file.mc_trks.pdgid[0] + ): + #################################################### + # NEUTRINO-FILE###################################### + #################################################### + primaries_jshower = ki.tools.best_jshower(file.trks) + primaries_jmuon = ki.tools.best_jmuon(file.trks) + dict_truth = self._construct_truth_dictionary( + primaries, + primaries_jshower, + primaries_jmuon, + file, + padding_value, + file_type="neutrino", + ) + else: + ValueError("Not a neutrino, muon, hnl, noise or data file.") + elif len(file.mc_trks.E[0]) == 0: + if file.header["calibration"] == "dynamical": + #################################################### + # DATA-FILE########################################## + #################################################### + primaries_jshower = ki.tools.best_jshower(file.trks) + primaries_jmuon = ki.tools.best_jmuon(file.trks) + dict_truth = self._construct_truth_dictionary( + primaries_jshower, + primaries_jshower, + primaries_jmuon, + file, + padding_value, + file_type="data", + ) + + elif file.header["calibration"] == "statical": + #################################################### + # NOISE-FILE########################################## + #################################################### + primaries_jshower = ki.tools.best_jshower(file.trks) + primaries_jmuon = ki.tools.best_jmuon(file.trks) + dict_truth = self._construct_truth_dictionary( + primaries_jshower, + primaries_jshower, + primaries_jmuon, + file, + padding_value, + file_type="noise", + ) + else: + ValueError("Not a neutrino, muon, hnl, noise or data file.") + else: + ValueError("Not a neutrino, muon, hnl, noise or data file.") + + # Ensure all arrays in dict_truth are 1-dimensional + for key, value in dict_truth.items(): + if isinstance(value, np.ndarray) and value.ndim > 1: + print(key) + dict_truth[key] = value.flatten() + truth_df = pd.DataFrame(dict_truth) + + return truth_df + + +class KM3NeTRegularTruthExtractor(KM3NeTTruthExtractor): + """Class for extracting the truth regular true information from file.""" + + def __init__(self, name: str = "truth"): + """Initialize the class to extract the truth regular information.""" + super().__init__(name, add_hnl_info=False) + + +class KM3NeTHNLTruthExtractor(KM3NeTTruthExtractor): + """Class for extracting the truth hnl true information from file.""" + + def __init__(self, name: str = "truth_hnl"): + """Initialize the class to extract the truth hnl information.""" + super().__init__(name, add_hnl_info=True) + + +class KM3NeTRegularRecoExtractor(KM3NeTTruthExtractor): + """Class for extracting the truth regular reco information from file.""" + + def __init__(self, name: str = "reco"): + """Initialize the class.""" + super().__init__(name, add_hnl_info=False, add_reco_info=True) + + +class KM3NeTHNLRecoExtractor(KM3NeTTruthExtractor): + """Class for extracting the truth hnl reco information from file.""" + + def __init__(self, name: str = "reco_hnl"): + """Initialize the class to extract the truth hnl reco information.""" + super().__init__(name, add_hnl_info=True, add_reco_info=True) diff --git a/src/graphnet/data/extractors/km3net/utilities/__init__.py b/src/graphnet/data/extractors/km3net/utilities/__init__.py new file mode 100644 index 000000000..12ef09032 --- /dev/null +++ b/src/graphnet/data/extractors/km3net/utilities/__init__.py @@ -0,0 +1 @@ +"""Code to initialize the utilities for KM3neT data extractors.""" diff --git a/src/graphnet/data/extractors/km3net/utilities/km3net_utilities.py b/src/graphnet/data/extractors/km3net/utilities/km3net_utilities.py new file mode 100644 index 000000000..f378b2853 --- /dev/null +++ b/src/graphnet/data/extractors/km3net/utilities/km3net_utilities.py @@ -0,0 +1,104 @@ +"""Code with some functionalities for the extraction.""" + +from typing import List, Tuple, Union + +import numpy as np +import pandas as pd + + +def create_unique_id_run_by_run( + file_type: str, + run_id: List[int], + evt_id: List[int], + hnl_model: str, +) -> List[int]: + """Create a unique ID for each event based on its parameters. + + Args: + file_type (str): 'neutrino', 'muon', 'noise', 'data', or 'hnl'. + run_id (List[int]): List of run IDs for the events. + evt_id (List[int]): List of event IDs within each run. + + Returns: + List[str]: A list of unique IDs for each event, formatted as strings. + """ + file_type_dict = { + "neutrino": 1, + "muon": 2, + "noise": 3, + "data": 4, + "hnl": 5, + } + + hnl_type_dict = { + "none": 0, # possibility of adding hnl models to break run-filetype degeneracy + } + unique_id = [] + for i in range(len(run_id)): + unique_id.append( + int( + str(run_id[i]) + + str(evt_id[i]) + + str(file_type_dict[file_type]) + + str(hnl_type_dict[hnl_model]) + ) + ) + + return unique_id + + +def filter_None_NaN( + values: Union[List[float], np.ndarray], + padding_value: float, +) -> np.ndarray: + """Remove None and NaN, transforming them to padding value.""" + values = [padding_value if v is None else v for v in values] + values = np.array(values, dtype=float) + values[np.isnan(values)] = padding_value + return values + + +def xyz_dir_to_zen_az( + dir_x: List[float], + dir_y: List[float], + dir_z: List[float], + padding_value: float, +) -> Tuple[List[float], List[float]]: + """Convert direction vector to zenith and azimuth angles.""" + # Compute zenith angle (elevation angle) + with np.errstate(invalid="ignore"): + zenith = np.arccos(dir_z) # zenith angle in radians + + # Compute azimuth angle + azimuth = np.arctan2(dir_y, dir_x) # azimuth angle in radians + az_centered = azimuth + np.pi * np.ones( + len(azimuth) + ) # Center the azimuth angle around zero + # check for NaN in the zenith and replace with padding_value + zenith[np.isnan(zenith)] = padding_value + # change the azimuth values to padding value if the zenith is padding value + az_centered[zenith == padding_value] = padding_value + + return zenith, az_centered + + +def creating_time_zero(df: pd.DataFrame) -> pd.DataFrame: + """Shift the event time so that the first hit has zero in time.""" + df = df.sort_values(by=["event_no", "t"]) + df["min_t"] = df.groupby("event_no")["t"].transform("min") + df["t"] = df["t"] - df["min_t"] + df = df.drop(["min_t"], axis=1) + + return df + + +def assert_no_uint_values(df: pd.DataFrame) -> pd.DataFrame: + """Assert no format no supported by sqlite is in the data.""" + for column in df.columns: + if df[column].dtype == "uint32": + df[column] = df[column].astype("int32") + elif df[column].dtype == "uint64": + df[column] = df[column].astype("int64") + else: + pass + return df diff --git a/src/graphnet/data/readers/__init__.py b/src/graphnet/data/readers/__init__.py index 911a00255..3252983db 100644 --- a/src/graphnet/data/readers/__init__.py +++ b/src/graphnet/data/readers/__init__.py @@ -5,3 +5,4 @@ from .internal_parquet_reader import ParquetReader from .liquido_reader import LiquidOReader from .prometheus_reader import PrometheusReader +from .km3netreader import KM3NeTReader diff --git a/src/graphnet/data/readers/graphnet_file_reader.py b/src/graphnet/data/readers/graphnet_file_reader.py index 6d145f1c0..00b5e004a 100644 --- a/src/graphnet/data/readers/graphnet_file_reader.py +++ b/src/graphnet/data/readers/graphnet_file_reader.py @@ -18,6 +18,7 @@ from graphnet.data.extractors.internal import ParquetExtractor from graphnet.data.extractors.liquido import H5Extractor from graphnet.data.extractors.prometheus import PrometheusExtractor +from graphnet.data.extractors.km3net import KM3NeTExtractor class GraphNeTFileReader(Logger, ABC): @@ -100,6 +101,7 @@ def set_extractors( List[ParquetExtractor], List[H5Extractor], List[PrometheusExtractor], + List[KM3NeTExtractor], ], ) -> None: """Set `Extractor`(s) as member variable. @@ -121,6 +123,7 @@ def _validate_extractors( List[ParquetExtractor], List[H5Extractor], List[PrometheusExtractor], + List[KM3NeTExtractor], ], ) -> None: for extractor in extractors: diff --git a/src/graphnet/data/readers/km3netreader.py b/src/graphnet/data/readers/km3netreader.py new file mode 100644 index 000000000..faaea9265 --- /dev/null +++ b/src/graphnet/data/readers/km3netreader.py @@ -0,0 +1,71 @@ +"""Module for reading KM3NeT files.""" + +from typing import TYPE_CHECKING, Union, List, Any, Dict + + +from graphnet.data.readers import GraphNeTFileReader +from graphnet.utilities.imports import has_km3net_package +from graphnet.data.extractors.km3net import ( + KM3NeTTruthExtractor, + KM3NeTFullPulseExtractor, + KM3NeTTriggPulseExtractor, + KM3NeTHNLTruthExtractor, + KM3NeTRegularRecoExtractor, + KM3NeTHNLRecoExtractor, +) + + +# km3net specific imports +if has_km3net_package() or TYPE_CHECKING: + import km3io as ki # pyright: reportMissingImports=false + + +class KM3NeTReader(GraphNeTFileReader): + """Class for reading KM3NeT files.""" + + _accepted_file_extensions = [".root"] + _accepted_extractors = [ + KM3NeTTruthExtractor, + KM3NeTFullPulseExtractor, + KM3NeTTriggPulseExtractor, + KM3NeTHNLTruthExtractor, + KM3NeTRegularRecoExtractor, + KM3NeTHNLRecoExtractor, + ] + + def __call__( + self, file_path: Union[str] + ) -> Dict[str, Union[Dict[Any, Any], Any]]: + """Open and apply extractors to a single root file. + + Args: + file_path: The path to the file to be read. + + Returns: + data in a list of ordered dataframes with a unique ID. + """ + file = ki.OfflineReader(file_path) + if len(file.trks) > 0: + data = {} + for extractor in self._extractors: + data[extractor._extractor_name] = extractor( + file + ) # extractor returns dataframe + + return data + else: + print(f"File {file_path} has no events.") + return {} + + def find_files( + self, path: Union[str, List[str]] + ) -> Union[List[str], List[Any]]: + """Find all files in a directory with the correct extension. + + Args: + path: The path to the directory to be searched. + + Returns: + A list of file paths. + """ + return super().find_files(path) diff --git a/src/graphnet/models/data_representation/data_representation.py b/src/graphnet/models/data_representation/data_representation.py index 667ce1d62..07cc3ba52 100644 --- a/src/graphnet/models/data_representation/data_representation.py +++ b/src/graphnet/models/data_representation/data_representation.py @@ -289,14 +289,14 @@ def _validate_input( # `DataRepresentation` was instantiated with. assert len(input_feature_names) == len( self._input_feature_names - ), f"""Input features ({input_feature_names}) is not what + ), f"""Input features ({input_feature_names}) is not what {self.__class__.__name__} was instatiated with ({self._input_feature_names})""" # noqa for idx in range(len(input_feature_names)): assert ( input_feature_names[idx] == self._input_feature_names[idx] ), f""" Order of representation features in data - are not the same as expected. Got {input_feature_names} + are not the same as expected. Got {input_feature_names} vs. {self._input_feature_names}""" # noqa def _perturb_input(self, input_features: np.ndarray) -> np.ndarray: diff --git a/src/graphnet/models/pretrained/icecube/kaggle/icemix/neutrino_direction/README.md b/src/graphnet/models/pretrained/icecube/kaggle/icemix/neutrino_direction/README.md index 18e69c5ee..aa669e789 100644 --- a/src/graphnet/models/pretrained/icecube/kaggle/icemix/neutrino_direction/README.md +++ b/src/graphnet/models/pretrained/icecube/kaggle/icemix/neutrino_direction/README.md @@ -11,7 +11,7 @@ Original repository can be found [here](https://github.com/DrHB/icecube-2nd-plac The state dict of the models is an exact copy of the one used in the competition, so further state_dict keys might be needed to load the models. Here is some pseudo-code which might help you load the weights into GraphNet: -```python +```python checkpoint = torch.load(path_to_checkpoint, torch.device("cpu")) if "state_dict" in checkpoint: checkpoint = checkpoint["state_dict"] @@ -65,7 +65,7 @@ class Direction_flipped(Label): ).reshape(-1, 1) z = torch.cos(graph[self._zenith_key]).reshape(-1, 1) return torch.cat((y, x, z), dim=1) # Flipped x and y -``` +``` ## Hard Local Coincidence feature @@ -89,7 +89,7 @@ def _hlc(self, x: torch.tensor) -> torch.tensor: The original models were stacked using a simple linear combination of the predictions. Here is some pseudo-code which can be used to stack the models: -```python +```python MODEL_NAMES = [ "B_d32", "B_d64", @@ -183,4 +183,4 @@ def inference( return results -``` \ No newline at end of file +``` diff --git a/src/graphnet/training/loss_functions.py b/src/graphnet/training/loss_functions.py index 71f994944..db55956be 100644 --- a/src/graphnet/training/loss_functions.py +++ b/src/graphnet/training/loss_functions.py @@ -15,6 +15,7 @@ from torch.nn.functional import ( one_hot, binary_cross_entropy, + binary_cross_entropy_with_logits, softplus, ) @@ -203,16 +204,30 @@ def _forward(self, prediction: Tensor, target: Tensor) -> Tensor: class BinaryCrossEntropyLoss(LossFunction): - """Compute binary cross entropy loss. + """Compute binary cross entropy loss.""" - Predictions are vector probabilities (i.e., values between 0 and 1), - and targets should be 0 and 1. - """ + def __init__(self, from_logits: bool = False, *args: Any, **kwargs: Any): + """Construct BinaryCrossEntropyLoss. + + Args: + from_logits: Whether the predictions are logits. + NOTE: If True, the predictions are expected to be raw scores + (i.e., not passed through a sigmoid function). If False, the + predictions are expected to be probabilities + (i.e., passed through a sigmoid function). + """ + super().__init__(*args, **kwargs) + self._from_logits = from_logits def _forward(self, prediction: Tensor, target: Tensor) -> Tensor: - return binary_cross_entropy( - prediction.float(), target.float(), reduction="none" - ) + if self._from_logits: + return binary_cross_entropy_with_logits( + prediction.float(), target.float(), reduction="none" + ) + else: + return binary_cross_entropy( + prediction.float(), target.float(), reduction="none" + ) class LogCMK(torch.autograd.Function): @@ -271,39 +286,39 @@ def backward( ctx: Any, grad_output: Tensor ) -> Tensor: # pylint: disable=invalid-name,arguments-differ """Backward pass for LogCMK computation. - + Mathematical Background: ----------------------- - For the von Mises-Fisher distribution, the gradient of log C_m(Îș) with + For the von Mises-Fisher distribution, the gradient of log C_m(Îș) with respect to Îș is given by the ratio of modified Bessel functions: - + ∂/∂Îș log C_m(Îș) = (m/2-1)/Îș - I_{m/2}(Îș)/I_{m/2-1}(Îș) - + For m=3, this simplifies to the exact formula: ∂/∂Îș log C_3(Îș) = 1/Îș - 1/tanh(Îș) - + For small Îș values, we use the Taylor series approximation: f(Îș) = -Îș/3 + ÎșÂł/45 - 2Îș⁔/945 + O(Îș⁷) - - The first-order approximation -Îș/3 provides sufficient accuracy for + + The first-order approximation -Îș/3 provides sufficient accuracy for |Îș| < 1e-6, with truncation error bounded by |Îș|Âł/45 â‰Č O(10⁻ÂČÂč). - + Implementation Details: ---------------------- Uses boolean masking to avoid double evaluation and RuntimeWarnings: - Small Îș: |Îș| < 1e-6 → gradient = -Îș/3 (Taylor approximation) - Large Îș: |Îș| ≄ 1e-6 → gradient = 1/Îș - 1/tanh(Îș) (exact formula) - + References: ---------- [1] von Mises-Fisher distribution: Wikipedia [2] arXiv:1812.04616, Section 8.2 [3] MIT License (c) 2019 Max Ryabinin - Modified for GraphNeT - + Args: ctx: Autograd context containing saved tensors and metadata. grad_output: Gradient with respect to the output tensor. - + Returns: Tuple of gradients: (None for m, gradient w.r.t. Îș). """ @@ -314,16 +329,16 @@ def backward( if np.isclose(m, 3, atol=1e-6): # Initialize gradient array grads = np.zeros_like(kappa) - + # Handle small kappa values (including zero) to avoid division by zero small_mask = np.abs(kappa) < 1e-6 grads[small_mask] = -kappa[small_mask] / 3 - + # Handle large kappa values large_mask = ~small_mask if np.any(large_mask): kappa_large = kappa[large_mask] - grads[large_mask] = 1/kappa_large - 1/np.tanh(kappa_large) + grads[large_mask] = 1 / kappa_large - 1 / np.tanh(kappa_large) else: grads = -( (scipy.special.iv(m / 2.0, kappa)) diff --git a/src/graphnet/utilities/imports.py b/src/graphnet/utilities/imports.py index 5ce5e45d0..0fb092eb4 100644 --- a/src/graphnet/utilities/imports.py +++ b/src/graphnet/utilities/imports.py @@ -47,6 +47,19 @@ def has_jammy_flows_package() -> bool: return False +def has_km3net_package() -> bool: + """Check whether the `km3net` packages are available.""" + try: + import km3io # pyright: reportMissingImports=false + + return True + except ImportError: + Logger(log_folder=None).warning_once( + "`km3net` not available. Some functionality may be missing.", + ) + return False + + def requires_icecube(test_function: Callable) -> Callable: """Decorate `test_function` for use only if `icecube` module is present.""" diff --git a/tests/training/test_loss_functions.py b/tests/training/test_loss_functions.py index 2d67628a8..16ac8cef0 100644 --- a/tests/training/test_loss_functions.py +++ b/tests/training/test_loss_functions.py @@ -185,28 +185,30 @@ def test_von_mises_fisher_approximation_large_kappa( ) -def test_logcmk_backward_zero_handling(dtype: torch.dtype = torch.float64) -> None: +def test_logcmk_backward_zero_handling( + dtype: torch.dtype = torch.float64, +) -> None: """Test LogCMK backward pass handles arrays with zero values correctly. - + This test ensures that the LogCMK.backward method correctly handles cases where the kappa tensor contains zero values without raising division by zero errors or warnings. The implementation uses boolean masking to conditionally apply different formulas for small (including zero) and large kappa values, avoiding double evaluation that would cause RuntimeWarnings. - + Args: dtype: PyTorch data type for the test tensors. """ # Test parameters m = 3 # Dimension for which we have the exact formula - + # Create kappa tensor with zeros and other values, including edge cases kappa_values = [0.0, 1e-7, 1e-6, 1e-5, 0.1, 1.0, 10.0] kappa = torch.tensor(kappa_values, dtype=dtype, requires_grad=True) - + # Forward pass using VonMisesFisherLoss.log_cmk_exact which internally uses LogCMK result = VonMisesFisherLoss.log_cmk_exact(m, kappa) - + # Test that backward pass doesn't raise any errors or warnings # Capture warnings to ensure no RuntimeWarnings are generated with warnings.catch_warnings(record=True) as caught_warnings: @@ -224,33 +226,45 @@ def test_logcmk_backward_zero_handling(dtype: torch.dtype = torch.float64) -> No except (ZeroDivisionError, RuntimeWarning) as e: backward_success = False error_msg = str(e) - + # Verify no errors occurred assert backward_success, f"Backward pass failed with error: {error_msg}" - + # Verify no RuntimeWarnings were generated - runtime_warnings = [w for w in caught_warnings if issubclass(w.category, RuntimeWarning)] - assert len(runtime_warnings) == 0, f"RuntimeWarnings were generated: {[str(w.message) for w in runtime_warnings]}" - + runtime_warnings = [ + w for w in caught_warnings if issubclass(w.category, RuntimeWarning) + ] + assert ( + len(runtime_warnings) == 0 + ), f"RuntimeWarnings were generated: {[str(w.message) for w in runtime_warnings]}" + # Verify gradients are finite - assert torch.all(torch.isfinite(grads)), "Gradients should be finite for all kappa values" - + assert torch.all( + torch.isfinite(grads) + ), "Gradients should be finite for all kappa values" + # Test specific values for correctness # For kappa=0, the gradient should be -kappa/3 = 0 zero_idx = 0 # Index where kappa=0 - assert torch.isclose(grads[zero_idx], torch.tensor(0.0, dtype=dtype)), \ - f"Gradient at kappa=0 should be 0, got {grads[zero_idx]}" - + assert torch.isclose( + grads[zero_idx], torch.tensor(0.0, dtype=dtype) + ), f"Gradient at kappa=0 should be 0, got {grads[zero_idx]}" + # For very small kappa (1e-7), should use -kappa/3 approximation small_kappa_idx = 1 # Index where kappa=1e-7 expected_small_grad = -kappa_values[small_kappa_idx] / 3 - assert torch.isclose(grads[small_kappa_idx], torch.tensor(expected_small_grad, dtype=dtype), atol=1e-10), \ - "Gradient for small kappa should use -kappa/3 approximation" - + assert torch.isclose( + grads[small_kappa_idx], + torch.tensor(expected_small_grad, dtype=dtype), + atol=1e-10, + ), "Gradient for small kappa should use -kappa/3 approximation" + # Test with array containing multiple zeros - kappa_multi_zero = torch.tensor([0.0, 0.0, 1.0, 0.0, 10.0], dtype=dtype, requires_grad=True) + kappa_multi_zero = torch.tensor( + [0.0, 0.0, 1.0, 0.0, 10.0], dtype=dtype, requires_grad=True + ) result_multi = VonMisesFisherLoss.log_cmk_exact(m, kappa_multi_zero) - + with warnings.catch_warnings(record=True) as caught_warnings_multi: warnings.simplefilter("always") try: @@ -262,14 +276,24 @@ def test_logcmk_backward_zero_handling(dtype: torch.dtype = torch.float64) -> No multi_zero_success = True except (ZeroDivisionError, RuntimeWarning): multi_zero_success = False - + assert multi_zero_success, "Should handle arrays with multiple zero values" - assert torch.all(torch.isfinite(grads_multi)), "All gradients should be finite with multiple zeros" - + assert torch.all( + torch.isfinite(grads_multi) + ), "All gradients should be finite with multiple zeros" + # Verify no RuntimeWarnings for multiple zeros case - runtime_warnings_multi = [w for w in caught_warnings_multi if issubclass(w.category, RuntimeWarning)] - assert len(runtime_warnings_multi) == 0, f"RuntimeWarnings were generated with multiple zeros: {[str(w.message) for w in runtime_warnings_multi]}" - + runtime_warnings_multi = [ + w + for w in caught_warnings_multi + if issubclass(w.category, RuntimeWarning) + ] + assert ( + len(runtime_warnings_multi) == 0 + ), f"RuntimeWarnings were generated with multiple zeros: {[str(w.message) for w in runtime_warnings_multi]}" + # Verify that zero elements have zero gradients zero_mask = kappa_multi_zero == 0.0 - assert torch.all(grads_multi[zero_mask] == 0.0), "Zero kappa values should have zero gradients" + assert torch.all( + grads_multi[zero_mask] == 0.0 + ), "Zero kappa values should have zero gradients"