File size: 5,830 Bytes
70ee5cc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d2d4ccf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70ee5cc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8344837
 
 
 
 
 
 
 
 
 
 
 
 
 
70ee5cc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
# GNN4Colliders quickstart

GNN4Colliders turns collider ROOT events into shared event samples, applies
physics features, builds a representation for a selected model family, and
runs training or inference through one CLI. The current production model
family is `root_gnn`; the shared data and task layers are intended to support
future representations as well.

```text
ROOT -> EventSample -> features -> GraphSampleCache -> ROOT-GNN -> logits
```

## Install

The supported environment is Python 3.12. From the repository root:

```bash
uv sync --dev --extra root-gnn
```

The `root-gnn` extra is required for DGL-backed graph workflows and ROOT-GNN
parity tests. Add `--extra onnx` when using export validation.

## Run a local end-to-end smoke test

This is the fastest way to verify the package without downloading production
data:

```bash
uv run python scripts/dev/smoke_end_to_end.py
```

The script creates a tiny temporary ROOT file, prepares a cache, trains one
CPU epoch, evaluates, predicts, and cleans up. It is the recommended first
command after installation or a code change.

## Run the public Hugging Face smoke configuration

The small pinned fixture is selected by `config_hf_smoke`:

```bash
uv run gnn4colliders prepare --config-name config_hf_smoke
uv run gnn4colliders train --config-name config_hf_smoke
```

The first command resolves and caches the public ROOT fixture. The default
debug trainer runs one epoch. Outputs are written below `outputs/hf_smoke/`;
the resolved configuration and checkpoint path are recorded there.

For the full 12-process pretraining source, use this only when the larger
download is intentional:

```bash
uv run gnn4colliders prepare --config-name config_hf_delphes
```

The data source, revision, checksum, tree, feature branches, and split rules
are documented in [configuration.md](configuration.md).

The inclusive no-selection ttH CP benchmark has a dedicated configuration. It
uses `ttH_NLO.root` as label 0 (CP-even) and `ttH_CPodd.root` as label 1
(CP-odd), with all configured object branches passed through without event or
object selections:

```bash
uv run gnn4colliders prepare \
  --config-name config_tth_cp_even_odd \
  data.num_workers=8
```

This task configuration enables absolute event weights for optimization
because the NLO files contain signed weights. The original signed weights are
still retained in event metadata and used by the configured evaluation
semantics.

## Use your own ROOT file

Preparation accepts Hydra overrides. Supply the ROOT file, tree, cache path,
feature branch specification, object types, and scales:

```bash
uv run gnn4colliders prepare \
  data.files=[data/events.root] \
  data.tree_name=Events \
  data.cache.path=cache/events.pt \
  'data.feature_branches=[[jet_pt],[jet_eta],[jet_phi],CALC_E,[1.0],[0.0],NODE_TYPE]' \
  data.object_types=[vector] \
  data.scales=[1,1,1,1,1,1,1]
```

Then train, evaluate, or predict using the same cache:

```bash
uv run gnn4colliders train data.cache.path=cache/events.pt trainer.max_epochs=1
uv run gnn4colliders evaluate \
  data.cache.path=cache/events.pt \
  inference.checkpoint=outputs/default/checkpoints/epoch_0000.pt
uv run gnn4colliders predict \
  data.cache.path=cache/events.pt \
  inference.checkpoint=outputs/default/checkpoints/epoch_0000.pt \
  inference.output=outputs/predictions.npz
```

The feature override above is illustrative: replace it with branches that
exist in the input tree. See [configuration.md](configuration.md) for the
full schema.

## Models, tasks, and checkpoints

Select behavior with semantic config groups rather than Python module paths:

```bash
uv run gnn4colliders train \
  model=root_gnn/edge_network \
  task=pretraining_multiclass \
  data.cache.path=cache/events.pt \
  trainer.max_epochs=20
```

To fine-tune a pretrained ROOT-GNN checkpoint, use the fine-tuning model and
provide `checkpoint.pretrained`:

```bash
uv run gnn4colliders train \
  model=root_gnn/fine_tuned_edge_network \
  task=binary_classification \
  data.cache.path=cache/target.pt \
  checkpoint.pretrained=/path/to/pretrained.pt \
  model.freeze_backbone=true
```

`checkpoint.resume` continues an existing run, restoring its training state;
`checkpoint.pretrained` starts a new task from model weights and creates a new
head and optimizer. Evaluation uses the held-out test split by default.

The first transformer model reuses node features as ordered tokens and the
same task/trainer lifecycle. It can be trained from an existing prepared cache:

```bash
uv run gnn4colliders train \
  model=root_transformer/transformer \
  task=binary_classification \
  data.cache.path=cache/target.pt \
  model.hid_size=64 model.n_heads=4
```

This bridge intentionally ignores graph edges. Native sequence-cache storage
and distributed sequence loading are follow-up work.

## Validate changes

Run focused tests first, then the full suite and style checks:

```bash
uv run pytest tests/unit -q
GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest tests/parity -q
uv run pytest -q
uv run ruff check .
uv run ruff format --check .
```

For the portable public-data regression:

```bash
uv run python -m validation.run_public_validation --no-download
```

The parity gate requires DGL and compares the new implementation against
committed frozen reference fixtures. Historical runtime code is not imported
from main; compatibility adapters are reserved for old checkpoint, metadata,
and output boundaries.

## Where to look next

- [Architecture](architecture.md) — package boundaries and data flow.
- [Configuration](configuration.md) — Hydra groups and data sources.
- [Migration](migration.md) — rewrite status and compatibility policy.
- [Developer guide](../README_PROJECT.md) — transfer learning, distributed
  training, ONNX export, and output contracts.