rewriting codebase (#7)

916755e about 1 month ago

preview code

Raw

History Blame

8.96 kB

Task: Initialize the GNN4Colliders project structure

We are rewriting an existing collider-ML repository from scratch.

The repository is named GNN4Colliders.

The installable Python package should be:

gnn4colliders

The initial model family being rewritten is called:

root_gnn

In the future, the project may contain other architectures such as:

root_transformer

The goal is therefore to build a reusable collider-ML package rather than make the entire package GNN-specific.

Important context

Read these files before making changes:

docs/architecture.md

and any existing repository-level documentation.

The legacy implementation should be treated as a behavioral reference only.

If a legacy/ directory exists:

do not modify it
do not reorganize it
do not copy its architecture blindly
do not begin migrating implementation code in this task

This task is only about establishing the new project structure and documentation.

Desired architecture

Create this initial structure:

GNN4Colliders/
├── README.md
├── AGENTS.md
├── pyproject.toml
│
├── configs/
│   ├── data/
│   ├── model/
│   │   └── root_gnn/
│   ├── task/
│   ├── trainer/
│   └── environment/
│
├── src/
│   └── gnn4colliders/
│       ├── __init__.py
│       │
│       ├── data/
│       │   └── __init__.py
│       │
│       ├── features/
│       │   └── __init__.py
│       │
│       ├── graphs/
│       │   └── __init__.py
│       │
│       ├── models/
│       │   ├── __init__.py
│       │   └── root_gnn/
│       │       └── __init__.py
│       │
│       ├── training/
│       │   └── __init__.py
│       │
│       ├── inference/
│       │   └── __init__.py
│       │
│       └── cli/
│           └── __init__.py
│
├── tests/
│   ├── unit/
│   ├── integration/
│   ├── parity/
│   └── fixtures/
│
├── notebooks/
│
├── scripts/
│   ├── slurm/
│   └── dev/
│
├── docs/
│   ├── architecture.md
│   └── migration.md
│
└── legacy/

Do not create unnecessary placeholder Python modules yet.

Empty directories may contain .gitkeep files where required.

Architectural intent

The package should follow these responsibilities.

`gnn4colliders.data`

Generic data access and dataset infrastructure.

Future responsibilities may include:

ROOT/Awkward input
dataset abstractions
cache handling
batching
folds/splits

This package should not contain GNN-specific logic.

`gnn4colliders.features`

Physics-domain transformations that can be shared across architectures.

Future responsibilities may include:

collider object features
selections
feature scaling
derived physics quantities

These should be usable by both GNNs and future transformer-like architectures.

`gnn4colliders.graphs`

Graph-specific representation logic.

Future responsibilities may include:

graph topology
graph construction
edge feature construction
graph padding

Keep graph representation separate from generic ROOT input where possible.

`gnn4colliders.models`

Architecture-specific ML models.

Initially:

gnn4colliders.models.root_gnn

Future architectures may include:

gnn4colliders.models.root_transformer

Shared infrastructure should not be placed inside root_gnn unless it is genuinely GNN-specific.

`gnn4colliders.training`

Architecture-independent training infrastructure where practical.

Future responsibilities may include:

training lifecycle
losses
metrics
checkpointing
distributed training
reproducibility utilities

Do not implement these yet.

`gnn4colliders.inference`

Architecture-independent inference/application infrastructure where practical.

Future responsibilities may include:

prediction
evaluation
output writers
model export

`gnn4colliders.cli`

Thin command-line entry points.

Eventually the project should support commands conceptually similar to:

gnn4colliders prepare
gnn4colliders train
gnn4colliders evaluate
gnn4colliders predict
gnn4colliders export

Do not implement these workflows yet unless minimal CLI scaffolding is necessary for packaging.

Configuration philosophy

Create the config directory structure, but do not reproduce the legacy dynamic-import configuration system.

Configuration should eventually describe experiments semantically.

For example:

model:
  type: root_gnn

rather than exposing Python internals such as:

module: some.python.module
class: SomeClass

We expect to use composable YAML configuration, likely with Hydra, but this task should not build the configuration system beyond any minimal dependency or documentation decision that is clearly justified.

`pyproject.toml`

Create a modern minimal pyproject.toml.

Requirements:

project name: gnn4colliders
use src/ layout
package discovery should find src/gnn4colliders
specify an appropriate modern Python minimum version
include only dependencies that are clearly required for the initial project skeleton
development tooling may include:
- pytest
- ruff

Do not prematurely add the full legacy dependency environment.

Do not pin CUDA, PyTorch, DGL, ROOT, or other scientific dependencies until their compatibility strategy is addressed separately.

If adding a CLI entry point now would require inventing implementation, leave it out and document the intended future CLI instead.

README.md

Create a useful initial README that explains:

What GNN4Colliders is.
That it is a collider-ML toolkit intended to support multiple model families.
That root_gnn is the first model family being rewritten.
That additional architectures such as root_transformer may be added later.
The high-level package layout.
The distinction between:
- shared collider/data infrastructure
- representation-specific code
- model-specific code
- experiment configuration
The intended future CLI/config workflow.
Current development status:
- architecture and migration scaffold
- implementation not yet complete
Basic developer setup using an editable install.
A short development philosophy emphasizing:
- testability
- reproducibility
- explicit interfaces
- incremental migration from legacy behavior

Do not claim features already work when they do not.

AGENTS.md

Create the repository-level AGENTS.md using the project instructions supplied in this task.

It should establish durable rules for future coding-agent work, including:

repository purpose
architecture boundaries
legacy code rules
testing expectations
reproducibility expectations
configuration conventions
code-quality conventions
incremental migration workflow
requirement to inspect relevant existing code before changing behavior
requirement to avoid broad unrelated refactors

The file should remain concise enough to serve as a practical agent instruction file.

migration.md

If docs/migration.md does not already exist, create a minimal one.

It should list the intended migration stages without implementing them:

project/package skeleton
characterization/parity tests
configuration schema
ROOT/Awkward I/O
physics feature extraction and selections
graph construction
dataset caching/loading/batching
active ROOT-GNN model
losses and metrics
training lifecycle
checkpoints
inference
export
distributed/HPC workflows
legacy removal after parity

If docs/migration.md already exists, preserve its content unless a small structural update is clearly necessary.

Constraints

Do not:

rewrite legacy model code
migrate training logic
implement datasets
implement graph construction
introduce Lightning
introduce Kedro
create speculative abstractions
add unnecessary dependencies
modify legacy code
remove existing documentation
perform broad cleanup unrelated to this task

Prefer the smallest structure that establishes clear long-term boundaries.

Validation

After making changes:

Show the resulting directory tree.
Verify the package can be discovered/imported if practical.
Run any formatter/linter/tests that are available and relevant.
Check that pyproject.toml is valid.
Review README.md and AGENTS.md for claims about functionality that does not exist yet.

Then report:

files created
files modified
validation performed
any decisions you intentionally deferred
any conflicts you found with the existing architecture documentation