Skip to content

CNN PR [3/4]: Add CNN training example#896

Draft
sevmag wants to merge 5 commits into
graphnet-team:mainfrom
sevmag:features/cnn-train-example
Draft

CNN PR [3/4]: Add CNN training example#896
sevmag wants to merge 5 commits into
graphnet-team:mainfrom
sevmag:features/cnn-train-example

Conversation

@sevmag
Copy link
Copy Markdown
Collaborator

@sevmag sevmag commented May 27, 2026

Summary

Fifth and final PR carved out of #813 (see the split tracking comment). Adds an end-to-end training example for the CNN models.

What's new

  • examples/04_training/09_train_cnn.py — builds an image data representation (ExamplePrometheusImage), feeds it to an LCSC backbone inside a StandardModel, and trains an EnergyReconstruction task on the bundled example data.

⚠️ Depends on #892, #894, #895 — draft until those merge

This example imports from all three earlier split PRs:

Because those base branches only exist in the fork, this PR targets main directly, so its current diff includes the files from #892/#894/#895 in addition to the net-new example. Once those merge into main, GitHub will narrow this diff to just examples/04_training/09_train_cnn.py.

Review/merge order: merge #892, #894, then #895 first; this PR is kept as a draft until then, after which it'll be rebased onto main and marked ready.

Test plan

  • Pre-commit clean (black reformatted, then all hooks pass)
  • python -m py_compile succeeds
  • Module-level imports resolve against the combined image-representation + CNN branches
  • CI green (after the dependency PRs merge)
  • Manual: run python examples/04_training/09_train_cnn.py on the bundled example data

Split from #813.

🤖 Generated with Claude Code

sevmag added 5 commits May 19, 2026 08:57
Introduce a new `data_representation/images/` module that lets
detectors be represented as multi-channel 3D image tensors rather than
graphs, as required by CNN backbones.

Components:

- `ImageRepresentation` base class — analogue of `GraphDefinition`
  for image-shaped inputs. Builds per-channel image tensors from raw
  pulse data using a `GridDefinition` that maps each DOM to a voxel.
- `GridDefinition` — abstract mapping from DOM string/dom indices to a
  3D grid. `IC86GridDefinition` covers the IceCube IC86 array
  (main array + upper/lower DeepCore), `ExamplePrometheusGridDefinition`
  covers the Prometheus example geometry. Mapping tables live in
  `cnn_mapping_tables.py` as Python literals (no parquet files).
- `IC86Image` and `ExamplePrometheusImage` — concrete image
  representations bundling a detector and its grid definition.

Also adds:

- `TEST_IMAGE_DIR` and the three `TEST_IC86*_IMAGE` constants in
  `graphnet.constants` for the new test fixtures.
- `.npy` test fixtures for the IC86 main array and the two DeepCore
  sub-arrays.
- Unit tests for `GridDefinition` and `ImageRepresentation`.

No model code yet — that lands in follow-up PRs (`IceCubeDNN`, `LCSC`).
Split from graphnet-team#813.
Introduce `graphnet.models.cnn`, the convolutional-network counterpart
to the existing GNN backbones, intended to consume the image data
representation added in the image-representation PR.

- `CNN` — abstract base class (analogue of `GNN`) defining the
  interface for convolutional backbones operating on image-shaped
  `Data` objects.
- `IceCubeDNN` — configurable CNN backbone following the IceCube DNN
  reconstruction architecture.

`cnn/__init__.py` exports `CNN` and `IceCubeDNN`; the `LCSC` model
lands in a follow-up PR.

Split from graphnet-team#813.
Add the LCSC convolutional backbone, building on the `CNN` base class.
Registered in `cnn/__init__.py` alongside `CNN` and `IceCubeDNN`.

Stacked on the CNN-base PR. Split from graphnet-team#813.
Add `examples/04_training/09_train_cnn.py`, an end-to-end training
example for the CNN models: it builds an image data representation
(`ExamplePrometheusImage`), feeds it to an `LCSC` backbone inside a
`StandardModel`, and trains an `EnergyReconstruction` task on the
bundled example data.

Depends on the image data representation and CNN model PRs. Split from
graphnet-team#813.
@sevmag sevmag mentioned this pull request May 27, 2026
@sevmag sevmag changed the title Add CNN training example CNN PR [3/4]: Add CNN training example May 27, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant