Fixing predict tools in pytorch branch

ctlearn/conftest.py

@@ -7,22 +7,21 @@
 import numpy as np
 import pytest
 import shutil
-from astropy import units as u
 from astropy.table import Column, Table
 from traitlets.config.loader import Config
 
 from ctapipe.core import run_tool
 from ctapipe.io import write_table
+from ctapipe.tools.process import ProcessorTool
 from ctapipe.utils import get_dataset_path
+
 from ctlearn.tools.keras.train_model import TrainCTLearnKerasModel
 from ctlearn.tools.pytorch.train_model import TrainCTLearnPyTorchModel
-from ctlearn.utils import get_lst1_subarray_description
+from ctlearn.tools.utils import get_lst1_subarray_description
 
-# TODO: ADD PyTorch here 
 TRAINING_TOOLS = {"Keras": TrainCTLearnKerasModel, "PyTorch": TrainCTLearnPyTorchModel}
 MODEL_FILE_FORMATS = {"Keras": "keras", "PyTorch": "pth"}
 
-
 @pytest.fixture(scope="session")
 def gamma_simtel_path():
     return get_dataset_path("gamma_test_large.simtel.gz")
@@ -138,8 +137,6 @@ def dl1_gamma_file(dl1_tmp_path, gamma_simtel_path):
     """
     DL1 file containing both images and parameters from a gamma simulation set.
     """
