Pytorch support

.flake8

@@ -0,0 +1,3 @@
+[flake8]
+max-line-length = 88
+extend-ignore = E203, W503

.github/workflows/python-package-conda.yml

@@ -2,12 +2,12 @@ name: CI
 
 on:
   push:
-    branches:
-      - "**"
-    tags:
-      - "**"
+    branches: ["**"]
+    tags: ["**"]
   pull_request:
   workflow_dispatch:
+  schedule:
+    - cron: "0 2 * * *"  # Daily at 02:00 UTC
 
 jobs:
   build:
@@ -17,14 +17,22 @@ jobs:
         python-version: ['3.12', '3.13', '3.14']
         dl1dh-version: ['latest', 'nightly']
         tensorflow-version: ['latest', '2.16.*']
+        torch-version: ['latest', '2.2.*']
         exclude:
           - python-version: '3.13'
             tensorflow-version: '2.16.*'
           - python-version: '3.14'
             tensorflow-version: '2.16.*'
+          - python-version: '3.13'
+            torch-version: '2.2.*'
+          - python-version: '3.14'
+            torch-version: '2.2.*'
       max-parallel: 6
     runs-on: ${{ matrix.os }}
     continue-on-error: ${{ matrix.dl1dh-version == 'nightly' || matrix.python-version == '3.14' }}
+    env:
+      PIP_NO_CACHE_DIR: "1"
+      PIP_EXTRA_INDEX_URL: "https://download.pytorch.org/whl/cpu"
     steps:
       - uses: actions/checkout@v4
 
@@ -61,6 +69,11 @@ jobs:
           else
             pip install "tensorflow==${{ matrix.tensorflow-version }}"
           fi
+          if [ "${{ matrix.torch-version }}" = "latest" ]; then
+            pip install --upgrade torch torchvision
+          else
+            pip install "torch==${{ matrix.torch-version }}" torchvision
+          fi
 
       - name: Add MKL_THREADING_LAYER variable
         run: echo "MKL_THREADING_LAYER=GNU" >> $GITHUB_ENV
@@ -76,8 +89,8 @@ jobs:
         run: |
           source $HOME/miniconda/etc/profile.d/conda.sh
           conda activate ctlearn
-          pip install -e .
-
+          pip install -e .[tests]
+          
       - name: Run pytest
         run: |
           source $HOME/miniconda/etc/profile.d/conda.sh

.gitignore

@@ -1,5 +1,5 @@
 ctlearn/_version.py
-
+output_dir2/
 *.swp
 *.swo
 *.gemini
@@ -12,13 +12,34 @@ ctlearn/_version.py
 *.png
 *.csv
 *~
+*.log
+.vscode/
+launcher_workspace/
+*.h5
 
+output_dir2/
+output_dir/
+.vscode/
 # Compiled Python files
 __pycache__/
 *.py[cod]
 .DS_Store
 *.egg-info/
 dist