-    from ctapipe.tools.process import ProcessorTool
-
     output = dl1_tmp_path / "gamma.dl1.h5"
     argv = [
         f"--input={gamma_simtel_path}",
@@ -174,10 +171,7 @@ def r1_gamma_file(r1_tmp_path, gamma_simtel_path):
     """
     R1 file containing both waveforms and parameters from a gamma simulation set.
     """
-    from ctapipe.tools.process import ProcessorTool
-
     output = r1_tmp_path / "gamma.r1.h5"
-
     allowed_tels = [1, 2]
     argv = [
         f"--input={gamma_simtel_path}",
@@ -195,8 +189,6 @@ def r1_proton_file(r1_tmp_path, proton_simtel_path):
     """
     R1 file containing both waveforms and parameters from a proton simulation set.
     """
-    from ctapipe.tools.process import ProcessorTool
-
     # Restrict to two LSTs for R1 tests to reduce computational load
     allowed_tels = [1, 2]
     output = r1_tmp_path / "proton.r1.h5"
@@ -403,14 +395,10 @@ def ctlearn_trained_dl1_stereo_models(
     # Loop over reconstruction tasks and train models for each combination
     ctlearn_trained_dl1_stereo_models = {}
     for reco_task in ["type", "energy", "skydirection"]:
-        # Output directory for trained model
-        output_dir = tmp_path / f"ctlearn_{telescope_type}_{reco_task}"
-
         # Build command-line arguments
         argv = [
             f"--signal={signal_dir}",
             "--pattern-signal=*.dl1.h5",
-            f"--output={output_dir}",
             f"--reco={reco_task}",
             "--TrainCTLearnModel.n_epochs=1",
             "--TrainCTLearnModel.batch_size=2",
@@ -432,7 +420,10 @@ def ctlearn_trained_dl1_stereo_models(
 
         # Run training tools
         for framework, training_tool in TRAINING_TOOLS.items():
-            assert run_tool(training_tool(config=config), argv=argv, cwd=tmp_path) == 0
+            framework_argv = argv.copy()
+            output_dir = tmp_path / f"ctlearn_{framework}_{telescope_type}_{reco_task}"
+            framework_argv.append(f"--output={output_dir}")
+            assert run_tool(training_tool(config=config), argv=framework_argv, cwd=tmp_path) == 0
             ctlearn_trained_dl1_stereo_models[f"{framework}_{telescope_type}_{reco_task}"] = (
                 output_dir / f"ctlearn_model.{MODEL_FILE_FORMATS[framework]}"
             )

ctlearn/core/ctlearn_enum.py

@@ -14,34 +14,6 @@
 
 from enum import Enum
 
-class FrameworkType(Enum):
-    """
-    Deep learning framework type enumeration.
-
-    This enumeration specifies which deep learning framework to use for
-    model training and inference. CTLearn supports both Keras (TensorFlow backend)
-    and PyTorch frameworks.
-
-    Attributes:
-        KERAS (int): Use Keras/TensorFlow framework (value: 1)
-            - Advantages: High-level API, easy to use, good for prototyping
-            - TensorFlow 2.x with Keras API
-            - Suitable for production deployment
-
-        PYTORCH (int): Use PyTorch framework (value: 2)
-            - Advantages: Dynamic computation graphs, flexible, research-friendly
-            - PyTorch 1.x or 2.x
-            - Better for custom architectures and experimental models
-
-    Example:
-        >>> from ctlearn.core.ctlearn_enum import FrameworkType
-        >>> framework = FrameworkType.PYTORCH
-        >>> print(framework.name)  # 'PYTORCH'
-        >>> print(framework.value)  # 2
-    """
-    KERAS = 1
-    PYTORCH = 2
-
 
 class Task(Enum):
     """

ctlearn/core/model.py

@@ -6,7 +6,7 @@
 
 from ctapipe.core import Component
 from ctapipe.core.traits import Bool, Int, CaselessStrEnum, List, Dict, Unicode, Path
-from ctlearn.utils import validate_trait_dict
+from ctlearn.tools.utils import validate_trait_dict
 
 __all__ = [
     "CTLearnModel",

ctlearn/core/pytorch/model.py

@@ -21,25 +21,27 @@
 __all__ = [
     "BasicBlock",
     "BottleneckBlock",
-    "MultiHeadClassifier",
+    "MultiFullyConnectedHead",
     "build_fully_connect_pytorch_head",
     "PyTorchSingleCNN",
     "PyTorchResNet",
     "PyTorchLoadedModel",
 ]
 
 
-class MultiHeadClassifier(nn.Module):
+class MultiFullyConnectedHead(nn.Module):
     """
     A PyTorch container module to hold the multi-task fully connected heads.
     """
 
     def __init__(self, heads_dict, single_output_task=None):
         super().__init__()
+        # Save the dict with tasks info in an attribute
+        self.heads_dict = heads_dict
         # Sanitize keys because 'type' conflicts with nn.Module.type() method
         self._task_mapping = {
             task: f"head_{task}" if hasattr(nn.Module, task) else task
-            for task in heads_dict.keys()
+            for task in self.heads_dict.keys()
         }
         sanitized_heads = {
             self._task_mapping[task]: module for task, module in heads_dict.items()
@@ -97,7 +99,7 @@ def build_fully_connect_pytorch_head(in_features, layers, activation_function, t
         heads[task] = nn.Sequential(*task_layers)
 
     single_output_task = tasks[0] if (len(tasks) == 1 and tasks[0] == "type") else None
-    return MultiHeadClassifier(heads, single_output_task=single_output_task)
+    return MultiFullyConnectedHead(heads, single_output_task=single_output_task)
 
 
 class FullModelPipeline(nn.Module):
@@ -111,7 +113,7 @@ def __init__(self, backbone, head):
 
     def forward(self, x):
         features = self.backbone(x)
-        return self.head(features)
+        return self.head(features), features
 
 
 class PyTorchSingleCNN(SingleCNN):

ctlearn/core/tests/test_attention.py

@@ -1,7 +1,4 @@
 import pytest
-pytest.importorskip("keras")
-import numpy as np
-import pytest
 import torch
 import keras
 import tensorflow as tf

ctlearn/core/tests/test_loaders.py

@@ -0,0 +1,54 @@
+import pytest
+from torch.utils.data import DataLoader
+from traitlets.config.loader import Config
+
+from dl1_data_handler.reader import DLImageReader
+from ctlearn.core.keras.sequence import KerasSequence
+from ctlearn.core.pytorch.dataset import PyTorchDataset
+
+
+@pytest.mark.parametrize(
+    ("dataloader_cls", "expected_features_shape"),
+    [
+        (KerasSequence, (1, 110, 110, 2)),
+        (PyTorchDataset, (1, 2, 110, 110)),
+    ],
+    ids=["Keras", "PyTorch"],
+)
+def test_data_loading(dl1_gamma_file, dataloader_cls, expected_features_shape):
+    """Check the data loading for the Keras and PyTorch frameworks"""
+    # Create a configuration suitable for the test
+    config = Config(
+        {
+            "DLImageReader": {
+                "allowed_tels": [4],
+                "focal_length_choice": "EQUIVALENT",
+            },
+        }
+    )
+    # Create an image reader
+    dl1_reader = DLImageReader(input_url_signal=[dl1_gamma_file], config=config)
+    # Initialize the PyTorch Dataset or the Keras Sequence
+    dl1_dataset = dataloader_cls(
+        DLDataReader=dl1_reader,
+        indices=[0],
+        tasks=["type", "energy", "cameradirection", "skydirection"],
+    )
+    # For PyTorch, wrap in DataLoader to apply batch dimension (dim=0)
+    if issubclass(dataloader_cls, PyTorchDataset):
+        batch_loader = DataLoader(dl1_dataset, batch_size=1)
+        features, labels = next(iter(batch_loader))
+        # Optional: convert PyTorch tensor to numpy if labels/features are Tensors
+        if hasattr(features, "numpy"):
+            features = features.numpy()
+    else:
+        features, labels = dl1_dataset[0]
+    #  Check that all the correct labels are present
+    assert (
+        "type" in labels
+        and "energy" in labels
+        and "cameradirection" in labels
+        and "skydirection" in labels
+    )
+    #  Check the shape of the features match the expected ones
+    assert features.shape == expected_features_shape

ctlearn/core/tests/test_models.py

@@ -1,9 +1,6 @@
 import pytest
-pytest.importorskip("keras")
-import re
-import keras
 import numpy as np
-import pytest
+import keras
 import torch
 import torch.nn as nn