-
+.pytest_cache/*.log
 # Sphinx documentation
 docs/build/
+build/
+# Default pytorch output
+run/
+test/
+*.se2
+*.jar
+*.puml
+mc_tjark/
+calibration/
+test_local_cristian/
+
+test/prepare_file.py
+test_local_cristian/
+*.txt

Dockerfile

@@ -2,11 +2,20 @@
 FROM python:3.12 AS builder
 
 # Install git (needed for setuptools_scm during build) and build tool
+
 RUN apt-get update \
     && apt-get install -y --no-install-recommends git \
     && rm -rf /var/lib/apt/lists/*
 RUN pip install --no-cache-dir build
 
+# Copy source code needed for the build
+WORKDIR /repo
+COPY ./pyproject.toml MANIFEST.in ./
+COPY ./ctlearn ./ctlearn/
+# If .git is truly needed for versioning by setuptools_scm, copy it. Otherwise, omit.
+COPY ./.git ./.git/
+RUN pip install --no-cache-dir build
+
 # Copy source code needed for the build
 WORKDIR /repo
 COPY ./pyproject.toml MANIFEST.in ./
@@ -18,14 +27,22 @@ COPY ./.git ./.git/
 RUN python -m build --wheel
 
 # Stage 2: Create the final runtime image BASED ON NVIDIA's TF image
+# TODO what version to use ? after 24.?? TF 2.14 is not found in the container
+FROM nvcr.io/nvidia/tensorflow:24.01-tf2-py3 
+
+# Copy only the built wheel from the builder stage's dist directory
+# Build the wheel
+
+# Stage 2: Create the final runtime image BASED ON NVIDIA's TF image
+
 FROM nvcr.io/nvidia/tensorflow:25.02-tf2-py3 
 # Copy only the built wheel from the builder stage's dist directory
 COPY --from=builder /repo/dist /tmp/dist
 
+# Install the ctlearn wheel using pip from the NVIDIA base image
 # Install the ctlearn wheel using pip from the NVIDIA base image
 RUN python -m pip install --no-cache-dir /tmp/dist/* \
     && rm -r /tmp/dist
 
 RUN addgroup --system ctlearn && adduser --system --group ctlearn
-USER ctlearn
-
+USER ctlearn

README.rst

@@ -47,6 +47,49 @@ The lastest version fo this package can be installed as a pip package:
 
 See the documentation for further information like `installation instructions for the IT-cluster <https://ctlearn.readthedocs.io/en/latest/installation.html#install-a-released-version>`_, `installation instructions for developers <https://ctlearn.readthedocs.io/en/latest/installation.html#installing-with-pip-setuptools-from-source-for-development>`_, `package usage <https://ctlearn.readthedocs.io/en/stable/usage.html>`_, and `dependencies <https://ctlearn.readthedocs.io/en/stable/installation.html#dependencies>`_ among other topics.
 
+Running CTLearn Training and Prediction
+---------------------------------------
+
+CTLearn provides a unified command line interface (CLI) using `ctapipe`'s ``Tool`` and ``Component`` systems, supporting both the **Keras** and **PyTorch** deep learning frameworks.
+
+Launching training
+~~~~~~~~~~~~~~~~~~
+
+You can launch a training run using the unified tool ``ctlearn-train``. To run with a specific framework, set the ``--framework`` option (choices: ``keras``, ``pytorch``):
+
+.. code-block:: bash
+
+   # Launch training with PyTorch
+   ctlearn-train --framework=pytorch --output ./my_output_dir --signal /path/to/signal/h5/ --pattern-signal "*.dl1.h5" --reco energy --n_epochs=10 --batch_size=32
+
+   # Launch training with Keras
+   ctlearn-train --framework=keras --output ./my_output_dir --signal /path/to/signal/h5/ --pattern-signal "*.dl1.h5" --reco energy --n_epochs=10 --batch_size=32
+
+Common Training Command Options:
+* ``--framework``: Deep learning framework to use (``keras`` or ``pytorch``).
+* ``-o``, ``--output``: Directory to save experiment checkpoints, parameters, and logs.
+* ``--signal``: Directory containing signal HDF5 data files.
+* ``--pattern-signal``: File name pattern for signal files (e.g. ``*.dl1.h5``).
+* ``--reco``: Tasks to train (e.g. ``type``, ``energy``, ``cameradirection``, ``skydirection``). Multiple tasks can be provided.
+* ``--n_epochs``: Number of epochs to train.
+* ``--batch_size``: Batch size for training.
+* ``--save_onnx=True``: Export the trained model to ONNX format.
+* ``--load_onnx_model=PATH``: Load an existing ONNX model to train/fine-tune.
+* ``--overwrite``: Overwrite the output directory if it already exists.
+
+Launching prediction
+~~~~~~~~~~~~~~~~~~~~
+
+Similarly, predictions on test/observation data can be executed using the unified prediction tools (for monoscopic or stereoscopic mode):
+
+.. code-block:: bash
+
+   # Monoscopic prediction with PyTorch
+   ctlearn-predict-mono-model --framework=pytorch --output ./pred_results --signal /path/to/test/h5/ --pattern-signal "*.dl1.h5" --energy_checkpoint /path/to/checkpoint.pth
+
+   # Monoscopic prediction with Keras
+   ctlearn-predict-mono-model --framework=keras --output ./pred_results --signal /path/to/test/h5/ --pattern-signal "*.dl1.h5" --energy_checkpoint /path/to/keras_model/
+
 Citing this software
 --------------------
 

condaenv.vowayxih.requirements.txt

@@ -0,0 +1,11 @@
+numba
+tensorflow>=2.14,<2.15
+dl1_data_handler>=0.14.1,<0.15
+pydot
+pytorch_lightning
+deepspeed
+onnx
+onnxsim
+torch==2.5.0
+torchvision==0.20.0
+torchaudio==2.5.0