diff --git a/.gitattributes b/.gitattributes index 48e1fe032537c7401a1b7f2740bf08603da43745..a6344aac8c09253b3b630fb776ae94478aa0275b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -33,4 +33,3 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text *.zip filter=lfs diff=lfs merge=lfs -text *.zst filter=lfs diff=lfs merge=lfs -text *tfevents* filter=lfs diff=lfs merge=lfs -text -training_time.png filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore index 6394ec02142783deceb8da2e61da24ebaa11796a..3f992ad4c0a1f5004e4c6b529902b67a7fc3c2c3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,16 +1,3 @@ -# Python tooling __pycache__/ -*.py[cod] -*.egg-info/ -.venv/ -.pytest_cache/ -.ruff_cache/ -.coverage -dist/ -build/ - -# Local data and generated outputs -data/ -outputs/ -logs/ -*.onnx +trainings/ +scores/ diff --git a/.python-version b/.python-version deleted file mode 100644 index 6324d401a069f4020efcf0ff07442724b52f47c2..0000000000000000000000000000000000000000 --- a/.python-version +++ /dev/null @@ -1 +0,0 @@ -3.14 diff --git a/AGENTS.md b/AGENTS.md deleted file mode 100644 index 2dc15d64245161fd05fa5a1da7d0a2d6d67e8704..0000000000000000000000000000000000000000 --- a/AGENTS.md +++ /dev/null @@ -1,52 +0,0 @@ -# GNN4Colliders agent instructions - -GNN4Colliders trains event classifiers (the ROOT-GNN edge network of Ho et -al., JINST 21 (2026) P08006, a transformer and Deep Sets) on reconstructed -collider objects. Keep it small and readable: a physicist should be able to follow a -run from `cli.py` to the model in a few files. - -## Layout - -```text -src/gnn4colliders/ flat modules; see docs/design.md for what each holds -configs/ one self-contained YAML file per experiment -scripts/ Slurm wrappers and small helpers (no core logic) -tests/ pytest, with small ROOT fixtures in tests/fixtures -Pretrained_GNN/ the published pretrained checkpoints (Git LFS) -``` - -* `README.md` is the paper's Hugging Face model card. Do not edit it; usage - goes in `docs/usage.md`, design and results in `docs/design.md`. -* The main benchmark is `configs/gamgam.yaml` (ATLAS Open Data GamGam, 5 - classes); compare with Table 3 of the paper (`docs/design.md`). - -## Rules - -* Add a module or an abstraction only when existing code cannot hold it - clearly. No new frameworks (Lightning, Hydra, DGL, PyG) without agreement. -* Configuration is plain dataclasses in `config.py`; every key has a default - or is required, and unknown keys are errors. Site settings (accounts, - paths, `UV_CACHE_DIR`) belong in scripts, not in configs or code. -* Name domain quantities (`labels`, `weight`, `fold`, `event_index`) instead - of relying on tensor column positions, except for the documented seven node - columns. -* Randomness is seeded explicitly (`train.seed`, `data.seed`); never mutate - global random state elsewhere. Do not claim bitwise reproducibility on GPUs. -* The node schema, edge construction, MLP layout and parameter names must - stay compatible with `Pretrained_GNN/`; a test loads both pretrained models. -* Changes to features, splits, loss or metrics change physics results: say so - explicitly and re-check the GamGam benchmark. - -## Checks - -```bash -uv sync --dev -uv run pytest # GPU tests run on a CUDA node -uv run ruff check . && uv run ruff format --check . -``` - -On Perlmutter, run GPU work on compute nodes (`salloc -C gpu -q -shared_interactive ...`) with `UV_CACHE_DIR=$SCRATCH/.cache/uv`. - -Report what changed, what was run and what was not, and any change to -results. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..f8bc5304ced57d6e9da89cdf00995433e2c3b716 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 LBL ATLAS + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Pretrained_GNN/multiclass_pretrained_model_12/config.yaml b/Pretrained_GNN/multiclass_pretrained_model_12/config.yaml deleted file mode 100755 index 1370d23875c785ac41bafb8011bb4324f761cca7..0000000000000000000000000000000000000000 --- a/Pretrained_GNN/multiclass_pretrained_model_12/config.yaml +++ /dev/null @@ -1,319 +0,0 @@ -Datasets: - SingleT_schan: - args: - buffer_size: 11 - chunks: 100 - file_names: SingleT_schan.root - fold_var: Number - label: 8 - name: SingleT_schan - node_branch_names: &id001 - - - jet_pt - - ele_pt - - mu_pt - - ph_pt - - MET_met - - - jet_eta - - ele_eta - - mu_eta - - ph_eta - - 0 - - - jet_phi - - ele_phi - - mu_phi - - ph_phi - - MET_phi - - CALC_E - - - jet_btag - - 0 - - 0 - - 0 - - 0 - - - 0 - - ele_charge - - mu_charge - - 0 - - 0 - - NODE_TYPE - node_branch_types: &id002 - - vector - - vector - - vector - - vector - - single - node_feature_scales: &id003 - - 1e-1 - - 1 - - 1 - - 1e-1 - - 1 - - 1 - - 1 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: &id004 - n_folds: 10 - test: - - 0 - - 1 - train: - - 2 - - 3 - - 4 - - 5 - - 6 - - 7 - - 8 - - 9 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - VBF: - args: - buffer_size: 11 - chunks: 100 - file_names: VBF_NLO_inc.root - fold_var: Number - label: 3 - name: VBF - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - WH: - args: - buffer_size: 11 - chunks: 100 - file_names: WH_NLO_inc.root - fold_var: Number - label: 4 - name: WH - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ZH: - args: - buffer_size: 11 - chunks: 100 - file_names: ZH_NLO_inc.root - fold_var: Number - label: 5 - name: ZH - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ggF: - args: - buffer_size: 11 - chunks: 100 - file_names: ggF_NLO_inc.root - fold_var: Number - label: 2 - name: ggF - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - tHjb: - args: - buffer_size: 11 - chunks: 100 - file_names: tHjb_NLO_inc.root - fold_var: Number - label: 1 - name: tHjb - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttH: - args: - buffer_size: 11 - chunks: 100 - file_names: ttH_NLO_inc.root - fold_var: Number - label: 0 - name: ttH - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttW: - args: - buffer_size: 11 - chunks: 100 - file_names: ttW.root - fold_var: Number - label: 10 - name: ttW - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttbar: - args: - buffer_size: 11 - chunks: 100 - file_names: ttbar.root - fold_var: Number - label: 9 - name: ttbar - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttt: - args: - buffer_size: 11 - chunks: 100 - file_names: ttt.root - fold_var: Number - label: 11 - name: ttt - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - tttt: - args: - buffer_size: 11 - chunks: 100 - file_names: tttt.root - fold_var: Number - label: 7 - name: tttt - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttyy: - args: - buffer_size: 11 - chunks: 100 - file_names: ttyy.root - fold_var: Number - label: 6 - name: ttyy_ch - node_branch_names: *id001 - node_branch_types: *id002 - node_feature_scales: *id003 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_pretraining_multiclass_12_process/ - tree_name: output - weight_var: weight - class: LazyDataset - folding: *id004 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 -Loss: - args: {} - class: CrossEntropyLoss - finish: - args: - dim: 1 - class: Softmax - module: torch.nn - module: torch.nn -Model: - args: - dropout: 0 - hid_size: 64 - in_size: 7 - n_layers: 4 - n_proc_steps: 4 - out_size: 12 - class: Edge_Network - module: models.GCN -Training: - batch_size: 1024 - epochs: 100 - gamma: 0.99 - learning_rate: 0.0001 -Training_Directory: trainings/pretraining_multiclass/multiclass_12_process/ -Training_Name: multiclass_12_process diff --git a/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt b/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt deleted file mode 100755 index 4c2025ee7cbd748ae0080a2b5648c4b1e17c9f3d..0000000000000000000000000000000000000000 --- a/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:f382f9f8b2d8c09b896657d397985f9149f018fcf3aa6d6f56ce3e51cf8a6351 -size 1707502 diff --git a/Pretrained_GNN/multiclass_pretrained_model_12/training.log b/Pretrained_GNN/multiclass_pretrained_model_12/training.log deleted file mode 100755 index 9412d7e31a69685e670faaa0eb986a936784b103..0000000000000000000000000000000000000000 --- a/Pretrained_GNN/multiclass_pretrained_model_12/training.log +++ /dev/null @@ -1,74 +0,0 @@ -Training multiclass_12_process 2024-09-06 23:50:52.835697 -Epoch 00000 | LR 1.0000e-04 | Loss 1.6853 | Accuracy 0.3672 | Test_Loss 1.6844 | Test_AUC 0.8323 | Time 2406.1502 s -Epoch 00001 | LR 9.9000e-05 | Loss 1.6216 | Accuracy 0.3743 | Test_Loss 1.6660 | Test_AUC 0.8344 | Time 2406.9453 s -Epoch 00002 | LR 9.8010e-05 | Loss 1.6066 | Accuracy 0.3769 | Test_Loss 1.6656 | Test_AUC 0.8389 | Time 2355.3162 s -Epoch 00003 | LR 9.7030e-05 | Loss 1.5989 | Accuracy 0.3752 | Test_Loss 1.6700 | Test_AUC 0.8379 | Time 2351.9926 s -Epoch 00004 | LR 9.6060e-05 | Loss 1.5936 | Accuracy 0.3809 | Test_Loss 1.6568 | Test_AUC 0.8435 | Time 2323.1655 s -Epoch 00005 | LR 9.5099e-05 | Loss 1.5898 | Accuracy 0.3784 | Test_Loss 1.6613 | Test_AUC 0.8405 | Time 2293.3972 s -Epoch 00006 | LR 9.4148e-05 | Loss 1.5869 | Accuracy 0.3832 | Test_Loss 1.6416 | Test_AUC 0.8416 | Time 2297.7492 s -Epoch 00007 | LR 9.3207e-05 | Loss 1.5844 | Accuracy 0.3917 | Test_Loss 1.6090 | Test_AUC 0.8455 | Time 2296.5126 s -Epoch 00008 | LR 9.2274e-05 | Loss 1.5822 | Accuracy 0.3895 | Test_Loss 1.6174 | Test_AUC 0.8429 | Time 2372.6853 s -Epoch 00009 | LR 9.1352e-05 | Loss 1.5810 | Accuracy 0.3965 | Test_Loss 1.6079 | Test_AUC 0.8453 | Time 2455.0305 s -Epoch 00010 | LR 9.0438e-05 | Loss 1.5795 | Accuracy 0.3966 | Test_Loss 1.6067 | Test_AUC 0.8468 | Time 2459.9180 s -Epoch 00011 | LR 8.9534e-05 | Loss 1.5782 | Accuracy 0.3943 | Test_Loss 1.6117 | Test_AUC 0.8465 | Time 2444.6930 s -Epoch 00012 | LR 8.8638e-05 | Loss 1.5770 | Accuracy 0.3973 | Test_Loss 1.6081 | Test_AUC 0.8471 | Time 2475.8966 s -Epoch 00013 | LR 8.7752e-05 | Loss 1.5759 | Accuracy 0.3973 | Test_Loss 1.6037 | Test_AUC 0.8486 | Time 2330.9825 s -Epoch 00014 | LR 8.6875e-05 | Loss 1.5749 | Accuracy 0.3974 | Test_Loss 1.6066 | Test_AUC 0.8491 | Time 2276.6652 s -Epoch 00015 | LR 8.6006e-05 | Loss 1.5741 | Accuracy 0.3981 | Test_Loss 1.6070 | Test_AUC 0.8475 | Time 2278.0207 s -Epoch 00016 | LR 8.5146e-05 | Loss 1.5733 | Accuracy 0.3958 | Test_Loss 1.6089 | Test_AUC 0.8471 | Time 2296.7680 s -Epoch 00017 | LR 8.4294e-05 | Loss 1.5726 | Accuracy 0.3978 | Test_Loss 1.6004 | Test_AUC 0.8491 | Time 2273.7842 s -Epoch 00018 | LR 8.3451e-05 | Loss 1.5718 | Accuracy 0.3847 | Test_Loss 1.6652 | Test_AUC 0.8448 | Time 2265.8179 s -Epoch 00019 | LR 8.2617e-05 | Loss 1.5713 | Accuracy 0.3858 | Test_Loss 1.6495 | Test_AUC 0.8434 | Time 2271.3266 s -Epoch 00020 | LR 8.1791e-05 | Loss 1.5708 | Accuracy 0.3842 | Test_Loss 1.6602 | Test_AUC 0.8431 | Time 2263.7548 s -Epoch 00021 | LR 8.0973e-05 | Loss 1.5703 | Accuracy 0.3829 | Test_Loss 1.6628 | Test_AUC 0.8439 | Time 2277.5963 s -Epoch 00022 | LR 8.0163e-05 | Loss 1.5699 | Accuracy 0.3873 | Test_Loss 1.6452 | Test_AUC 0.8436 | Time 2381.2508 s -Epoch 00023 | LR 7.9361e-05 | Loss 1.5694 | Accuracy 0.3841 | Test_Loss 1.6477 | Test_AUC 0.8446 | Time 2371.5616 s -Epoch 00024 | LR 7.8568e-05 | Loss 1.5690 | Accuracy 0.3917 | Test_Loss 1.6253 | Test_AUC 0.8448 | Time 2412.1963 s -Epoch 00025 | LR 7.7782e-05 | Loss 1.5685 | Accuracy 0.3997 | Test_Loss 1.5952 | Test_AUC 0.8483 | Time 2361.6398 s -Epoch 00026 | LR 7.7004e-05 | Loss 1.5680 | Accuracy 0.3975 | Test_Loss 1.6004 | Test_AUC 0.8491 | Time 2370.2956 s -Epoch 00027 | LR 7.6234e-05 | Loss 1.5679 | Accuracy 0.4008 | Test_Loss 1.6004 | Test_AUC 0.8496 | Time 2401.6121 s -Epoch 00028 | LR 7.5472e-05 | Loss 1.5678 | Accuracy 0.3993 | Test_Loss 1.6004 | Test_AUC 0.8492 | Time 2403.7899 s -Epoch 00029 | LR 7.4717e-05 | Loss 1.5676 | Accuracy 0.3993 | Test_Loss 1.5997 | Test_AUC 0.8497 | Time 2407.6645 s -Epoch 00030 | LR 7.3970e-05 | Loss 1.5672 | Accuracy 0.4022 | Test_Loss 1.5928 | Test_AUC 0.8498 | Time 2401.9651 s -Epoch 00031 | LR 7.3230e-05 | Loss 1.5669 | Accuracy 0.3991 | Test_Loss 1.6006 | Test_AUC 0.8475 | Time 2397.6024 s -Epoch 00032 | LR 7.2498e-05 | Loss 1.5667 | Accuracy 0.4011 | Test_Loss 1.5956 | Test_AUC 0.8483 | Time 2404.3273 s -Epoch 00033 | LR 7.1773e-05 | Loss 1.5664 | Accuracy 0.3999 | Test_Loss 1.5982 | Test_AUC 0.8468 | Time 2397.5160 s -Epoch 00034 | LR 7.1055e-05 | Loss 1.5662 | Accuracy 0.4012 | Test_Loss 1.5956 | Test_AUC 0.8491 | Time 2396.6267 s -Epoch 00035 | LR 7.0345e-05 | Loss 1.5659 | Accuracy 0.4003 | Test_Loss 1.5962 | Test_AUC 0.8473 | Time 2393.2753 s -Training multiclass_12_process 2024-09-08 07:29:43.909566 -Epoch 00036 | LR 7.0345e-05 | Loss 1.5654 | Accuracy 0.3859 | Test_Loss 1.6487 | Test_AUC 0.8440 | Time 2350.2070 s -Epoch 00037 | LR 6.9641e-05 | Loss 1.5654 | Accuracy 0.3855 | Test_Loss 1.6569 | Test_AUC 0.8450 | Time 2376.8111 s -Epoch 00038 | LR 6.8945e-05 | Loss 1.5652 | Accuracy 0.3876 | Test_Loss 1.6477 | Test_AUC 0.8444 | Time 2388.1118 s -Epoch 00039 | LR 6.8255e-05 | Loss 1.5650 | Accuracy 0.3856 | Test_Loss 1.6473 | Test_AUC 0.8441 | Time 2367.7247 s -Epoch 00040 | LR 6.7573e-05 | Loss 1.5648 | Accuracy 0.3893 | Test_Loss 1.6384 | Test_AUC 0.8448 | Time 2354.1136 s -Epoch 00041 | LR 6.6897e-05 | Loss 1.5646 | Accuracy 0.3861 | Test_Loss 1.6461 | Test_AUC 0.8458 | Time 2340.8834 s -Epoch 00042 | LR 6.6228e-05 | Loss 1.5644 | Accuracy 0.3932 | Test_Loss 1.6217 | Test_AUC 0.8482 | Time 2349.6574 s -Epoch 00043 | LR 6.5566e-05 | Loss 1.5641 | Accuracy 0.4010 | Test_Loss 1.5895 | Test_AUC 0.8484 | Time 2360.5786 s -Epoch 00044 | LR 6.4910e-05 | Loss 1.5638 | Accuracy 0.4019 | Test_Loss 1.5900 | Test_AUC 0.8483 | Time 2468.9202 s -Epoch 00045 | LR 6.4261e-05 | Loss 1.5640 | Accuracy 0.4008 | Test_Loss 1.5990 | Test_AUC 0.8481 | Time 2473.0072 s -Epoch 00046 | LR 6.3619e-05 | Loss 1.5639 | Accuracy 0.4011 | Test_Loss 1.5962 | Test_AUC 0.8498 | Time 2448.4408 s -Epoch 00047 | LR 6.2982e-05 | Loss 1.5638 | Accuracy 0.4017 | Test_Loss 1.5946 | Test_AUC 0.8485 | Time 2381.1478 s -Epoch 00048 | LR 6.2353e-05 | Loss 1.5636 | Accuracy 0.4011 | Test_Loss 1.5942 | Test_AUC 0.8480 | Time 2393.2382 s -Epoch 00049 | LR 6.1729e-05 | Loss 1.5635 | Accuracy 0.4022 | Test_Loss 1.5950 | Test_AUC 0.8482 | Time 2350.8688 s -Epoch 00050 | LR 6.1112e-05 | Loss 1.5634 | Accuracy 0.4007 | Test_Loss 1.5968 | Test_AUC 0.8500 | Time 2334.9762 s -Epoch 00051 | LR 6.0501e-05 | Loss 1.5632 | Accuracy 0.4026 | Test_Loss 1.5919 | Test_AUC 0.8487 | Time 2269.9354 s -Epoch 00052 | LR 5.9896e-05 | Loss 1.5631 | Accuracy 0.4018 | Test_Loss 1.5961 | Test_AUC 0.8493 | Time 2303.6344 s -Epoch 00053 | LR 5.9297e-05 | Loss 1.5630 | Accuracy 0.4022 | Test_Loss 1.5928 | Test_AUC 0.8501 | Time 2275.7180 s -Epoch 00054 | LR 5.8704e-05 | Loss 1.5626 | Accuracy 0.3868 | Test_Loss 1.6539 | Test_AUC 0.8439 | Time 2282.3519 s -Epoch 00055 | LR 5.8117e-05 | Loss 1.5626 | Accuracy 0.3848 | Test_Loss 1.6525 | Test_AUC 0.8449 | Time 2338.0367 s -Epoch 00056 | LR 5.7535e-05 | Loss 1.5625 | Accuracy 0.3861 | Test_Loss 1.6500 | Test_AUC 0.8444 | Time 2349.1735 s -Epoch 00057 | LR 5.6960e-05 | Loss 1.5624 | Accuracy 0.3883 | Test_Loss 1.6474 | Test_AUC 0.8438 | Time 2352.2581 s -Epoch 00058 | LR 5.6391e-05 | Loss 1.5623 | Accuracy 0.3869 | Test_Loss 1.6484 | Test_AUC 0.8453 | Time 2300.8949 s -Epoch 00059 | LR 5.5827e-05 | Loss 1.5621 | Accuracy 0.3885 | Test_Loss 1.6456 | Test_AUC 0.8452 | Time 2296.1956 s -Epoch 00060 | LR 5.5268e-05 | Loss 1.5620 | Accuracy 0.3891 | Test_Loss 1.6347 | Test_AUC 0.8457 | Time 2311.4265 s -Epoch 00061 | LR 5.4716e-05 | Loss 1.5618 | Accuracy 0.4018 | Test_Loss 1.5879 | Test_AUC 0.8495 | Time 2307.2661 s -Epoch 00062 | LR 5.4169e-05 | Loss 1.5615 | Accuracy 0.4029 | Test_Loss 1.5857 | Test_AUC 0.8507 | Time 2356.9390 s -Epoch 00063 | LR 5.3627e-05 | Loss 1.5617 | Accuracy 0.4008 | Test_Loss 1.5969 | Test_AUC 0.8497 | Time 2315.4540 s -Epoch 00064 | LR 5.3091e-05 | Loss 1.5618 | Accuracy 0.4014 | Test_Loss 1.5944 | Test_AUC 0.8490 | Time 2322.3426 s -Epoch 00065 | LR 5.2560e-05 | Loss 1.5618 | Accuracy 0.4013 | Test_Loss 1.5960 | Test_AUC 0.8498 | Time 2337.3654 s -Epoch 00066 | LR 5.2034e-05 | Loss 1.5616 | Accuracy 0.4038 | Test_Loss 1.5905 | Test_AUC 0.8505 | Time 2354.6974 s -Epoch 00067 | LR 5.1514e-05 | Loss 1.5615 | Accuracy 0.4028 | Test_Loss 1.5932 | Test_AUC 0.8491 | Time 2346.2346 s -Epoch 00068 | LR 5.0999e-05 | Loss 1.5615 | Accuracy 0.4026 | Test_Loss 1.5919 | Test_AUC 0.8492 | Time 2364.3581 s -Epoch 00069 | LR 5.0489e-05 | Loss 1.5614 | Accuracy 0.4014 | Test_Loss 1.5947 | Test_AUC 0.8485 | Time 2321.0254 s -Epoch 00070 | LR 4.9984e-05 | Loss 1.5613 | Accuracy 0.4018 | Test_Loss 1.5942 | Test_AUC 0.8504 | Time 2360.4706 s -Epoch 00071 | LR 4.9484e-05 | Loss 1.5612 | Accuracy 0.4016 | Test_Loss 1.5946 | Test_AUC 0.8487 | Time 2348.2632 s diff --git a/Pretrained_GNN/multilabel_pretrained_model_41/config.yaml b/Pretrained_GNN/multilabel_pretrained_model_41/config.yaml deleted file mode 100644 index d5c63f0079c57faf98f9371a339639e66300e5e1..0000000000000000000000000000000000000000 --- a/Pretrained_GNN/multilabel_pretrained_model_41/config.yaml +++ /dev/null @@ -1,569 +0,0 @@ -Datasets: - SingleT_schan: - args: - buffer_size: 11 - chunks: 100 - file_names: SingleT_schan_multilabel.root - fold_var: Number - label: &id001 - - n_higgs - - n_top - - n_V - - n_W - - n_Z - - higgs_exists - - higgs_pt_low - - higgs_pt_med - - higgs_pt_high - - higgs_eta_0 - - higgs_eta_1 - - higgs_eta_2 - - higgs_eta_3 - - higgs_phi_0 - - higgs_phi_1 - - higgs_phi_2 - - higgs_phi_3 - - top0_exists - - top0_pt_low - - top0_pt_med - - top0_pt_high - - top0_eta_0 - - top0_eta_1 - - top0_eta_2 - - top0_eta_3 - - top0_phi_0 - - top0_phi_1 - - top0_phi_2 - - top0_phi_3 - - top1_exists - - top1_pt_low - - top1_pt_med - - top1_pt_high - - top1_eta_0 - - top1_eta_1 - - top1_eta_2 - - top1_eta_3 - - top1_phi_0 - - top1_phi_1 - - top1_phi_2 - - top1_phi_3 - name: SingleT_schan - node_branch_names: &id002 - - - jet_pt - - ele_pt - - mu_pt - - ph_pt - - MET_met - - - jet_eta - - ele_eta - - mu_eta - - ph_eta - - 0 - - - jet_phi - - ele_phi - - mu_phi - - ph_phi - - MET_phi - - CALC_E - - - jet_btag - - 0 - - 0 - - 0 - - 0 - - - 0 - - ele_charge - - mu_charge - - 0 - - 0 - - NODE_TYPE - node_branch_types: &id003 - - vector - - vector - - vector - - vector - - single - node_feature_scales: &id004 - - 1e-1 - - 1 - - 1 - - 1e-1 - - 1 - - 1 - - 1 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: &id005 - n_folds: 10 - test: - - 0 - - 1 - train: - - 2 - - 3 - - 4 - - 5 - - 6 - - 7 - - 8 - - 9 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - VBF: - args: - buffer_size: 11 - chunks: 100 - file_names: VBF_NLO_inc_multilabel.root - fold_var: Number - label: *id001 - name: VBF - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - WH: - args: - buffer_size: 11 - chunks: 100 - file_names: WH_NLO_inc_multilabel.root - fold_var: Number - label: *id001 - name: WH - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ZH: - args: - buffer_size: 11 - chunks: 100 - file_names: ZH_NLO_inc_multilabel.root - fold_var: Number - label: *id001 - name: ZH - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ggF: - args: - buffer_size: 11 - chunks: 100 - file_names: ggF_NLO_inc_multilabel.root - fold_var: Number - label: *id001 - name: ggF - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - tHjb: - args: - buffer_size: 11 - chunks: 100 - file_names: tHjb_NLO_inc_multilabel.root - fold_var: Number - label: *id001 - name: tHjb - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttH: - args: - buffer_size: 11 - chunks: 100 - file_names: ttH_NLO_inc_multilabel.root - fold_var: Number - label: *id001 - name: ttH - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttW: - args: - buffer_size: 11 - chunks: 100 - file_names: ttW_multilabel.root - fold_var: Number - label: *id001 - name: ttW - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttbar: - args: - buffer_size: 11 - chunks: 100 - file_names: ttbar_multilabel.root - fold_var: Number - label: *id001 - name: ttbar - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttt: - args: - buffer_size: 11 - chunks: 100 - file_names: ttt_multilabel.root - fold_var: Number - label: *id001 - name: ttt - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - tttt: - args: - buffer_size: 11 - chunks: 100 - file_names: tttt_multilabel.root - fold_var: Number - label: *id001 - name: tttt - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 - ttyy: - args: - buffer_size: 11 - chunks: 100 - file_names: ttyy_multilabel.root - fold_var: Number - label: *id001 - name: ttyy - node_branch_names: *id002 - node_branch_types: *id003 - node_feature_scales: *id004 - raw_dir: /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/data/ntuples/Hyy_pretraining/multilabel/ - save_dir: /pscratch/sd/j/joshuaho/root_gnn/root_gnn_dgl/data/processed_diverse_multilabel/ - tree_name: output - weight_var: weight - class: LazyMultiLabelDataset - folding: *id005 - module: root_gnn_base.dataset - padding_mode: NONE - shuffle_chunks: 10 -Loss: - args: - label_names: - - n_higgs - - n_top - - n_V - - n_W - - n_Z - - higgs_exists - - higgs_pt_low - - higgs_pt_med - - higgs_pt_high - - higgs_eta_0 - - higgs_eta_1 - - higgs_eta_2 - - higgs_eta_3 - - higgs_phi_0 - - higgs_phi_1 - - higgs_phi_2 - - higgs_phi_3 - - top0_exists - - top0_pt_low - - top0_pt_med - - top0_pt_high - - top0_eta_0 - - top0_eta_1 - - top0_eta_2 - - top0_eta_3 - - top0_phi_0 - - top0_phi_1 - - top0_phi_2 - - top0_phi_3 - - top1_exists - - top1_pt_low - - top1_pt_med - - top1_pt_high - - top1_eta_0 - - top1_eta_1 - - top1_eta_2 - - top1_eta_3 - - top1_phi_0 - - top1_phi_1 - - top1_phi_2 - - top1_phi_3 - label_types: - - r - - r - - r - - r - - r - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - label_weights: - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - - 1 - class: MultiLabelLoss - finish: - args: - label_names: - - n_higgs - - n_top - - n_V - - n_W - - n_Z - - higgs_exists - - higgs_pt_low - - higgs_pt_med - - higgs_pt_high - - higgs_eta_0 - - higgs_eta_1 - - higgs_eta_2 - - higgs_eta_3 - - higgs_phi_0 - - higgs_phi_1 - - higgs_phi_2 - - higgs_phi_3 - - top0_exists - - top0_pt_low - - top0_pt_med - - top0_pt_high - - top0_eta_0 - - top0_eta_1 - - top0_eta_2 - - top0_eta_3 - - top0_phi_0 - - top0_phi_1 - - top0_phi_2 - - top0_phi_3 - - top1_exists - - top1_pt_low - - top1_pt_med - - top1_pt_high - - top1_eta_0 - - top1_eta_1 - - top1_eta_2 - - top1_eta_3 - - top1_phi_0 - - top1_phi_1 - - top1_phi_2 - - top1_phi_3 - label_types: - - r - - r - - r - - r - - r - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - - c - class: MultiLabelFinish - module: models.loss - module: models.loss -Model: - args: - dropout: 0.1 - hid_size: 64 - in_size: 7 - n_layers: 4 - n_proc_steps: 4 - out_size: 41 - class: Edge_Network - module: models.GCN -Training: - batch_size: 1024 - epochs: 100 - gamma: 0.99 - learning_rate: 0.0001 -Training_Directory: trainings/pretraining_multilabel/multilabel_41_higgs_tops_all_kinematics/ -Training_Name: multilabel_41_higgs_tops_all_kinematics diff --git a/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_44.pt b/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_44.pt deleted file mode 100644 index 3eb94c98a2f4e4d31ea9d256e140a21a594f94c3..0000000000000000000000000000000000000000 --- a/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_44.pt +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c3fdf8d12e9c1bece6e22a91b10a6ede5e7ce8103e247b912fbfda580461446b -size 1730158 diff --git a/Pretrained_GNN/multilabel_pretrained_model_41/training.log b/Pretrained_GNN/multilabel_pretrained_model_41/training.log deleted file mode 100644 index ee51d69c2134cc11052c7292eff01c421b771b3a..0000000000000000000000000000000000000000 --- a/Pretrained_GNN/multilabel_pretrained_model_41/training.log +++ /dev/null @@ -1,91 +0,0 @@ -Training multilabel_41_higgs_tops_all_kinematics 2024-08-23 19:20:10.496271 -MultiLabel_Accuracy | 0.7455 | 0.5281 | 0.3708 | 0.3795 | 0.9044 | 0.4787 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.7644 | 0.9852 | 0.5181 | 0.8782 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.7744 | 0.9731 | 0.5523 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00000 | LR 1.0000e-04 | Loss 21.5584 | Accuracy 0.8039 | Test_Loss 29.8678 | Test_AUC 0.0000 | Time 352.4659 s -MultiLabel_Accuracy | 0.7517 | 0.5461 | 0.4684 | 0.4625 | 0.9044 | 0.7094 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8195 | 0.9852 | 0.6451 | 0.8782 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8309 | 0.9731 | 0.7521 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00001 | LR 9.9000e-05 | Loss 16.0182 | Accuracy 0.8252 | Test_Loss 28.8779 | Test_AUC 0.0000 | Time 303.0560 s -MultiLabel_Accuracy | 0.7517 | 0.5593 | 0.4727 | 0.4625 | 0.9044 | 0.7153 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8259 | 0.9852 | 0.7177 | 0.8782 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8330 | 0.9731 | 0.7826 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00002 | LR 9.8010e-05 | Loss 15.3191 | Accuracy 0.8285 | Test_Loss 28.6484 | Test_AUC 0.0000 | Time 288.8502 s -MultiLabel_Accuracy | 0.7560 | 0.5629 | 0.4914 | 0.4665 | 0.9044 | 0.7099 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8284 | 0.9852 | 0.7207 | 0.9055 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8376 | 0.9731 | 0.7980 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00003 | LR 9.7030e-05 | Loss 15.0311 | Accuracy 0.8304 | Test_Loss 28.5576 | Test_AUC 0.0000 | Time 286.0547 s -MultiLabel_Accuracy | 0.7625 | 0.5731 | 0.5011 | 0.4725 | 0.9044 | 0.7103 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8325 | 0.9852 | 0.7283 | 0.9086 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8426 | 0.9731 | 0.8055 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00004 | LR 9.6060e-05 | Loss 14.8461 | Accuracy 0.8319 | Test_Loss 28.4603 | Test_AUC 0.0000 | Time 284.5482 s -MultiLabel_Accuracy | 0.7645 | 0.5736 | 0.5196 | 0.5005 | 0.9044 | 0.7025 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8364 | 0.9852 | 0.7657 | 0.9338 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8469 | 0.9731 | 0.8034 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00005 | LR 9.5099e-05 | Loss 14.7059 | Accuracy 0.8346 | Test_Loss 28.4169 | Test_AUC 0.0000 | Time 281.6354 s -MultiLabel_Accuracy | 0.7648 | 0.5748 | 0.5223 | 0.5069 | 0.9044 | 0.7103 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8350 | 0.9852 | 0.7655 | 0.9281 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8430 | 0.9731 | 0.8036 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00006 | LR 9.4148e-05 | Loss 14.6290 | Accuracy 0.8347 | Test_Loss 28.3604 | Test_AUC 0.0000 | Time 281.7016 s -MultiLabel_Accuracy | 0.7626 | 0.5821 | 0.5289 | 0.5139 | 0.9044 | 0.7184 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8283 | 0.9852 | 0.7678 | 0.9335 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8344 | 0.9731 | 0.7900 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00007 | LR 9.3207e-05 | Loss 14.5512 | Accuracy 0.8349 | Test_Loss 28.3178 | Test_AUC 0.0000 | Time 283.5411 s -MultiLabel_Accuracy | 0.7629 | 0.5778 | 0.5311 | 0.5153 | 0.9044 | 0.7141 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8294 | 0.9852 | 0.7725 | 0.9309 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8272 | 0.9731 | 0.7831 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00008 | LR 9.2274e-05 | Loss 14.4881 | Accuracy 0.8345 | Test_Loss 28.3158 | Test_AUC 0.0000 | Time 301.1452 s -MultiLabel_Accuracy | 0.7641 | 0.5809 | 0.5265 | 0.5111 | 0.9044 | 0.7141 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8289 | 0.9852 | 0.7678 | 0.9249 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8278 | 0.9731 | 0.7936 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00009 | LR 9.1352e-05 | Loss 14.4349 | Accuracy 0.8344 | Test_Loss 28.2953 | Test_AUC 0.0000 | Time 290.1668 s -MultiLabel_Accuracy | 0.7630 | 0.5775 | 0.5313 | 0.5147 | 0.9044 | 0.7133 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8268 | 0.9852 | 0.7712 | 0.9161 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8224 | 0.9731 | 0.7899 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00010 | LR 9.0438e-05 | Loss 14.3933 | Accuracy 0.8341 | Test_Loss 28.2877 | Test_AUC 0.0000 | Time 282.7054 s -MultiLabel_Accuracy | 0.7631 | 0.5796 | 0.5274 | 0.5106 | 0.9044 | 0.7121 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8272 | 0.9852 | 0.7771 | 0.9241 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8227 | 0.9731 | 0.7846 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00011 | LR 8.9534e-05 | Loss 14.3596 | Accuracy 0.8341 | Test_Loss 28.2753 | Test_AUC 0.0000 | Time 281.5382 s -MultiLabel_Accuracy | 0.7640 | 0.5816 | 0.5218 | 0.5071 | 0.9044 | 0.7125 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8278 | 0.9852 | 0.7836 | 0.9215 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8261 | 0.9731 | 0.7864 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00012 | LR 8.8638e-05 | Loss 14.3304 | Accuracy 0.8342 | Test_Loss 28.2619 | Test_AUC 0.0000 | Time 356.4046 s -MultiLabel_Accuracy | 0.7643 | 0.5823 | 0.5164 | 0.5030 | 0.9044 | 0.7158 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8284 | 0.9852 | 0.7898 | 0.9320 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8290 | 0.9731 | 0.7799 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00013 | LR 8.7752e-05 | Loss 14.3099 | Accuracy 0.8344 | Test_Loss 28.2455 | Test_AUC 0.0000 | Time 319.6655 s -MultiLabel_Accuracy | 0.7639 | 0.5792 | 0.5255 | 0.5108 | 0.9044 | 0.7151 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8270 | 0.9852 | 0.7780 | 0.9251 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8258 | 0.9731 | 0.7858 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00014 | LR 8.6875e-05 | Loss 14.2862 | Accuracy 0.8343 | Test_Loss 28.2518 | Test_AUC 0.0000 | Time 285.1502 s -MultiLabel_Accuracy | 0.7647 | 0.5830 | 0.5255 | 0.5100 | 0.9044 | 0.7184 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8260 | 0.9852 | 0.7850 | 0.9223 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8261 | 0.9731 | 0.7829 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00015 | LR 8.6006e-05 | Loss 14.2648 | Accuracy 0.8345 | Test_Loss 28.2293 | Test_AUC 0.0000 | Time 284.2648 s -MultiLabel_Accuracy | 0.7655 | 0.5830 | 0.5284 | 0.5133 | 0.9044 | 0.7184 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8285 | 0.9852 | 0.7904 | 0.9221 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8304 | 0.9731 | 0.7856 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00016 | LR 8.5146e-05 | Loss 14.2550 | Accuracy 0.8350 | Test_Loss 28.2367 | Test_AUC 0.0000 | Time 285.0594 s -MultiLabel_Accuracy | 0.7646 | 0.5796 | 0.5293 | 0.5135 | 0.9044 | 0.7168 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8288 | 0.9852 | 0.7897 | 0.9253 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8290 | 0.9731 | 0.7913 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00017 | LR 8.4294e-05 | Loss 14.2427 | Accuracy 0.8351 | Test_Loss 28.2483 | Test_AUC 0.0000 | Time 285.6626 s -MultiLabel_Accuracy | 0.7650 | 0.5815 | 0.5352 | 0.5199 | 0.9044 | 0.7206 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8246 | 0.9852 | 0.7882 | 0.9255 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8242 | 0.9731 | 0.7825 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00018 | LR 8.3451e-05 | Loss 14.2201 | Accuracy 0.8351 | Test_Loss 28.2372 | Test_AUC 0.0000 | Time 281.2534 s -MultiLabel_Accuracy | 0.7651 | 0.5869 | 0.5295 | 0.5168 | 0.9044 | 0.7228 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8277 | 0.9852 | 0.7926 | 0.9242 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8306 | 0.9731 | 0.7825 | 0.9525 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00019 | LR 8.2617e-05 | Loss 14.2050 | Accuracy 0.8353 | Test_Loss 28.2081 | Test_AUC 0.0000 | Time 280.0656 s -MultiLabel_Accuracy | 0.7646 | 0.5755 | 0.5386 | 0.5236 | 0.9044 | 0.7172 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8265 | 0.9852 | 0.7885 | 0.9215 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8277 | 0.9731 | 0.7928 | 0.9599 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00020 | LR 8.1791e-05 | Loss 14.1920 | Accuracy 0.8355 | Test_Loss 28.2547 | Test_AUC 0.0000 | Time 281.6702 s -MultiLabel_Accuracy | 0.7666 | 0.5848 | 0.5371 | 0.5224 | 0.9044 | 0.7205 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8267 | 0.9852 | 0.7904 | 0.9215 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8255 | 0.9731 | 0.7858 | 0.9597 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00021 | LR 8.0973e-05 | Loss 14.1798 | Accuracy 0.8356 | Test_Loss 28.2137 | Test_AUC 0.0000 | Time 283.8978 s -MultiLabel_Accuracy | 0.7645 | 0.5835 | 0.5386 | 0.5260 | 0.9044 | 0.7215 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8252 | 0.9852 | 0.7857 | 0.9235 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8265 | 0.9731 | 0.7921 | 0.9599 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00022 | LR 8.0163e-05 | Loss 14.1698 | Accuracy 0.8357 | Test_Loss 28.2437 | Test_AUC 0.0000 | Time 284.9271 s -MultiLabel_Accuracy | 0.7659 | 0.5885 | 0.5385 | 0.5258 | 0.9044 | 0.7210 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8287 | 0.9852 | 0.7920 | 0.9255 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8275 | 0.9731 | 0.7913 | 0.9594 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00023 | LR 7.9361e-05 | Loss 14.1528 | Accuracy 0.8361 | Test_Loss 28.1935 | Test_AUC 0.0000 | Time 282.3253 s -MultiLabel_Accuracy | 0.7656 | 0.5849 | 0.5377 | 0.5230 | 0.9044 | 0.7188 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8275 | 0.9852 | 0.7909 | 0.9259 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8289 | 0.9731 | 0.7928 | 0.9598 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00024 | LR 7.8568e-05 | Loss 14.1443 | Accuracy 0.8359 | Test_Loss 28.2195 | Test_AUC 0.0000 | Time 295.7619 s -MultiLabel_Accuracy | 0.7652 | 0.5848 | 0.5384 | 0.5230 | 0.9044 | 0.7158 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8280 | 0.9852 | 0.7879 | 0.9215 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8293 | 0.9731 | 0.7947 | 0.9594 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00025 | LR 7.7782e-05 | Loss 14.1351 | Accuracy 0.8357 | Test_Loss 28.2209 | Test_AUC 0.0000 | Time 305.7035 s -MultiLabel_Accuracy | 0.7657 | 0.5860 | 0.5367 | 0.5201 | 0.9044 | 0.7158 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8297 | 0.9852 | 0.7926 | 0.9228 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8298 | 0.9731 | 0.7950 | 0.9594 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00026 | LR 7.7004e-05 | Loss 14.1285 | Accuracy 0.8359 | Test_Loss 28.2225 | Test_AUC 0.0000 | Time 294.8985 s -MultiLabel_Accuracy | 0.7661 | 0.5872 | 0.5370 | 0.5208 | 0.9044 | 0.7178 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8288 | 0.9852 | 0.7931 | 0.9269 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8294 | 0.9731 | 0.7926 | 0.9601 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00027 | LR 7.6234e-05 | Loss 14.1134 | Accuracy 0.8360 | Test_Loss 28.1849 | Test_AUC 0.0000 | Time 282.3505 s -MultiLabel_Accuracy | 0.7663 | 0.5863 | 0.5395 | 0.5243 | 0.9044 | 0.7188 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8283 | 0.9852 | 0.7875 | 0.9197 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8304 | 0.9731 | 0.7959 | 0.9589 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00028 | LR 7.5472e-05 | Loss 14.1056 | Accuracy 0.8359 | Test_Loss 28.2178 | Test_AUC 0.0000 | Time 286.5912 s -MultiLabel_Accuracy | 0.7660 | 0.5886 | 0.5383 | 0.5231 | 0.9044 | 0.7219 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8268 | 0.9852 | 0.7913 | 0.9258 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8262 | 0.9731 | 0.7927 | 0.9599 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00029 | LR 7.4717e-05 | Loss 14.0994 | Accuracy 0.8361 | Test_Loss 28.2007 | Test_AUC 0.0000 | Time 281.7022 s -MultiLabel_Accuracy | 0.7675 | 0.5887 | 0.5363 | 0.5219 | 0.9044 | 0.7175 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8311 | 0.9852 | 0.7943 | 0.9255 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8295 | 0.9731 | 0.7942 | 0.9597 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00030 | LR 7.3970e-05 | Loss 14.0880 | Accuracy 0.8362 | Test_Loss 28.1905 | Test_AUC 0.0000 | Time 279.7372 s -MultiLabel_Accuracy | 0.7677 | 0.5906 | 0.5366 | 0.5232 | 0.9044 | 0.7194 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8312 | 0.9852 | 0.7940 | 0.9255 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8308 | 0.9731 | 0.7967 | 0.9594 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00031 | LR 7.3230e-05 | Loss 14.0795 | Accuracy 0.8364 | Test_Loss 28.1945 | Test_AUC 0.0000 | Time 285.0377 s -MultiLabel_Accuracy | 0.7687 | 0.5940 | 0.5357 | 0.5232 | 0.9044 | 0.7206 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8299 | 0.9852 | 0.7924 | 0.9242 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8324 | 0.9731 | 0.7960 | 0.9596 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00032 | LR 7.2498e-05 | Loss 14.0707 | Accuracy 0.8364 | Test_Loss 28.1790 | Test_AUC 0.0000 | Time 282.6988 s -MultiLabel_Accuracy | 0.7673 | 0.5884 | 0.5350 | 0.5214 | 0.9044 | 0.7151 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8321 | 0.9852 | 0.7913 | 0.9197 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8279 | 0.9731 | 0.7922 | 0.9590 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00033 | LR 7.1773e-05 | Loss 14.0643 | Accuracy 0.8358 | Test_Loss 28.1895 | Test_AUC 0.0000 | Time 278.8458 s -MultiLabel_Accuracy | 0.7681 | 0.5947 | 0.5338 | 0.5211 | 0.9044 | 0.7213 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8298 | 0.9852 | 0.7941 | 0.9262 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8269 | 0.9731 | 0.7919 | 0.9600 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00034 | LR 7.1055e-05 | Loss 14.0574 | Accuracy 0.8362 | Test_Loss 28.1609 | Test_AUC 0.0000 | Time 283.4880 s -MultiLabel_Accuracy | 0.7690 | 0.5982 | 0.5342 | 0.5229 | 0.9044 | 0.7235 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8299 | 0.9852 | 0.7948 | 0.9273 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8274 | 0.9731 | 0.7912 | 0.9601 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00035 | LR 7.0345e-05 | Loss 14.0464 | Accuracy 0.8365 | Test_Loss 28.1496 | Test_AUC 0.0000 | Time 285.9291 s -MultiLabel_Accuracy | 0.7684 | 0.5962 | 0.5339 | 0.5212 | 0.9044 | 0.7213 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8349 | 0.9852 | 0.7960 | 0.9255 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8302 | 0.9731 | 0.7935 | 0.9597 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00036 | LR 6.9641e-05 | Loss 14.0393 | Accuracy 0.8365 | Test_Loss 28.1716 | Test_AUC 0.0000 | Time 284.9658 s -MultiLabel_Accuracy | 0.7686 | 0.5963 | 0.5337 | 0.5208 | 0.9044 | 0.7184 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8339 | 0.9852 | 0.7956 | 0.9259 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8305 | 0.9731 | 0.7932 | 0.9596 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00037 | LR 6.8945e-05 | Loss 14.0344 | Accuracy 0.8364 | Test_Loss 28.1794 | Test_AUC 0.0000 | Time 289.9682 s -MultiLabel_Accuracy | 0.7692 | 0.5996 | 0.5330 | 0.5190 | 0.9044 | 0.7203 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8346 | 0.9852 | 0.7971 | 0.9260 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8321 | 0.9731 | 0.7949 | 0.9598 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00038 | LR 6.8255e-05 | Loss 14.0218 | Accuracy 0.8367 | Test_Loss 28.1637 | Test_AUC 0.0000 | Time 287.4777 s -MultiLabel_Accuracy | 0.7691 | 0.5970 | 0.5330 | 0.5201 | 0.9044 | 0.7190 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8345 | 0.9852 | 0.7968 | 0.9276 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8342 | 0.9731 | 0.7962 | 0.9601 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00039 | LR 6.7573e-05 | Loss 14.0097 | Accuracy 0.8367 | Test_Loss 28.1543 | Test_AUC 0.0000 | Time 290.9069 s -MultiLabel_Accuracy | 0.7694 | 0.5998 | 0.5333 | 0.5203 | 0.9044 | 0.7209 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8350 | 0.9852 | 0.7975 | 0.9277 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8353 | 0.9731 | 0.7976 | 0.9601 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00040 | LR 6.6897e-05 | Loss 13.9961 | Accuracy 0.8369 | Test_Loss 28.1429 | Test_AUC 0.0000 | Time 284.9635 s -MultiLabel_Accuracy | 0.7692 | 0.5994 | 0.5328 | 0.5191 | 0.9044 | 0.7226 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8346 | 0.9852 | 0.7977 | 0.9274 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8365 | 0.9731 | 0.7987 | 0.9602 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00041 | LR 6.6228e-05 | Loss 13.9844 | Accuracy 0.8370 | Test_Loss 28.1354 | Test_AUC 0.0000 | Time 283.8730 s -MultiLabel_Accuracy | 0.7694 | 0.6001 | 0.5347 | 0.5204 | 0.9044 | 0.7256 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8345 | 0.9852 | 0.7974 | 0.9282 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8360 | 0.9731 | 0.7990 | 0.9599 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00042 | LR 6.5566e-05 | Loss 13.9704 | Accuracy 0.8371 | Test_Loss 28.1307 | Test_AUC 0.0000 | Time 284.3014 s -MultiLabel_Accuracy | 0.7695 | 0.5988 | 0.5348 | 0.5199 | 0.9044 | 0.7225 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8366 | 0.9852 | 0.7988 | 0.9278 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8378 | 0.9731 | 0.8020 | 0.9602 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00043 | LR 6.4910e-05 | Loss 13.9571 | Accuracy 0.8372 | Test_Loss 28.1316 | Test_AUC 0.0000 | Time 282.0029 s -MultiLabel_Accuracy | 0.7695 | 0.5994 | 0.5348 | 0.5192 | 0.9044 | 0.7233 | 0.8561 | 0.6573 | 0.9652 | 0.8321 | 0.9076 | 0.9074 | 0.8316 | 0.8933 | 0.8935 | 0.7987 | 0.8932 | 0.8354 | 0.9852 | 0.7979 | 0.9274 | 0.8829 | 0.8077 | 0.8077 | 0.8832 | 0.8453 | 0.8454 | 0.8456 | 0.8453 | 0.8359 | 0.9731 | 0.8007 | 0.9602 | 0.8843 | 0.8548 | 0.8545 | 0.8843 | 0.8695 | 0.8695 | 0.8696 | 0.8693 -Epoch 00044 | LR 6.4261e-05 | Loss 13.9445 | Accuracy 0.8371 | Test_Loss 28.1224 | Test_AUC 0.0000 | Time 282.0396 s diff --git a/README.md b/README.md index df7cb3370d911c78badcfe5ef5eb6ec03996cbee..7be5fc7f47d5db027d120b8024982df93db95b74 100644 --- a/README.md +++ b/README.md @@ -1,1351 +1,3 @@ ---- -library_name: gnn4colliders -tags: -- graph-neural-network -- high-energy-physics -- particle-physics -- event-classification -- pytorch -datasets: -- HWresearch/Delphes ---- - -# Pretrained Event Classification Model for High Energy Physics Analysis - -Joshua Ho, Ryan Roberts, Shuo Han, and Haichen Wang - -Department of Physics, University of California, Berkeley, CA 94720; Physics -Division, Lawrence Berkeley National Laboratory, Berkeley, CA 94720 - -Code: [Hugging Face model repository](https://huggingface.co/HWresearch/GNN4Colliders) - -## Abstract - -We introduce a foundation model for event classification in high-energy physics, built on a Graph Neural Network architecture and trained on 120 million simulated proton-proton collision events spanning 12 distinct physics processes. The model is pretrained to learn a general and robust representation of collision data using challenging multiclass and multilabel classification tasks. Its performance is evaluated across seven event classification tasks, which include new physics processes not encountered during pretraining as well as ATLAS Open Data to demonstrate generalizability across different simulation frameworks, from Delphes fast simulation to full ATLAS detector simulation. Fine-tuning the pretrained model significantly improves classification performance, particularly in scenarios with limited training data, demonstrating gains in both accuracy and computational efficiency. To investigate the underlying mechanisms behind these performance improvements, we employ a representational similarity evaluation framework based on Centered Kernel Alignment. This analysis reveals that encoder-stage representations of the fine-tuned model remain similar to those of the baseline, while intermediate graph processing layers diverge substantially, indicating that fine-tuning preserves general-purpose encoders while developing fundamentally different message-passing pathways to arrive at superior task performance. - - -# Introduction - -Machine learning has become a ubiquitous tool in particle physics, employed in a variety of tasks including triggering, simulation, reconstruction, and offline analysis. While its utility spans classification, regression, and generative tasks, the current paradigm of developing machine learning models from scratch for each specific application presents several challenges. This approach not only demands specialized expertise and substantial computing resources but can also result in suboptimal performance due to limited training data. The from-scratch development of models necessitates individual validation studies to ensure that neural networks utilize well-modeled information from training samples, whether derived from Monte Carlo simulations or control samples from experimental data. - -Foundation models offer a promising direction to address these limitations. These models, pretrained on large, diverse datasets across various tasks, provide robust and general representations of underlying data structures. Notable examples in other fields include GPT-4 (OpenAI et al. 2024) and BERT (Devlin et al. 2019) in natural language processing, Stable Diffusion (Rombach et al. 2022; Podell et al. 2023) in image processing, and AlphaFold (Jumper et al. 2021) in structural biology. The foundation model approach offers several advantages for particle physics applications: reduced computing resources for fine-tuning (Yosinski et al. 2014) compared to training from scratch, superior performance on specific tasks (particularly with limited training data), and potentially simplified validation procedures as the pretrained representations have already been studied on a broad set of physics processes, providing a foundation that may inform but does not replace task-specific validation. - -Current literature on pretrained models for particle physics can be categorized based on the data representation they handle. Models operating on particle- or event-level numerical data use features like particle four momenta or jets, leveraging self-supervised or generative methods to learn versatile representations. Detector-focused model operates on high-dimensional responses such as calorimeter deposits or pixel hits, employing geometry-aware techniques for accurate simulation and analysis. Finally, models using textual or code representations apply large language model architectures to integrate domain knowledge, enabling tasks like question answering and code generation. - -Recent studies have begun exploring foundation models tailored to particle physics data, which has a variety of distinct structures and properties across many experiments and data processing stages, including particle-level & event-level numeric data  (Wildridge et al. 2024; Katel et al. 2024; Golling et al. 2024; Mikuni and Nachman 2024; Harris et al. 2024; Birk et al. 2024; Vigl et al. 2024), detector-level & geometry-aware data  (Araz et al. 2024; Liu et al. 2023; Hashemi et al. 2024; Huang et al. 2024), and textual or code data  (Zhang et al. 2024). - -While recent work has demonstrated the promise of foundation models in high-energy physics, existing approaches differ from ours in important ways. The majority of prior foundation models operate at the jet level, treating individual jets as the primary object of study. MPM (Leigh et al. 2024) and OmniJet-\\(\alpha\\) (Birk et al. 2024) focus on learning representations of jet constituents using masked modeling and autoregressive generation respectively, while OmniLearn (Mikuni and Nachman 2024) demonstrates strong cross-detector and cross-task generalization across a range of final states. RS3L (Harris et al. 2024) proposes a self-supervised contrastive learning strategy at the jet level using re-simulation-based data augmentation. The work of Vigl et al. (Vigl et al. 2024) demonstrates the value of pretraining and fine-tuning workflows for joint reconstruction and analysis optimization in HEP, but focuses on a specific search topology rather than developing a general-purpose event-level representation. Bumblebee (Wildridge et al. 2024) is perhaps the closest in spirit to our approach, employing a BERT-inspired architecture at the event level by embedding particle four-vectors with permutation invariance. However, Bumblebee uses a Transformer-based architecture and a generative pretraining objective. In contrast, our model uses a Graph Neural Network architecture that explicitly encodes relational structure between particles in the event, and is pretrained on a substantially larger and more diverse dataset of 120 million events spanning 12 physics processes using discriminative multiclass and multilabel objectives. These distinctions make our approach complementary to existing jet-level foundation models, addressing the event-level classification setting that prior work has not systematically explored. We further evaluate generalizability across simulation frameworks, from Delphes fast simulation to full ATLAS detector simulation, a setting not addressed by prior work. - -This paper presents a foundation model designed specifically for collider event-level data. In modern collider experiments, final-stage analysis processes information from reconstructed objects that either directly correspond to particles in collision final states (such as leptons and photons) or serve as proxies (such as jets and missing transverse energy). While traditional approaches often relied on “high-level” variables calculated from object features, recent trends favor direct input of event objects and their features into neural networks for analysis tasks. A notable example is Ref. (ATLAS Collaboration 2023b), which established the observation of simultaneous production of four top quarks with the ATLAS experiment by employing a graph neural network (GNN) architecture to process event-level object information. - -We present foundation models that adopt an architecture similar to that used in Ref. (ATLAS Collaboration 2023b). Our models are pretrained using either multiclass classification or multilabel learning tasks across 12 distinct physics processes and 120 million simulated proton-proton collision events. We evaluate these models through fine-tuning and testing on seven classification tasks, including novel processes not seen during pretraining as well as ATLAS Open Data, the latter serving to demonstrate generalizability across simulation frameworks. Our analysis benchmarks the models’ performance improvements, their scaling behavior with training sample size, and computational efficiency. To investigate the mechanisms underlying these improvements, we employ a representational similarity framework based on Centered Kernel Alignment, revealing that pretrained and scratch-trained models develop similar intermediate representations while diverging at the task-specific output level. This work represents the first prototype of a foundation model operating on collider final-state object data. - -We use the term “foundation model” in the sense of a model trained on broad and diverse data at scale and adapted to a wide range of downstream tasks, without requiring very large model capacity or scaling-driven emergent behavior. In the High Energy Physics (HEP) analysis ecosystem, where task-specific classifiers are typically trained from scratch on limited simulated samples, a model pretrained on 120 million events spanning 12 physics processes and successfully adapted to seven downstream tasks across multiple simulation frameworks represents a qualitatively different paradigm, and we argue this paradigm is appropriately characterized as a foundation model approach in this context. - -# Data Samples - -To provide a diverse set of physics processes for the pretraining, we use Madgraph@NLO 2.7.3 (Alwall et al. 2014) to generate proton-proton collision events at next-to-leading order (NLO) in Quantum Chromodynamics (QCD). We generate 12 distinct Standard Model (SM) physics processes, including six major Higgs boson production mechanisms: gluon fusion production (\\(ggF\\)), vector boson fusion (\\(VBF\\)), associated production of the Higgs boson with a W boson (\\(WH\\)) or a Z boson (\\(ZH\\)), associated production of the Higgs boson with a top-quark pair (\\(t\bar{t}H\\)), and associated production of the Higgs boson with a single top quark and a forward quark (\\(tHq\\)). Additionally, we simulate six top quark production processes: single top production, top-quark pair production (\\(t\bar{t}\\)), top quark pair production in association with a pair of photons (\\(t\bar{t}\gamma\gamma\\)), associated production of a top-quark pair with a W boson (\\(t\bar{t}W\\)), simultaneous production of three top quarks (\\(t\bar{t}t\\)), and simultaneous production of four top quarks (\\(t\bar{t}t\bar{t}\\)). In these samples, the Higgs boson and top quarks decay inclusively. These 12 Higgs and top quark production processes constitute the pretraining dataset. - -To test the pretrained model, we further generated four processes including three beyond Standard Model (SM) processes: a SM \\(t\bar{t}H\\) production where the Higgs boson decays exclusively to a pair of photons, a \\(t\bar{t}H\\) production with the Higgs boson decaying to a pair of photons, where the top-Yukawa coupling is CP-odd, implemented using the Higgs Characterization model (P. Artoisenet et al. 2013), the production of a pair of superpartners of the top quark (s-top) using the Minimal Supersymmetric Standard Model (MSSM) (Rosiek 1990; Allanach et al. 2009), and flavor changing neutral current (FCNC) processes (Degrande et al. 2015; Durieux et al. 2015). For the s-top process, we simulate the production of heavier s-top pairs (\\(t_2\bar{t_2}\\)), where each heavier s-top (mass 582 GeV) decays into a lighter s-top (\\(t_1\\) or \\(\bar{t_1}\\), mass 400 GeV) and a Higgs boson. The FCNC process involves \\(t\bar{t}\\) production where one top quark decays to a Higgs boson and a light quark. We generate 10 million events for each process, except for \\(tHq\\) and \\(t\bar{t}t\bar{t}\\), where 5 million events were produced. - -In all simulation samples, the center of mass energy of the proton-proton collision is set to 13 TeV. The Higgs boson, top quarks, and vector bosons are set to decay inclusively (except the \\(t\bar{t}H \rightarrow \gamma\gamma\\) samples), with MadSpin (Pierre Artoisenet et al. 2013) handling the decays of top quarks and W bosons. The generated events are processed through Pythia 8.235 (Sjöstrand et al. 2015) for parton showering and heavy particle decays, followed by Delphes 3.4.2 (Favereau et al. 2014) configured to emulate the ATLAS detector (Collaboration 2008) for fast detector simulation. - -The detector-level object selection criteria are defined to align with typical experimental conditions. Photons are required to have transverse momentum \\(p_T \geq 20\\) GeV and pseudorapidity \\(|\eta| \leq 2.37\\), excluding the electromagnetic calorimeter crack region (\\(1.37 < |\eta| < 1.52\\)). Electrons must have \\(p_T \geq 10\\) GeV and \\(|\eta| \leq 2.47\\) (excluding the same crack region), while muons are selected with \\(p_T \geq 10\\) GeV and \\(|\eta| \leq 2.7\\). Jets are reconstructed using the anti-\\(k_t\\) algorithm (Cacciari et al. 2008) with radius parameter \\(\Delta R=0.4\\), where \\(\Delta R\\) is defined as \\(\sqrt{\Delta\eta ^2 + \Delta\phi^2}\\), with \\(\Delta\eta\\) being the difference in pseudorapidity and \\(\Delta\phi\\) the difference in azimuthal angle. Jets must satisfy \\(p_T \geq 25\\) GeV and \\(|\eta| \leq 2.5\\). To avoid double-counting, jets are removed if they are within \\(\Delta R < 0.4\\) of a photon or lepton. The identification of jets originating from b-quark decays (b-tagging) is performed by matching jets within \\(\Delta R = 0.4\\) of a b-quark, with efficiency corrections applied to match the performance of the ATLAS experiment’s b-tagging algorithm (Aad et al. 2019). - -In addition to the simulated samples described above, we make use of the ATLAS Open Data release at 13 TeV (ATLAS Collaboration 2025), which provides proton-proton collision data collected by the ATLAS detector during 2015 and 2016, corresponding to an integrated luminosity of 36 \\(\text{fb}^{-1}\\). Unlike the Delphes-based fast simulation used in our pretraining and other fine-tuning samples, the ATLAS Open Data are processed through the full ATLAS detector simulation and reconstruction chain, making this dataset a stringent test of the generalizability of our pretrained models across different simulation frameworks. We used two event collections from this release: the `GamGam` collection, which requires at least two photons with \\(p_T \geq 25~\text{GeV}\\) and is enriched in Higgs boson decays to diphotons, containing gluon fusion (\\(ggF\\)), vector boson fusion (\\(VBF\\)), and associated Higgs production (\\(WH\\), \\(ZH\\), and \\(t\bar{t}H\\)) processes; and the `1LMET30` collection, which requires at least one lepton with \\(p_T \geq 7~\text{GeV}\\) and missing transverse energy of at least \\(30~\text{GeV}\\), and is enriched with leptonically-decaying W bosons, containing triboson production processes (\\(WWW\\), \\(ZZZ\\), and other mixed triboson production \\(WWZ\\) and \\(WZZ\\)). - -The ATLAS Open Data triboson samples used in this analysis differ from our Delphes-based samples both at the event-generation and detector-simulation levels. Whereas our pretraining and Delphes fine-tuning samples are generated with MadGraph@NLO interfaced to Pythia and processed with Delphes fast detector simulation (Alwall et al. 2014; Sjöstrand et al. 2015; Favereau et al. 2014), the ATLAS triboson samples used here are generated with Sherpa 2.2.2 and processed through the ATLAS detector simulation and reconstruction chain, with the detector response modeled using Geant4 (Bothmann et al. 2019; ATLAS Collaboration 2010; Agostinelli et al. 2003). Therefore, the ATLAS Open Data triboson task probes not only a different detector simulation and reconstruction chain, but also a different event-generation and matrix-element setup from the MadGraph+Pythia+Delphes pipeline used for the Delphes samples. - -# Methods - -## Overview - -We present a methodology for developing and evaluating a foundation model for particle collision event analysis. The approach centers on pretraining a Graph Neural Network (GNN) architecture using a comprehensive dataset that spans multiple physics tasks, enabling the model to learn robust and transferable features. For task-specific applications, we employ a fine-tuning strategy that combines output layer adaptation with carefully calibrated learning rates for updating the pretrained parameters. - -Given the prevalence of classification problems in particle physics data analysis, we evaluate the model’s efficacy through a systematic assessment across five binary classification tasks and two multiclass classification tasks: - -- \\(t\bar{t}H(\rightarrow \gamma\gamma)\\) with CP-even versus CP-odd t-H interaction (Delphes) - -- \\(t\bar{t}\\) with FCNC top quark decays versus \\(tHq\\) processes (Delphes) - -- \\(t\bar{t}W\\) versus \\(t\bar{t}t\\) processes (Delphes) - -- s-top pair production with Higgs bosons in the decay chain versus \\(t\bar{t}H\\) processes (Delphes) - -- \\(WH\\) versus \\(ZH\\) production modes (Delphes) - -- 5-class multiclass classification of Higgs production modes where the Higgs decays exclusively to two photons: \\(t\bar{t}H\\), \\(ggF\\), \\(VBF\\), \\(WH\\), and \\(ZH\\) (ATLAS Open Data) - -- 3-class multiclass classification of triboson events: \\(WWW\\), \\(ZZZ\\) and Mixed (\\(WWZ\\) and \\(WZZ\\)) (ATLAS Open Data) - -Our evaluation metrics encompass classification performance, computational efficiency, and model interpretability. The investigation extends to analyzing the model’s scaling behavior with respect to training dataset size, benchmarked against models trained without pretraining. Although we explored transfer learning through parameter freezing of pretrained layers, this approach did not yield performance improvements, leading us to focus our detailed analysis on fine-tuning strategies. - -This methodological framework demonstrates the potential of foundation models to enhance the efficiency of particle physics analyses while improving task-specific performance, offering a promising direction for future high-energy physics research. - -## GNN Architecture - -We implement a Graph Neural Network (GNN) architecture that naturally accommodates the point-cloud structure of particle physics data, employing the DGL framework with a PyTorch backend (Wang et al. 2019; Paszke et al. 2019). The GNN naturally handles graphs of varying node and edge counts through the message-passing framework, without requiring padding or truncation, making it well suited to collision events where the number of reconstructed objects varies from event to event. A fully connected graph is constructed for each event, with nodes corresponding to reconstructed jets, electrons, muons, photons, and the missing transverse-energy vector `E⃗_T^miss`. The features of each node include the four-momentum `(pT, η, ϕ, E)` of the object with a massless assumption (`E = pT cosh η`), the b-tagging label (for jets), the charge (for leptons), and an integer labeling the type of object represented by the node. We use a placeholder value of 0 for features which are not defined for every node type such as the b-jet tag, lepton charge, or the pseudorapidity of `E⃗_T^miss`. An explicit masking mechanism was not tested in this work and represents a potential avenue for future improvement. We assign the angular distances `(Δη, Δϕ, ΔR)` as edge features and the number of nodes `N` in the graph as a global feature. We denote the node features `{x⃗_i}`, edge features `{y⃗_ij}`, and global features `{z⃗}`. - -The GNN model is based on the graph network architecture described in (Battaglia et al. 2018) using simple multilayer perceptron (MLP) feature functions and summation aggregation. The model is comprised of three primary components: an encoder, the graph network, and a decoder. In the encoder, three MLPs embed the nodes, edges, and global features into a latent space of dimension 64. The graph network block, which is designed to facilitate message passing between different domains of the graph, performs an edge update \\(f_e\\), followed by a node update \\(f_n\\), and finally a global update \\(f_g\\), all defined below. The inputs to each update MLP are concatenated. - -$$ -\vec {y'}_{ij} = f_e\left(\{\vec x_k\},\vec y_{ij},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\vec x_j,\vec y_{ij},\vec z\right) -$$ - -$$ -\vec{x'}_{i} = f_n\left(\vec x_i,\{\vec{y'}_{jk}\},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\sum_j\vec{y'}_{ij},\vec z\right) -$$ - -$$ -\vec{z'} = f_g\left(\{\vec{x'}_i\},\{\vec{y'}_{ij}\},\vec z\right) = \mathrm{MLP}\left(\sum_i\vec{x'}_i,\sum_{i,j}\vec{y'}_{ij},\vec z\right) -$$ - -This graph block is iterated four times with the same update MLPs. Finally, the global features are passed through a decoder MLP and a final linear layer to produce the desired model outputs. Each MLP consists of 4 linear layers, each with an output width of 64, with the `ReLU` activation function. The output of the MLP is then passed through a `LayerNorm` layer(Ba et al. 2016). The total number of trainable parameters in this model is about 400,000. - -As a performance benchmark, a baseline GNN model is trained from scratch for each classification task. The Adam optimizer is used, with an initial learning rate of \\(10^{-4}\\) and an exponential decay schedule given by \\(\mathrm{LR}(x) = \mathrm{LR}_{\mathrm{initial}}\cdot(0.99)^x\\), where \\(x\\) denotes the epoch number. - -The baseline hyperparameters of the GNN are summarized in Table 1. - -| **Hyperparameter** | **Value** | -|:---------------------------|:-----------:| -| Hidden Layer Size | 64 | -| MLP Layers | 4 | -| Graph Processing Steps | 4 | -| Batch Size | 1024 | -| Optimizer | Adam | -| Initial Learning Rate | \\(10^{-4}\\) | -| Learning Rate Decay Factor | \\(0.99\\) | - -Baseline GNN Hyperparameters. - -A total of 400 GPU hours on a single NVIDIA A100 were invested for systematic hyperparameter optimization, with full results reported in Appendix 6. We perform a sweep over batch size, learning rate, MLP depth and width, and hidden dimension. These results show that the GNN architecture exhibits stability across hyperparameter configurations: excluding a single outlier, variations in the test AUC remain within \\(\pm 0.2\%\\). - -## Pretraining Strategy - -We explore two complementary pretraining approaches to develop robust representations of collision events: (1) multiclass classification, which trains the model to distinguish between different physics processes, and (2) multilabel classification, which predicts the existence and kinematics of heavy particles with prompt decays. The pretraining dataset consists of approximately 120 million events, evenly distributed across 12 distinct physics processes, including all major Higgs boson production mechanisms and top quark processes as described in Section 2. This large-scale pretraining effort was conducted on the Perlmutter supercomputer at NERSC. - -### Multiclass Classification - -For Monte Carlo simulated events, the underlying physics process that generated each event is known precisely, providing natural labels for supervised learning. However, the challenge lies in the complexity of collision events: different physics processes can produce similar kinematics and event topologies, particularly in certain regions of phase space. No single observable can unambiguously identify the underlying process. By training the model to distinguish between 12 different processes simultaneously, we challenge it to learn subtle differences in kinematics and topology that collectively characterize each process. The model is trained using categorical cross entropy as the loss function. The output layer of the multiclass classification model has 832 trainable parameters. - -### Multilabel Classification - -This approach combines both classification and regression tasks to characterize collision events. For discrete properties like particle presence in specific kinematic regions, we employ classification labels with binary cross-entropy loss. For continuous quantities like particle multiplicities, we use regression labels with mean-squared error loss. This hybrid approach enables the model to learn both categorical and continuous aspects of the physics processes simultaneously. - -We develop a comprehensive set of 41 labels that capture both particle multiplicities and kinematic properties. This approach increases prediction granularity and enhances model interpretability. By training the model to predict event kinematics rather than event identification, we create a task-independent framework that can potentially generalize better to novel scenarios not seen during pretraining. - -The particle multiplicity labels count the number of Higgs bosons (`n_higgs`), top quarks (`n_tops`), vector bosons (`n_V`), W bosons (`n_W`), and Z bosons (`n_Z`). The kinematic labels characterize the transverse momentum (`pT`), pseudorapidity (`η`), and azimuthal angle (`ϕ`) of Higgs bosons and top quarks through binned classifications. - -For Higgs bosons, \\(p_T\\) is categorized into three ranges: (0, 30) GeV, (30, 200) GeV, and (200, \\(\infty\\)) GeV, with the upper range particularly sensitive to potential BSM effects. Similarly, both leading and subleading top quarks have \\(p_T\\) classifications spanning (0, 30) GeV, (30, 300) GeV, and (300, \\(\infty\\)) GeV. When no particle exists within a specific \\(p_T\\) range, the corresponding label is set to \\([0, 0, 0]\\). For all particles, \\(\eta\\) measurements are divided into 4 bins with boundaries at \\([-1.5, 0, 1.5]\\), while \\(\phi\\) measurements use 4 bins with boundaries at \\([-\frac{\pi}{2}, 0, \frac{\pi}{2}]\\). As with \\(p_T\\), both \\(\eta\\) and \\(\phi\\) labels default to \\([0, 0, 0, 0]\\) in the absence of a particle. This comprehensive labeling schema enables fine-grained learning of kinematic distributions and particle multiplicities, essential for characterizing complex collision events. - -The loss function combines individual losses from all 41 labels through weighted averaging. Binary cross-entropy is applied to classification labels, while mean-squared error is used for regression labels. The model generates predictions for all labels simultaneously, with individual losses calculated according to their respective types. The final loss is computed as an equally-weighted average across all labels, with weights set to 1 to ensure uniform contribution to the optimization process. The output layer of the multilabel model has 2,688 trainable parameters. - -### Pretraining - -During pretraining, the initial learning rate is \\(10^{-4}\\), and the learning rate decays by 1% each epoch following the exponential decay function \\(LR(x) = 10^{-4}\cdot(0.99)^x\\), where \\(x\\) is the number of epochs. The model is trained until the training loss shows no improvement for \\(5\\) consecutive epochs, indicating convergence in the training. The epoch with the maximal validation AUC is then selected for fine-tuning. - -## Fine-tuning Methodology - -For downstream tasks, we adjust the model architecture for fine-tuning by replacing the original output layer (final linear layer) with a newly initialized linear layer while retaining the pretrained weights for all other layers. This modification allows the model to specialize in the specific downstream task while leveraging the general features learned during pretraining. - -The fine-tuning process begins with distinct learning rate setups for different parts of the model. The newly initialized linear layer is trained with an initial learning rate of \\(10^{-4}\\), matching the rate used for models trained from scratch. Meanwhile, the pretrained layers are fine-tuned more cautiously with a lower initial learning rate of \\(10^{-5}\\). This approach ensures that the pretrained layers adapt gradually without losing their general features, while the new layer learns effectively from scratch. Both learning rates decay over time following the same exponential decay function, \\(LR(x) = LR_{initial} \cdot (0.99)^x\\), to promote stable convergence as training progresses. - -We also evaluated a transfer learning setup in which either the decoder MLP or the final linear layer was replaced with a newly initialized component. During this process, all other model parameters remained frozen, leveraging the pretrained features without further updating them. However, we did not observe performance improvements using the transfer learning setup. Consequently, we focus on reporting results obtained with the fine-tuning approach. - -## Performance Evaluation - -We assess model performance using two figures of merit: the classification accuracy and the Area Under the Curve (AUC) of the Receiver Operating Characteristic (ROC) curve. For binary tasks, accuracy is computed after applying a threshold of 0.5 to the neural-network output score. For multiclass tasks, the predicted class is taken to be the class with the largest output score. Both metrics demonstrate consistent trends in our analysis. - -We employ an ensemble training approach where 5 independent models are trained for each configuration with random weight initialization and random subsets of the training dataset. This enables us to evaluate both the model’s sensitivity to initial parameters and to quantify uncertainties in their performance. All models are trained to the same stopping condition: \\(5\\) consecutive epochs where there are no improvements to the training loss, and the epoch with maximal validation AUC is then selected for further performance evaluation using the held-out test set. To investigate how model performance scales with training data, we conduct training runs using sample sizes ranging from \\(10^3\\) to \\(10^7\\) events per class (\\(10^3\\), \\(10^4\\), \\(10^5\\), \\(10^6\\), and \\(10^7\\)) for each model setup: the from-scratch baseline and models fine-tuned from multiclass or multilabel pretrained models. For the \\(10^7\\) Delphes case and the ATLAS Open Data tasks, only the initialization was randomized due to dataset size limitations. - -Additional performance metrics, including F1, precision, and recall, are reported in Appendix 7 and show trends consistent with the accuracy and AUC. - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Name of TaskMetricPretrainingSample Size
4-8103104105106107
tt̄H( → γγ) CP-even vs CP-oddAccuracy (%)Baseline56.4 ± 1.062.3 ± 0.164.2 ± 0.165.7 ± 0.066.2 ± 0.0
Multiclass+2.9 ± 1.1+2.0 ± 0.1+1.0 ± 0.1+0.1 ± 0.0-0.0 ± 0.0
Multilabel+1.3 ± 1.2+1.0 ± 0.2+0.6 ± 0.1+0.0 ± 0.0-0.1 ± 0.0
AUC (x 100)Baseline59.1 ± 1.466.9 ± 0.269.5 ± 0.071.2 ± 0.071.9 ± 0.0
Multiclass+3.9 ± 1.4+2.5 ± 0.2+1.0 ± 0.0+0.1 ± 0.0-0.0 ± 0.0
Multilabel+0.6 ± 1.6+1.1 ± 0.2+0.6 ± 0.1-0.0 ± 0.0-0.1 ± 0.0
FCNC vs tHqAccuracy (%)Baseline62.2 ± 0.667.9 ± 0.468.0 ± 0.469.2 ± 0.267.7 ± 0.1
Multiclass+4.6 ± 0.8+0.7 ± 0.4+1.4 ± 0.5+0.4 ± 0.2-0.2 ± 0.1
Multilabel-1.9 ± 0.8-1.1 ± 0.5+1.0 ± 0.5-0.0 ± 0.3+0.0 ± 0.1
AUC (x 100)Baseline67.6 ± 0.572.1 ± 0.074.3 ± 0.075.6 ± 0.073.8 ± 0.0
Multiclass+2.5 ± 1.0+2.6 ± 0.1+1.4 ± 0.1+0.6 ± 0.0-0.0 ± 0.0
Multilabel-3.4 ± 1.0-0.3 ± 0.1+0.4 ± 0.0+0.0 ± 0.0-0.1 ± 0.0
tt̄W vs tt̄tAccuracy (%)Baseline75.7 ± 0.177.5 ± 0.179.0 ± 0.079.8 ± 0.080.3 ± 0.0
Multiclass+2.7 ± 0.2+2.2 ± 0.1+1.0 ± 0.0+0.4 ± 0.0+0.0 ± 0.0
Multilabel+1.6 ± 0.1+0.5 ± 0.2+0.3 ± 0.0+0.0 ± 0.0-0.1 ± 0.0
AUC (x 100)Baseline83.3 ± 0.185.5 ± 0.187.2 ± 0.088.0 ± 0.088.5 ± 0.0
Multiclass+2.9 ± 0.2+2.2 ± 0.1+1.0 ± 0.0+0.3 ± 0.0+0.0 ± 0.0
Multilabel+1.1 ± 0.2+0.9 ± 0.1+0.4 ± 0.0+0.0 ± 0.0-0.1 ± 0.0
s-top vs tt̄HAccuracy (%)Baseline82.9 ± 0.286.3 ± 0.187.5 ± 0.188.5 ± 0.088.8 ± 0.0
Multiclass+0.4 ± 0.3+1.6 ± 0.1+0.9 ± 0.1+0.3 ± 0.0+0.0 ± 0.0
Multilabel+1.9 ± 0.2+0.8 ± 0.1+0.5 ± 0.1+0.0 ± 0.0-0.0 ± 0.0
AUC (x 100)Baseline90.3 ± 0.293.7 ± 0.094.7 ± 0.095.4 ± 0.095.6 ± 0.0
Multiclass+0.4 ± 0.2+1.2 ± 0.0+0.6 ± 0.0+0.2 ± 0.0+0.0 ± 0.0
Multilabel+0.9 ± 0.2+0.6 ± 0.1+0.3 ± 0.0+0.0 ± 0.0-0.0 ± 0.0
WH vs ZHAccuracy (%)Baseline51.2 ± 0.154.0 ± 0.155.8 ± 0.157.6 ± 0.058.0 ± 0.0
Multiclass+2.8 ± 0.2+2.8 ± 0.1+1.7 ± 0.1+0.4 ± 0.0+0.1 ± 0.0
Multilabel-0.4 ± 0.2-0.5 ± 0.2+0.3 ± 0.1+0.0 ± 0.0-0.1 ± 0.0
AUC (x 100)Baseline51.8 ± 0.156.1 ± 0.259.5 ± 0.062.1 ± 0.062.7 ± 0.0
Multiclass+4.6 ± 0.2+4.8 ± 0.2+2.5 ± 0.0+0.5 ± 0.0+0.1 ± 0.0
Multilabel-0.9 ± 0.2-0.3 ± 0.3+0.4 ± 0.2+0.0 ± 0.0-0.1 ± 0.0
- - -# Results - -## Classification Performance - -In general, the fine-tuned pretrained model achieves at least the same level of classification performance as the baseline model. Notably, there are significant improvements, particularly when the sample size is small, ranging from \\(10^3\\) to \\(10^5\\) events. In some cases, the accuracy (AUC) improvements exceed 4 percentage points (2.5 AUC points), demonstrating that pretrained models provide a strong initial representation that compensates for limited data. The numerical values of the improvements in accuracy may not fully capture the impact on the sensitivity of the measurements for which the neural network classifier is used, and the improvement in physics analysis sensitivity is likely to be greater. - -As the training sample size grows to \\(10^6\\) and eventually \\(10^7\\) events, the added benefit of pretraining diminishes. With abundant data, models trained from scratch approach or even match the accuracy of fine-tuned pretrained models. This suggests that large datasets enable effective learning from scratch, rendering the advantage of pretraining negligible in such scenarios. - -Although both pretraining approaches offer benefits, multiclass pretraining tends to provide more consistent improvements across tasks, especially in the low-data regime. In contrast, multilabel pretraining can sometimes lead to neutral or even slightly negative effects for certain tasks and data sizes. This highlights the importance of the pretraining task design, as the similarity between pretraining and fine-tuning tasks in the multiclass approach appears to yield better-aligned representations. - -We note that training sample sizes of \\(10^3\\) to \\(10^5\\) events per class reflect realistic operating conditions in HEP analyses. Searches for new physics involve simulated signal samples that are expensive to generate due to stringent physics selections, and standard analysis workflows further reduce per-category training statistics by partitioning events into multiple exclusive phase space regions, each requiring a dedicated classifier. The observed performance gains in this regime are therefore of direct practical relevance. - -Finally, the spread of accuracy across the five tasks for the baseline model is quite large, offering a robust test of fine-tuning across tasks of varying difficulty. The consistent observation of these trends across tasks confirms the reliability and robustness of the findings. - -## Generalizability to ATLAS Open Data - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
TaskDataset SizePretrainingAccuracy (%)AUC (×100)
Higgs Production2,500,000Baseline71.8 ± 0.291.2 ± 0.0
Multiclass−0.1 ± 0.2+0.1 ± 0.0
Multilabel−1.6 ± 0.2−0.6 ± 0.0
Triboson300,000Baseline54.8 ± 2.372.9 ± 3.3
Multiclass+3.7 ± 2.3+4.6 ± 3.3
Multilabel−1.7 ± 2.5−1.2 ± 3.8
- - -We evaluate the generalizability of the pretrained models on two multiclass classification tasks using the ATLAS Open Data, described in Section 2. These tasks are designed to probe whether representations learned on Delphes-simulated events transfer to data processed through the full ATLAS detector simulation and reconstruction chain, which differs significantly in its treatment of detector effects, pile-up, and object reconstruction. - -The first task uses the `GamGam` collection to classify five Higgs boson production modes (\\(ggF\\), \\(VBF\\), \\(WH\\), \\(ZH\\), and \\(t\bar{t}H\\)) where the Higgs boson decays exclusively to a pair of photons. Distinguishing between these production modes is of direct physics interest, as each is sensitive to different couplings of the Higgs boson to other SM particles. The \\(VBF\\) and \\(t\bar{t}H\\) production modes in particular are important probes of the Higgs boson’s couplings to vector bosons and top quarks respectively, and improving their separation from other production modes has direct implications for coupling measurements at the LHC. - -The second task uses the `1LMET30` collection to classify three triboson production processes: \\(WWW\\), \\(ZZZ\\), and Mixed (\\(WWZ\\), and \\(WZZ\\)). Triboson production is a rare SM process that provides a direct probe of quartic gauge boson couplings, and any deviation from SM predictions would be a strong indication of new physics. The relatively small dataset of \\(300,000\\) events reflect the rarity of these processes, making this task a particularly stringent test of the pretrained model’s ability to compensate for limited training data. - -The results are summarized in Table [tab:opendata_results]. For the Higgs production task, multiclass pretraining yields no improvement with \\(-0.1 \pm 0.2\\) percentage points in accuracy and a small increase of \\(+0.1 \pm 0.0\\) AUC points over the baseline, while multilabel pretraining leads to a slight degradation in both metrics. For the triboson task, multiclass pretraining provides a more substantial improvement of \\(+3.7 \pm 2.3\\) percentage points in accuracy and \\(+4.6 \pm 3.3\\) AUC points, consistent with the trend observed in the Delphes-based evaluation where pretraining benefits are largest when the training dataset is small. Once again, the multilabel pretraining again performs poorly with a slight degradation in both metrics. - -The consistent improvement from multiclass pretraining across both simulation frameworks suggests that the pretrained model learns physics-driven representations that are robust to differences in detector simulation. The failure of multilabel pretraining in both the Delphes and ATLAS Open Data evaluations points to a systematic limitation of this pretraining strategy, possibly reflecting a misalignment between the multilabel pretraining objective and the fine-tuning tasks. - -The poor performance of the multilabel pretraining in the low statistics regime for both the Delphes and ATLAS Open Data evaluations point to a systematic limitation of this pretraining strategy. We interpret this behavior as evidence that not all physically motivated supervised pretraining objectives produce equally transferable representations. In the multilabel setup, the model is explicitly optimized to predict a fixed set of hand-designed labels, including heavy-particle multiplicities and coarse kinematic properties. Although these labels are physically meaningful, they impose a specific structure on the learned representation: the network is encouraged to organize events according to predefined particle-level quantities rather than to discover the event-level features most useful for downstream classification. This can lead to negative transfer when the fine-tuning task depends on information that is weakly captured by, or not well aligned with, the multilabel targets. For example, tasks such as \\(WH\\) versus \\(ZH\\), FCNC versus \\(tHq\\), and the ATLAS Open Data triboson classification may depend on subtle correlations among reconstructed objects, detector-level features, or global event topology that are not fully captured by the chosen multilabel objective. In contrast, the multiclass pretraining objective directly requires the model to discriminate among complete physics processes, allowing it to learn whichever combination of object kinematics, multiplicities, and graph-level correlations is most useful for event separation. These results suggest that future pretraining objectives should avoid overly prescriptive auxiliary labels unless those labels are closely aligned with the intended downstream tasks. More flexible event-level objectives, or objectives based on discrimination among complete event topologies, may produce representations with better transferability. - -## Model Interpretability - -We aim to understand whether pretrained and baseline models learn the same underlying representations. To investigate this, we perform a layer-wise CKA analysis across three model comparisons: (1) the multiclass pretrained model *before* fine-tuning versus the best-performing baseline, (2) the best fine-tuned model versus the best-performing baseline, and (3) the best fine-tuned model versus the pretrained model before fine-tuning. Together, these three comparisons allow us to disentangle the effects of pretraining initialization from those of fine-tuning adaptation and to understand how representations evolve during fine-tuning. - -We use Centered Kernel Alignment (CKA)  (Kornblith et al. 2019) to quantify the similarity between the internal representations of neural networks by comparing their feature matrices in a manner that is invariant to scaling, rotation, and alignment. This invariance makes CKA particularly effective for studying relationships between network layers, even across networks trained from varying initialization. - -To provide an intuitive understanding of CKA values, we construct a table of the CKA scores for various transformations performed on a set of dummy data, shown in Table 2. - -- \\(A\\): randomly initialized matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu=0\\)) - -- \\(B\\): matrix with shape (1000, 64) constructed via various transformations performed on \\(A\\) - -- \\(Noise\\): randomly initialized noise matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu=0\\)) - -| Dataset | CKA Score | -|:------------------------------------------------------------|:---------:| -| \\(A, B = A\\) | 1.00 | -| \\(A, B =\\) permutation on columns of \\(A\\) | 1.00 | -| \\(A, B = A + \mathrm{Noise}(0.1)\\) | 0.99 | -| \\(A, B = A + \mathrm{Noise}(0.5)\\) | 0.80 | -| \\(A, B = A + \mathrm{Noise}(0.75)\\) | 0.77 | -| \\(A, B = A\cdot \mathrm{Noise}(1)\\) (Linear Transformation) | 0.76 | -| \\(A, B = A + \mathrm{Noise}(1)\\) | 0.69 | -| \\(A, B = A + \mathrm{Noise}(2)\\) | 0.51 | -| \\(A, B = A + \mathrm{Noise}(5)\\) | 0.39 | - -CKA scores for a dummy dataset \\(A\\) and \\(B\\), where \\(B\\) is created via various transformations performed on \\(A\\). This provides an interpretable reference scale: values above \\(\sim 0.8\\) indicate highly similar representations, values around \\(\sim 0.5\\) are comparable to adding noise of approximately 2 standard deviations, and values below \\(\sim 0.4\\) indicate substantially different representations. - -We train ensembles of models for each training task to observe how the CKA score changes due to the random initialization. The CKA score between two models is then defined as: -$$ -\begin{align} - CKA(A, B) = \frac{1}{n^2}\sum_i^n \sum_j^nCKA(A_i, B_j) -\end{align} -$$ -where \\(A_i\\) is the representation learned by the \\(i^{th}\\) model in an ensemble with \\(n\\) total models. The uncertainty in CKA is the standard deviation of \\(CKA(A_i, B_j)\\). - -The GNN architecture consists of three initial encoders (Node Encoder, Edge Encoder, Global Encoder), followed by graph processing stages applied iteratively in a message passing loop (Node Update, Edge Update, Global Update), and a final Global Decoder. The layer-wise CKA similarity for all three comparisons is shown in Figure 1 and reveals three distinct and complementary perspectives on how pretraining and fine-tuning shape model representations. - -![Layer-wise CKA similarity across three model comparisons: (top) multiclass pretrained model before fine-tuning versus the best-performing baseline, (middle) best fine-tuned model versus the best-performing baseline, and (bottom) best fine-tuned model versus the pretrained model before fine-tuning. Each panel shows the CKA similarity at each stage of the GNN architecture across all downstream tasks. For each model pair, the CKA is computed using a fixed evaluation sample of 1024 events selected from the held-out test set, and the reported CKA value is averaged over all pairwise combinations of the five independently trained models in each ensemble. Error bars represent the standard deviation across these pairwise model comparisons. The evaluation sample size was chosen for computational tractability, since the layer-wise model comparison is expensive to evaluate; although this sample is small compared to the full dataset, the small uncertainties indicate that the observed representation-level trends are stable across independently trained models.](plots/cka_layerwise.svg) - -*Layer-wise CKA similarity across three model comparisons: (top) multiclass pretrained model before fine-tuning versus the best-performing baseline, (middle) best fine-tuned model versus the best-performing baseline, and (bottom) best fine-tuned model versus the pretrained model before fine-tuning. Each panel shows the CKA similarity at each stage of the GNN architecture across all downstream tasks. For each model pair, the CKA is computed using a fixed evaluation sample of 1024 events selected from the held-out test set, and the reported CKA value is averaged over all pairwise combinations of the five independently trained models in each ensemble. Error bars represent the standard deviation across these pairwise model comparisons. The evaluation sample size was chosen for computational tractability, since the layer-wise model comparison is expensive to evaluate; although this sample is small compared to the full dataset, the small uncertainties indicate that the observed representation-level trends are stable across independently trained models.* - -### Pretrained vs Best Baseline - -The first comparison examines the representational similarity between the pretrained model, before any fine-tuning, and the best-performing baseline trained directly on each downstream task. Three distinct regimes are visible. In the encoding stages (Node Encoder, Edge Encoder, Global Encoder), CKA values are consistently high across all tasks, in the range of \\(0.9\\)–\\(1.0\\), indicating that the pretrained model has already developed low-level feature representations that closely resemble those of the best baseline, despite never having been trained on the downstream task, suggesting that pretraining on a diverse set of physics processes produces encoder representations that are nearly identical to those learned by a model fully optimized for the target classification. - -In the message passing stages (Node Update, Edge Update, Global Update), the similarity drops dramatically to values in the range of \\(0.2\\)–\\(0.5\\) across all tasks. Referencing Table 2, these values are comparable to adding noise of approximately 2 standard deviations, indicating that the pretrained and baseline models process and aggregate graph information through fundamentally different computational pathways. This suggests that the diverse pretraining objectives lead the model to develop more general intermediate representations that differ substantially from those optimized purely for the downstream task. - -At the Global Decoder, the similarity is task-dependent and correlates with task difficulty. For the harder tasks (\\(t\bar{t}H(\rightarrow\gamma\gamma)\\) CP-even vs CP-odd, FCNC vs \\(tHq\\), and \\(WH\\) vs \\(ZH\\)) the decoder representations remain substantially different from the baseline, with CKA values as low as \\(0.1\\)–\\(0.2\\). For the easier tasks (\\(t\bar{t}W\\) vs \\(t\bar{t}t\\) and s-top vs \\(t\bar{t}H\\)) the similarity partially recovers to \\(0.7\\)–\\(0.8\\). This pattern suggests that for difficult tasks, the pretrained model’s divergent intermediate representations encode inductive biases that are not accessible to the baseline model, which may explain the larger performance improvements observed for these tasks in the low-data regime. - -### Best Fine-Tuned vs Best Baseline - -The second comparison examines whether fine-tuning drives the pretrained model toward the same representational solution as the baseline, or whether it retains a distinct learned structure. The encoding stages again show high similarity in the range of \\(0.9\\)–\\(1.0\\), consistent with the first comparison and confirming that these low-level representations are stable across all training strategies. The message passing stages remain similarly low, with CKA values of \\(0.2\\)–\\(0.5\\), indicating that fine-tuning does not recover the baseline’s intermediate representations. The fine-tuned model retains the structurally different message passing pathway of the pretrained model even after task-specific training. - -The Global Decoder shows a marked improvement in similarity compared to the pretrained vs baseline comparison, with values consistently in the range of \\(0.8\\)–\\(1.0\\) across nearly all tasks. This indicates that fine-tuning primarily acts on the final decoder stage, driving the model’s output representations into close alignment with those of the best baseline regardless of task difficulty. The fine-tuned model therefore arrives at a similar decision boundary to the baseline through a fundamentally different intermediate computational pathway, which may explain why fine-tuned models can match or exceed baseline performance while maintaining distinct internal representations. - -### Best Fine-Tuned vs Pretrained - -The third comparison directly measures how much fine-tuning modifies the pretrained representations. In the encoding stages, the CKA similarity is nearly \\(1.0\\) across all tasks, confirming that fine-tuning leaves the low-level encoder representations essentially unchanged. This is consistent with the standard intuition that early layers in a pretrained network encode general features that do not need to be modified for downstream tasks. - -In the message passing stages, the similarity remains high, in the range of \\(0.8\\)–\\(0.95\\) for most tasks, indicating that fine-tuning makes only modest adjustments to the intermediate representations. This is in sharp contrast to the very low similarity observed between the pretrained and fine-tuned models versus the baseline in these same stages, and reveals that the distinctive message passing pathway of the pretrained model is largely preserved through fine-tuning rather than being restructured toward the baseline solution. - -At the Global Decoder, however, the similarity drops substantially, to values of \\(0.4\\) – \\(0.9\\) depending on the task, indicating that fine-tuning concentrates its representational changes in the final decoding stage. This is consistent with the picture emerging from the second comparison: fine-tuning primarily reorganizes decoder representations to align with the downstream task, while leaving the encoder and message passing representations largely intact. The tasks with the lowest fine-tuned vs pretrained decoder similarity ( \\(t\bar{t}H(\rightarrow\gamma\gamma)\\) CP-even vs CP-odd, FCNC vs \\(tHq\\), and \\(WH\\) vs \\(ZH\\)) are the same tasks that show the largest performance improvements from pretraining (Table [tab:results]), suggesting that the decoder reorganization is the key mechanism through which fine-tuning adapts the pretrained model to difficult downstream tasks. - -### Summary - -The three comparisons together reveal a consistent and interpretable picture of how pretraining and fine-tuning interact. The pretrained model begins with encoder representations already closely aligned with the best baseline, which is a strong initialization advantage. It then processes graph information through a fundamentally different message passing pathway that is retained through fine-tuning, suggesting that pretraining instills a distinct and general computational strategy for aggregating physics information. Fine-tuning then primarily reorganizes the final decoder representations to align with the downstream task, leaving the encoder and message passing stages largely unchanged. This picture is consistent with pretraining providing not merely a better parameter initialization, but a richer and more flexible representational structure that fine-tuning can efficiently specialize for diverse downstream tasks. - -## Computational Efficiency - -We evaluate computational efficiency using two ratios of fine-tuned training time to baseline training time, focusing only on the multiclass pretrained model. The first metric is the *time-to-target*, defined as the wall-clock time required to reach a test AUC within \\(10^{-3}\\) of the baseline model’s best test AUC. For this metric, we compare the fine-tuned model’s time-to-target against the baseline model’s time to reach the same target: - -$$ -\begin{equation} - \eta_{\mathrm{target}} = - \frac{T_{\mathrm{target}}^{\mathrm{fine\text{-}tuned}}} - {T_{\mathrm{target}}^{\mathrm{baseline}}}. -\end{equation} -$$ -This represents an optimistic target-aware early-stopping scenario. Of the 135 fine-tuning runs evaluated, only 2 fail to reach the target threshold: one FCNC vs. \\(tHq\\) ensemble member at \\(10^3\\) events, which falls short of the target by \\(0.014\\) AUC, and one s-top vs. \\(t\bar{t}H\\) ensemble member at \\(10^3\\) events, which falls short by \\(0.002\\) AUC. These two runs are excluded from the time-to-target analysis. - -The second metric compares the full training time of the fine-tuned and baseline models: -$$ -\begin{equation} - \eta_{\mathrm{full}} = - \frac{T_{\mathrm{full}}^{\mathrm{fine\text{-}tuned}}} - {T_{\mathrm{full}}^{\mathrm{baseline}}}. -\end{equation} -$$ -Here, full training time is defined as the wall-clock time required to reach a stopping condition of 5 consecutive epochs with no improvement in training loss. - -![Ratio of multiclass fine-tuned training time to baseline training time, evaluated using two stopping definitions. The top panel uses time-to-target, defined as the wall-clock time required to reach a test AUC within 10−3 of the baseline model’s best test AUC. This represents an optimistic target-aware early-stopping scenario for training. The bottom panel uses full training time, defined as the time required to reach a stopping condition of 5 consecutive epochs with no improvement in training loss. Points show individual ensemble members; bars and error bars show the ensemble mean and standard deviation. Two ensemble members that did not reach the target threshold (one FCNC vs. tHq run and one s-top vs. tt̄H run, both at 103 events) are excluded from the top panel.](plots/trainingtime.svg) - -*Ratio of multiclass fine-tuned training time to baseline training time, evaluated using two stopping definitions. The top panel uses time-to-target, defined as the wall-clock time required to reach a test AUC within 10−3 of the baseline model’s best test AUC. This represents an optimistic target-aware early-stopping scenario for training. The bottom panel uses full training time, defined as the time required to reach a stopping condition of 5 consecutive epochs with no improvement in training loss. Points show individual ensemble members; bars and error bars show the ensemble mean and standard deviation. Two ensemble members that did not reach the target threshold (one FCNC vs. tHq run and one s-top vs. tt̄H run, both at 103 events) are excluded from the top panel.* - -Figure 2 shows the two complementary measures of computational efficiency. The top panel reports the time-to-target ratio, which captures the speed at which fine-tuning reaches near-optimal performance and represents an optimistic target-aware early-stopping scenario. The bottom panel reports the full-training-time ratio, which captures the total wall-clock cost of fine-tuning under a standard loss-based stopping condition and represents a more conservative, practical estimate of computational savings. - -The time-to-target metric demonstrates the primary advantage of the pretrained model. At \\(10^5\\) events, fine-tuning reaches the AUC target in just \\(3\\)–\\(8\%\\) of the baseline training time across all tasks, corresponding to speedups exceeding \\(12\times\\) for every task. This advantage persists across all sample sizes: even for the seven full statistics (5 Delphes and 2 Open Data) tasks, fine-tuning reaches the target in an average of \\(43.5\%\\) of the baseline time. The key characteristic driving this behavior is that the fine-tuned model reaches near-optimal performance rapidly before entering a slow plateau phase, suggesting that aggressive early stopping can reduce computational costs substantially beyond what is reported under the full-training-time metric. - -Under the full-training-time metric, fine-tuning can require longer total training time than the baseline at small sample sizes. Because the fine-tuned model uses a learning rate one tenth that of the baseline for all layers except the final linear layer, it converges more slowly and continues training through a plateau long after the AUC target has already been reached, before the loss-based stopping condition eventually triggers. This additional training is not wasted, however: the fine-tuned model typically achieves higher test AUC and accuracy than the baseline by the end of training. This overhead diminishes as the training set grows, and at full statistics the full-training-time ratio falls below unity for every task, with a 7-task mean of \\(64.6\%\\). - -The main exception across both metrics is the \\(t\bar{t}H(\rightarrow\gamma\gamma)\\) CP-even vs. CP-odd task at \\(10^7\\) events, where the fine-tuned model requires \\(178\%\\) of the baseline time to reach the target, despite ultimately training to completion faster under the full-training-time metric. We attribute this to the inherent difficulty of the task and the absence of photon objects in the pretraining data, consistent with the lower classification performance and larger representational divergence observed for this task throughout the analysis. - -Taken together, the two metrics bracket the realistic range of computational savings: the time-to-target ratio gives the best-case saving achievable with well-tuned early stopping, while the full-training-time ratio gives the saving under a standard stopping condition with no additional optimization. Fine-tuning provides consistent time-to-target speedups across all tasks and sample sizes, while the full-training-time advantage is most pronounced at large sample sizes, with fine-tuning being generally slower than the baseline at small sample sizes under the loss-based stopping condition. - -To evaluate the total computational cost of the foundation model approach, the pretraining time must be amortized over all fine-tuned tasks. The pretraining times are: - -- Multiclass pretraining: 45.5 GPU hours - -- Multilabel pretraining: 60.0 GPU hours - -The multilabel pretraining time is higher due to a model synchronization step when training in parallel on 16 GPUs. The foundation model approach becomes more computationally efficient than training from scratch once the pretraining cost is offset by the cumulative time savings across fine-tuned tasks. - -We compute the crossover point under both efficiency metrics at full statistics (\\(10^7\\) events for Delphes and the full ATLAS Open Data sample) fine-tuned using the multiclass pretrained model. Under the full-training-time metric, baseline training times range from \\(0.17\\) GPU hours (triboson) to \\(12.24\\) GPU hours (\\(t\bar{t}W\\) vs. \\(t\bar{t}t\\)), with a 7-task mean of \\(7.77\\) GPU hours. Directly averaging the fine-tuned full training time across all tasks gives \\(4.50\\) GPU hours, saving approximately \\(3.27\\) GPU hours per task. Dividing the multiclass pretraining cost of \\(45.5\\) GPU hours by this per-task saving yields a crossover at approximately \\(14\\) tasks. - -Under the time-to-target metric, baseline times to target range from \\(0.19\\) GPU hours (triboson) to \\(2.93\\) GPU hours (\\(t\bar{t}W\\) vs. \\(t\bar{t}t\\)), with a 7-task mean of \\(1.74\\) GPU hours. At a mean time-to-target ratio of \\(43.5\%\\) across all 7 tasks, each fine-tuned task requires on average \\(0.88\\) GPU hours, saving approximately \\(0.87\\) GPU hours per task and yielding a crossover at approximately \\(52\\) tasks. The two crossover estimates of \\(14\\) and \\(52\\) tasks bracket the realistic range depending on how aggressively early stopping is applied. - -As a practical example, the ATLAS measurement of Higgs boson couplings using the \\(H \rightarrow \gamma\gamma\\) decay channel (ATLAS Collaboration 2023a) involved training 42 classifiers for event categorization, above the full-training estimate of \\(\sim14\\) tasks, though below the time-to-target estimate of \\(\sim52\\) tasks, with the larger time-to-target estimate driven in part by the outlier \\(t\bar{t}H(\rightarrow\gamma\gamma)\\) CP-even vs. CP-odd. This suggests that the foundation model approach can potentially meaningfully reduce the total computational cost of a realistic high-energy physics analysis, even when pretraining costs are fully accounted for. - -# Conclusions - -We presented a systematic study of a particle physics foundation model designed to operate on the four-momentum and identification properties of event final-state objects. This model is built on a Graph Neural Network (GNN) architecture and trained on a dataset comprising 120 million simulated proton-proton collision events across 12 distinct physics processes. - -Fine-tuning the pretrained model on downstream classification tasks yields consistent improvements in classification performance relative to models trained from scratch, with the largest gains observed in the low-data regime. The pretrained representations transfer effectively even to processes not seen during pretraining, and performance gains persist at full statistics for most tasks, with the \\(t\bar{t}H(\rightarrow\gamma\gamma)\\) CP-even vs. CP-odd task proving the most challenging due to the inherent difficulty of the task and the absence of photon-related events in the pretraining data. The poor performance of the multilabel pretraining in the low statistics regime further indicates that pretraining objectives must be chosen carefully, as overly prescriptive auxiliary labels can bias the learned representation toward features that are not optimally aligned with downstream event-classification tasks. - -The generalizability of the foundation model is further supported by its performance on the ATLAS Open Data tasks. Despite being produced with a different simulation pipeline and detector conditions than the Delphes samples used during pretraining, the model transfers effectively to both tasks, demonstrating that the learned representations are not tied to a specific simulation framework and can generalize across different datasets and event topologies. Although the ATLAS Open Data triboson task provides a test involving a different event-generation and detector-simulation setup from the MadGraph+Pythia+Delphes pipeline used for the Delphes samples, we do not perform a controlled study isolating individual generator-level effects, such as alternative parton shower tunes or matrix-element generators for the same physics process. Such controlled generator-variation studies remain an important direction for future work. - -The foundation model approach also offers substantial computational advantages. Fine-tuning reaches the AUC target in a small fraction of the baseline training time across nearly all tasks and sample sizes, with the pretraining cost recovered after approximately \\(14\\)–\\(52\\) fine-tuned tasks depending on the stopping condition applied. This is well within the scope of a single realistic physics analysis; the ATLAS measurement of Higgs boson couplings in the \\(H\rightarrow\gamma\gamma\\) channel, for example, involved training 42 classifiers for event categorization. - -To better understand the learned representations of the pretrained model and guide future optimization efforts, we employed a representational similarity evaluation framework using Centered Kernel Alignment (CKA). This metric allowed us to investigate the source of the performance gains observed in the foundation model. The layer-wise CKA analysis reveals that fine-tuning preserves the pretrained model’s encoder representations while concentrating representational changes in the final decoder stage, suggesting that pretraining provides not merely a better parameter initialization but a richer intermediate computational structure that fine-tuning efficiently specializes for downstream tasks. Future studies could build on this by directly probing whether the distinct message-passing pathways of the fine-tuned and baseline models encode complementary physical information, for example through input feature attribution or ablation studies, which could further inform the design of pretraining objectives and architectures. - -# Acknowledgments - -This work is supported by the U.S. National Science Foundation under the Award No.2046280, and by U.S. Department of Energy, Office of Science under contract DE-AC02-05CH11231. - -Joshua Ho would like to thank the support of UC Berkeley Summer Undergraduate Research Fellowships (SURF) and its donors for this work. - -We acknowledge the work of the ATLAS Collaboration to record or simulate, reconstruct, and distribute the Open Data used in this paper, and to develop and support the software with which it was analysed. - -# Hyperparameter Optimization - -![Change in best test AUC relative to the baseline configuration for systematic hyperparameter variations, evaluated on multiclass pretraining with 12 classes on Delphes samples (left) and Higgs production multiclass classification with 5 classes on ATLAS Open Data (right). Each point represents a single training run with one hyperparameter varied at a time. The grey band in the right panel shows the baseline ±1σ variation across ensemble members. The learning rate ×100 configuration is excluded from the ±0.2% stability claim in the main text as a clear outlier attributable to training instability at excessively large learning rates.](plots/hyperparameter_optimization.svg) - -*Change in best test AUC relative to the baseline configuration for systematic hyperparameter variations, evaluated on multiclass pretraining with 12 classes on Delphes samples (left) and Higgs production multiclass classification with 5 classes on ATLAS Open Data (right). Each point represents a single training run with one hyperparameter varied at a time. The grey band in the right panel shows the baseline ±1σ variation across ensemble members. The learning rate ×100 configuration is excluded from the ±0.2% stability claim in the main text as a clear outlier attributable to training instability at excessively large learning rates.* - -To identify a robust architecture configuration and confirm that the results presented in this work are not sensitive to the choice of hyperparameters, we performed a systematic hyperparameter sweep investing a total of 400 GPU hours on a single NVIDIA A100. The sweep was conducted on two tasks: multiclass pretraining with 12 classes on Delphes samples, and Higgs production multiclass classification with 5 classes on ATLAS Open Data. The latter was chosen as a representative downstream task to ensure that hyperparameter stability holds beyond the pretraining setting. - -We swept over four hyperparameter axes independently, varying one parameter at a time relative to the baseline configuration: batch size (\\(\times 2\\), \\(\times 4\\)), hidden dimension (\\(\times 2\\), \\(\times 4\\)), learning rate (\\(\times 10\\), \\(\times 100\\)), and MLP depth (\\(+1\\), \\(+2\\) layers). The results are shown in Figure 3, where each point shows the change in best test AUC relative to the baseline configuration. - -The GNN architecture exhibits broad stability across hyperparameter configurations. For the multiclass pretraining task, all configurations except a learning rate scaling of \\(\times 100\\) remain within \\(\pm 0.2\%\\) of the baseline AUC, with the \\(\times 100\\) learning rate representing a clear outlier at approximately \\(-0.4\%\\), consistent with training instability at excessively large learning rates. The remaining variations show no systematic trend with increasing batch size, hidden dimension, or depth, indicating that the model is not operating near a capacity bottleneck and that the baseline configuration is not a particularly sensitive operating point. - -The Higgs production task shows qualitatively similar behavior, with most configurations falling within the baseline \\(\pm 0.2\%\\) band. The \\(\times 100\\) learning rate again performs poorly, and adding two additional layers produces a modest degradation of approximately \\(-1.4\%\\), which may reflect overfitting or optimization difficulty with deeper networks on this task. Configurations with increased hidden dimension and moderate learning rate scaling show small positive deviations, suggesting the baseline is slightly conservative in capacity but not meaningfully so. - -Based on these results, we adopt the baseline configuration for all experiments reported in this work, as it provides stable and competitive performance without requiring task-specific tuning. - -# Additional Performance Metrics - -The F1, precision, and recall values are presented for all tasks in Table [tab:delphes_f1] and Table [tab:open_data_f1]. These metrics show the same general trends as the accuracy and AUC results described in the body of the paper. For binary outputs, events with scores greater than or equal to \\(0.5\\) are assigned to the positive class, while for multiclass outputs the predicted class is taken to be the class with the largest score. For each class \\(c\\), precision is computed as \\(P_c = TP_c/(TP_c+FP_c)\\), recall as \\(R_c = TP_c/(TP_c+FN_c)\\), and F1 as \\(F1_c = 2P_cR_c/(P_c+R_c)\\), where \\(TP_c\\), \\(FP_c\\), and \\(FN_c\\) are the true positives, false positives, and false negatives for that class. The reported precision, recall, and F1 values are macro-averaged, meaning that the metric is first computed separately for each class and then averaged over classes with equal weight. This gives each class equal importance regardless of its sample size. For each task, pretraining strategy, and training sample size, the metrics are computed separately for each independently trained model and then averaged over all runs. - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
Name of TaskMetricPretrainingSample Size
4-8103104105106107
tt̄H( → γγ) CP-even vs CP-oddF1 (%)Baseline56.2 ± 1.062.2 ± 0.264.1 ± 0.165.7 ± 0.066.2 ± 0.0
Multiclass+3.0 ± 1.1+2.1 ± 0.2+1.1 ± 0.1+0.1 ± 0.0-0.0 ± 0.0
Multilabel+1.3 ± 1.2+1.0 ± 0.2+0.7 ± 0.2+0.0 ± 0.0-0.1 ± 0.0
Precision (%)Baseline56.7 ± 1.162.5 ± 0.164.7 ± 0.165.9 ± 0.066.4 ± 0.0
Multiclass+2.7 ± 1.1+2.1 ± 0.1+0.7 ± 0.1+0.1 ± 0.0-0.1 ± 0.0
Multilabel+1.0 ± 1.2+0.9 ± 0.2+0.3 ± 0.1-0.1 ± 0.0-0.1 ± 0.0
Recall (%)Baseline56.6 ± 1.062.4 ± 0.164.4 ± 0.165.8 ± 0.066.3 ± 0.0
Multiclass+2.8 ± 1.0+2.1 ± 0.1+0.9 ± 0.1+0.1 ± 0.0-0.0 ± 0.0
Multilabel+1.1 ± 1.2+0.9 ± 0.2+0.5 ± 0.1-0.0 ± 0.0-0.1 ± 0.0
FCNC vs tHqF1 (%)Baseline61.4 ± 0.565.6 ± 0.166.7 ± 0.367.9 ± 0.166.6 ± 0.0
Multiclass+2.9 ± 0.8+1.7 ± 0.1+1.3 ± 0.3+0.5 ± 0.1-0.1 ± 0.1
Multilabel-2.1 ± 0.8-0.3 ± 0.2+0.7 ± 0.3+0.0 ± 0.2-0.0 ± 0.1
Precision (%)Baseline62.2 ± 0.365.7 ± 0.266.8 ± 0.167.8 ± 0.166.7 ± 0.0
Multiclass+2.1 ± 0.7+1.5 ± 0.2+1.1 ± 0.1+0.4 ± 0.1-0.0 ± 0.0
Multilabel-2.4 ± 0.8-0.5 ± 0.2+0.5 ± 0.1+0.0 ± 0.1-0.1 ± 0.0
Recall (%)Baseline63.1 ± 0.465.9 ± 0.167.8 ± 0.068.9 ± 0.067.9 ± 0.0
Multiclass+1.4 ± 0.7+2.4 ± 0.1+1.1 ± 0.0+0.4 ± 0.0-0.0 ± 0.0
Multilabel-2.6 ± 0.9+0.3 ± 0.1+0.3 ± 0.1+0.1 ± 0.0-0.1 ± 0.0
tt̄W vs tt̄tF1 (%)Baseline75.7 ± 0.177.4 ± 0.178.9 ± 0.079.8 ± 0.080.3 ± 0.0
Multiclass+2.7 ± 0.2+2.2 ± 0.1+1.0 ± 0.0+0.4 ± 0.0+0.0 ± 0.0
Multilabel+1.6 ± 0.1+0.5 ± 0.2+0.3 ± 0.0-0.0 ± 0.1-0.1 ± 0.0
Precision (%)Baseline75.8 ± 0.177.6 ± 0.079.0 ± 0.079.8 ± 0.080.3 ± 0.0
Multiclass+2.6 ± 0.2+2.0 ± 0.0+1.0 ± 0.0+0.3 ± 0.0+0.0 ± 0.0
Multilabel+1.5 ± 0.1+0.6 ± 0.1+0.3 ± 0.0+0.0 ± 0.0-0.1 ± 0.0
Recall (%)Baseline75.7 ± 0.177.5 ± 0.179.0 ± 0.079.8 ± 0.080.3 ± 0.0
Multiclass+2.7 ± 0.2+2.2 ± 0.1+1.0 ± 0.0+0.4 ± 0.0+0.0 ± 0.0
Multilabel+1.6 ± 0.1+0.6 ± 0.2+0.3 ± 0.0-0.0 ± 0.1-0.1 ± 0.0
s-top vs tt̄HF1 (%)Baseline82.8 ± 0.286.3 ± 0.187.5 ± 0.188.4 ± 0.088.8 ± 0.0
Multiclass+0.4 ± 0.2+1.6 ± 0.1+0.9 ± 0.1+0.3 ± 0.0+0.0 ± 0.0
Multilabel+1.9 ± 0.2+0.8 ± 0.1+0.5 ± 0.1+0.0 ± 0.0-0.0 ± 0.0
Precision (%)Baseline83.0 ± 0.386.3 ± 0.187.6 ± 0.088.4 ± 0.088.8 ± 0.0
Multiclass+0.3 ± 0.3+1.6 ± 0.1+0.9 ± 0.0+0.3 ± 0.0+0.0 ± 0.0
Multilabel+1.9 ± 0.3+0.8 ± 0.1+0.5 ± 0.1+0.0 ± 0.0-0.0 ± 0.0
Recall (%)Baseline82.8 ± 0.286.3 ± 0.187.5 ± 0.188.5 ± 0.088.9 ± 0.0
Multiclass+0.4 ± 0.2+1.7 ± 0.1+0.9 ± 0.1+0.3 ± 0.0+0.0 ± 0.0
Multilabel+2.0 ± 0.2+0.8 ± 0.1+0.5 ± 0.1+0.0 ± 0.0-0.0 ± 0.0
WH vs ZHF1 (%)Baseline50.4 ± 0.653.9 ± 0.255.4 ± 0.457.5 ± 0.158.0 ± 0.0
Multiclass+3.5 ± 0.6+2.9 ± 0.2+2.1 ± 0.4+0.4 ± 0.1+0.1 ± 0.0
Multilabel-0.4 ± 0.7-2.5 ± 1.1+0.4 ± 0.5-0.0 ± 0.1-0.1 ± 0.0
Precision (%)Baseline51.3 ± 0.154.1 ± 0.156.1 ± 0.257.6 ± 0.058.1 ± 0.0
Multiclass+2.7 ± 0.2+2.7 ± 0.1+1.5 ± 0.2+0.3 ± 0.0+0.1 ± 0.0
Multilabel-0.5 ± 0.2+0.3 ± 0.5+0.2 ± 0.2+0.1 ± 0.1-0.0 ± 0.0
Recall (%)Baseline51.2 ± 0.154.0 ± 0.155.8 ± 0.157.6 ± 0.058.0 ± 0.0
Multiclass+2.8 ± 0.2+2.8 ± 0.1+1.7 ± 0.1+0.4 ± 0.0+0.1 ± 0.0
Multilabel-0.4 ± 0.2-0.5 ± 0.2+0.3 ± 0.1+0.0 ± 0.0-0.1 ± 0.0
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
TaskDataset SizePretrainingF1 (%)Precision (%)Recall (%)
Higgs Production2,500,000Baseline63.0 ± 0.162.2 ± 0.066.3 ± 0.0
Multiclass+0.1 ± 0.1-0.1 ± 0.1+0.5 ± 0.1
Multilabel-1.4 ± 0.2-1.4 ± 0.1-0.9 ± 0.1
Triboson300,000Baseline49.7 ± 3.555.4 ± 2.355.6 ± 2.8
Multiclass+5.6 ± 3.5+2.4 ± 2.3+4.2 ± 2.8
Multilabel-0.3 ± 3.7-3.3 ± 2.7-1.0 ± 3.5
- - -
- -
- -Aad, Georges et al. 2019. “ATLAS b-jet identification performance and efficiency measurement with \\(t{\bar{t}}\\) events in pp collisions at \\(\sqrt{s}=13\\) TeV.” *Eur. Phys. J. C* 79 (11): 970. . - - -
- -Agostinelli, S. et al. 2003. “GEANT4—a simulation toolkit.” *Nucl. Instrum. Meth. A* 506: 250–303. . - - -
- -Allanach, B. C., C. Balázs, G. Bélanger, et al. 2009. “SUSY Les Houches Accord 2.” *Computer Physics Communications* 180 (1): 8–25. https://doi.org/. - - -
- -Alwall, J., R. Frederix, S. Frixione, et al. 2014. “The automated computation of tree-level and next-to-leading order differential cross sections, and their matching to parton shower simulations.” *JHEP* 07: 079. . - - -
- -Araz, Jack Y., Vinicius Mikuni, Felix Ringer, Nobuo Sato, Fernando Torales Acosta, and Richard Whitehill. 2024. *Point Cloud-Based Diffusion Models for the Electron-Ion Collider*. . - - -
- -Artoisenet, P. et al. 2013. “A framework for Higgs characterisation.” *JHEP* 11: 043. . - - -
- -Artoisenet, Pierre, Rikkert Frederix, Olivier Mattelaer, and Robbert Rietkerk. 2013. “Automatic spin-entangled decays of heavy resonances in Monte Carlo simulations.” *JHEP* 03: 015. . - - -
- -ATLAS Collaboration. 2010. “The ATLAS Simulation Infrastructure.” *Eur. Phys. J. C* 70: 823–74. . - - -
- -ATLAS Collaboration. 2023a. “Measurement of the Properties of Higgs Boson Production at \\(\sqrt{s}=13\\) TeV in the \\(H\to\gamma\gamma\\) Channel Using 139 Fb\\(^{-1}\\) of \\(pp\\) Collision Data with the ATLAS Experiment.” *JHEP* 07: 088. . - - -
- -ATLAS Collaboration. 2023b. “Observation of four-top-quark production in the multilepton final state with the ATLAS detector.” *Eur. Phys. J. C* 83 (6): 496. . - - -
- -ATLAS Collaboration. 2025. *ROOT ntuple format 2015–2016 proton-proton Open Data for Education and Outreach beta release from the ATLAS experiment*. CERN Open Data Portal. . - - -
- -Ba, Jimmy Lei, Jamie Ryan Kiros, and Geoffrey E. Hinton. 2016. “Layer Normalization.” *arXiv Preprint arXiv:1607.06450*. . - - -
- -Battaglia, Peter W., Jessica B. Hamrick, Victor Bapst, et al. 2018. “Relational Inductive Biases, Deep Learning, and Graph Networks.” *arXiv Preprint arXiv:1806.01261*. . - - -
- -Birk, Joschka, Anna Hallin, and Gregor Kasieczka. 2024. “OmniJet-\\(\alpha\\): The First Cross-Task Foundation Model for Particle Physics.” *Machine Learning: Science and Technology* 5 (3): 035031. . - - -
- -Bothmann, Enrico et al. 2019. “Event generation with Sherpa 2.2.” *SciPost Phys.* 7 (3): 034. . - - -
- -Cacciari, Matteo, Gavin P. Salam, and Gregory Soyez. 2008. “The anti-\\(k_t\\) jet clustering algorithm.” *JHEP* 04: 063. . - - -
- -Collaboration, The ATLAS. 2008. “The ATLAS Experiment at the CERN Large Hadron Collider.” *Journal of Instrumentation* 3 (08): S08003. . - - -
- -Degrande, Celine, Fabio Maltoni, Jian Wang, and Cen Zhang. 2015. “Automatic Computations at Next-to-Leading Order in QCD for Top-Quark Flavor-Changing Neutral Processes.” *Phys. Rev. D* 91 (February): 034024. . - - -
- -Devlin, Jacob, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. “BERT: Pre-Training of Deep Bidirectional Transformers for Language Understanding.” *Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies*, 4171–86. . - - -
- -Durieux, Gauthier, Fabio Maltoni, and Cen Zhang. 2015. “Global Approach to Top-Quark Flavor-Changing Interactions.” *Phys. Rev. D* 91 (April): 074017. . - - -
- -Favereau, J. de, C. Delaere, P. Demin, et al. 2014. “DELPHES 3, A modular framework for fast simulation of a generic collider experiment.” *JHEP* 02: 057. . - - -
- -Golling, Tobias, Lukas Heinrich, Michael Kagan, et al. 2024. *Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models*. . - - -
- -Harris, Philip, Michael Kagan, Jeffrey Krupa, Benedikt Maier, and Nathaniel Woodward. 2024. *Re-Simulation-Based Self-Supervised Learning for Pre-Training Foundation Models*. . - - -
- -Hashemi, Baran, Nikolai Hartmann, Sahand Sharifzadeh, James Kahn, and Thomas Kuhr. 2024. “Ultra-High-Granularity Detector Simulation with Intra-Event Aware Generative Adversarial Network and Self-Supervised Relational Reasoning.” *Nature Communications* 15 (1). . - - -
- -Huang, Andris, Yash Melkani, Paolo Calafiura, et al. 2024. *A Language Model for Particle Tracking*. . - - -
- -Jumper, John, Richard Evans, Alexander Pritzel, et al. 2021. “Highly Accurate Protein Structure Prediction with AlphaFold.” *Nature* 596: 583–89. . - - -
- -Katel, Subash, Haoyang Li, Zihan Zhao, Raghav Kansal, Farouk Mokhtar, and Javier Duarte. 2024. *Learning Symmetry-Independent Jet Representations via Jet-Based Joint Embedding Predictive Architecture*. . - - -
- -Kornblith, Simon, Mohammad Norouzi, Honglak Lee, and Geoffrey E. Hinton. 2019. “Similarity of Neural Network Representations Revisited.” *Proceedings of the 36th International Conference on Machine Learning*, Proceedings of machine learning research, vol. 97: 3519–29. . - - -
- -Leigh, Matthew, Samuel Klein, François Charton, et al. 2024. *Is Tokenization Needed for Masked Particle Modelling?* . - - -
- -Liu, Junze, Aishik Ghosh, Dylan Smith, Pierre Baldi, and Daniel Whiteson. 2023. “Generalizing to New Geometries with Geometry-Aware Autoregressive Models (GAAMs) for Fast Calorimeter Simulation.” *Journal of Instrumentation* 18 (11): P11003. . - - -
- -Mikuni, Vinicius, and Benjamin Nachman. 2024. *OmniLearn: A Method to Simultaneously Facilitate All Jet Physics Tasks*. . - - -
- -OpenAI, Josh Achiam, Steven Adler, et al. 2024. *GPT-4 Technical Report*. . - - -
- -Paszke, Adam, Sam Gross, Francisco Massa, et al. 2019. “PyTorch: An Imperative Style, High-Performance Deep Learning Library.” *Advances in Neural Information Processing Systems* 32. . - - -
- -Podell, Dustin, Zion English, Kyle Lacey, et al. 2023. *SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis*. . - - -
- -Rombach, Robin, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. 2022. “High-Resolution Image Synthesis with Latent Diffusion Models.” *Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition*, 10684–95. . - - -
- -Rosiek, Janusz. 1990. “Complete Set of Feynman Rules for the Minimal Supersymmetric Extension of the Standard Model.” *Phys. Rev. D* 41 (June): 3464–501. . - - -
- -Sjöstrand, Torbjörn, Stefan Ask, Jesper R. Christiansen, et al. 2015. “An introduction to PYTHIA 8.2.” *Comput. Phys. Commun.* 191: 159–77. . - - -
- -Vigl, Matthias, Nicole Hartman, and Lukas Heinrich. 2024. *Finetuning Foundation Models for Joint Analysis Optimization*. . - - -
- -Wang, Minjie, Lingfan Yu, Da Zheng, et al. 2019. “Deep Graph Library: Towards Efficient and Scalable Deep Learning on Graphs.” *arXiv Preprint arXiv:1909.01315*. . - - -
- -Wildridge, Andrew J., Jack P. Rodgers, Ethan M. Colbert, Yao yao, Andreas W. Jung, and Miaoyuan Liu. 2024. *Bumblebee: Foundation Model for Particle Physics Discovery*. . - - -
- -Yosinski, Jason, Jeff Clune, Yoshua Bengio, and Hod Lipson. 2014. “How Transferable Are Features in Deep Neural Networks?” *Advances in Neural Information Processing Systems* 27. . - - -
- -Zhang, Zhengde, Yiyu Zhang, Haodong Yao, et al. 2024. *Xiwu: A Basis Flexible and Learnable LLM for High Energy Physics*. . +--- +license: mit +--- diff --git a/configs/delphes_contrastive.yaml b/configs/delphes_contrastive.yaml deleted file mode 100644 index d26e695143fd00748add2dfe1768301faa6403bb..0000000000000000000000000000000000000000 --- a/configs/delphes_contrastive.yaml +++ /dev/null @@ -1,81 +0,0 @@ -# Contrastive (SimCLR) pretraining of the object transformer on the 12 -# Delphes samples of configs/delphes_jepa.yaml, without their labels (same -# event cache): two random views of every event (rotation, eta flip, energy -# scale, object dropout) are pulled together (src/gnn4colliders/contrastive.py). -# Fine-tune with configs/gamgam_jepa_finetune.yaml -# model.pretrained=outputs/delphes_contrastive/last.pt, or compare encoders -# without fine-tuning: gnn4colliders probe configs/gamgam_jepa_finetune.yaml -# data.train_events=1000 split=validation model.pretrained=... -# -# gnn4colliders train configs/delphes_contrastive.yaml -name: delphes_contrastive - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - # JEPA ignores labels and weights; the cache is that of - # configs/delphes_pretraining.yaml. - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: contrastive - backbone: transformer - # The encoder; configs/gamgam_jepa_finetune.yaml must match it. - hid_size: 128 - n_layers: 4 - n_heads: 8 - ffn_size: 512 - dropout: 0.0 - pair_bias: true - temperature: 0.1 - scale_jitter: 0.1 - object_dropout: 0.1 - -task: - name: contrastive - -train: - max_epochs: 30 - optimizer: adamw - learning_rate: 5.0e-4 - weight_decay: 0.05 - lr_decay: 0.95 - early_stopping: null - # Fine-tune from last.pt. - best_metric: loss diff --git a/configs/delphes_jepa.yaml b/configs/delphes_jepa.yaml deleted file mode 100644 index 3c85ef3992e4ddf964f580b08d3dab1cebc86300..0000000000000000000000000000000000000000 --- a/configs/delphes_jepa.yaml +++ /dev/null @@ -1,77 +0,0 @@ -# JEPA pretraining of the object transformer on the 12 Delphes samples of -# the supervised pretraining (configs/delphes_pretraining.yaml), without -# their labels. Fine-tune with configs/gamgam_jepa_finetune.yaml. -# -# gnn4colliders prepare configs/delphes_jepa.yaml -# gnn4colliders train configs/delphes_jepa.yaml -name: delphes_jepa - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - # JEPA ignores labels and weights; the cache is that of - # configs/delphes_pretraining.yaml. - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: jepa - # The encoder; configs/gamgam_jepa_finetune.yaml must match it. - hid_size: 128 - n_layers: 4 - n_heads: 8 - ffn_size: 512 - dropout: 0.0 - pair_bias: true - predictor_layers: 2 - mask_fraction: 0.4 - ema_momentum: 0.996 - -task: - name: jepa - -train: - max_epochs: 30 - optimizer: adamw - learning_rate: 5.0e-4 - weight_decay: 0.05 - lr_decay: 0.95 - early_stopping: null - # The target drifts with the EMA, so the loss does not compare epochs; - # fine-tune from last.pt. - best_metric: loss diff --git a/configs/delphes_jepa_edge.yaml b/configs/delphes_jepa_edge.yaml deleted file mode 100644 index 4d9e7a0e76aaaf0d0fe12b74aff3e48d90163963..0000000000000000000000000000000000000000 --- a/configs/delphes_jepa_edge.yaml +++ /dev/null @@ -1,82 +0,0 @@ -# JEPA pretraining of the edge network, of the size of the supervised -# pretrained models (Pretrained_GNN/), on the same 12 Delphes samples without -# their labels. Fine-tune with configs/gamgam_finetune.yaml and -# model.pretrained=outputs/delphes_jepa_edge/last.pt. -# -# gnn4colliders prepare configs/delphes_jepa_edge.yaml -# gnn4colliders train configs/delphes_jepa_edge.yaml -name: delphes_jepa_edge - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - # JEPA ignores labels and weights; the cache is that of - # configs/delphes_pretraining.yaml. - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: jepa - backbone: edge_network - # The encoder: as in Pretrained_GNN/*/config.yaml. - hid_size: 64 - n_layers: 4 - n_proc_steps: 4 - dropout: 0.0 - # The predictor: transformer blocks. - predictor_layers: 2 - n_heads: 4 - ffn_size: 256 - mask_fraction: 0.4 - ema_momentum: 0.996 - # Without it the encodings collapse (docs/design.md, JEPA pretraining). - variance_weight: 0.1 - -task: - name: jepa - -train: - max_epochs: 30 - optimizer: adamw - learning_rate: 5.0e-4 - weight_decay: 0.05 - lr_decay: 0.95 - early_stopping: null - # The target drifts with the EMA, so the loss does not compare epochs; - # fine-tune from last.pt. - best_metric: loss diff --git a/configs/delphes_masked.yaml b/configs/delphes_masked.yaml deleted file mode 100644 index 2c5a4e2b256e277b71833d9ad80f268565907610..0000000000000000000000000000000000000000 --- a/configs/delphes_masked.yaml +++ /dev/null @@ -1,78 +0,0 @@ -# Masked object modelling of the object transformer on the 12 Delphes -# samples of configs/delphes_jepa.yaml, without their labels (same event -# cache): JEPA's masking and predictor, but the predictor predicts the hidden -# objects' kinematics instead of their target encodings (model.jepa_target: -# features; src/gnn4colliders/jepa.py). Fine-tune or probe with -# configs/gamgam_jepa_finetune.yaml model.pretrained=outputs/delphes_masked/last.pt. -# -# gnn4colliders train configs/delphes_masked.yaml -name: delphes_masked - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - # JEPA ignores labels and weights; the cache is that of - # configs/delphes_pretraining.yaml. - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: jepa - jepa_target: features - # The encoder; configs/gamgam_jepa_finetune.yaml must match it. - hid_size: 128 - n_layers: 4 - n_heads: 8 - ffn_size: 512 - dropout: 0.0 - pair_bias: true - predictor_layers: 2 - mask_fraction: 0.4 - -task: - name: jepa - -train: - max_epochs: 30 - optimizer: adamw - learning_rate: 5.0e-4 - weight_decay: 0.05 - lr_decay: 0.95 - early_stopping: null - # Fine-tune from last.pt. - best_metric: loss diff --git a/configs/delphes_pfn.yaml b/configs/delphes_pfn.yaml deleted file mode 100644 index 793b672ebf2e9049afba2d17cdf1f7c395409c50..0000000000000000000000000000000000000000 --- a/configs/delphes_pfn.yaml +++ /dev/null @@ -1,86 +0,0 @@ -# Meta-training of a prior-fitted network (src/gnn4colliders/pfn.py) on -# classification tasks drawn from the 12 Delphes samples: it learns to -# classify events given labelled context events, and is then evaluated in -# context on new data (configs/gamgam_pfn.yaml). The encoder has the sizes of -# configs/delphes_jepa.yaml, so model.pretrained=outputs/delphes_jepa/last.pt -# starts it from the JEPA encoder. task.prior_classes=[1,6,7,8,9,10,11] -# leaves out the five processes of the GamGam task (ttH, ggF, VBF, WH, ZH). -# -# gnn4colliders train configs/delphes_pfn.yaml -name: delphes_pfn - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - # Tasks use the labels, not the weights; the cache is that of - # configs/delphes_pretraining.yaml. - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - # Events per step (tasks of context and query events). - batch_size: 4096 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: pfn - hid_size: 128 - n_layers: 4 - n_heads: 8 - ffn_size: 512 - dropout: 0.0 - pair_bias: true - pretrained: null - max_classes: 10 - pfn_layers: 6 - -task: - name: pfn - # The prior: classes per task, context events (log-uniform) and queries. - prior_classes: null - task_classes: [2, 10] - context_events: [16, 2048] - query_events: 256 - energy_jitter: 0.2 - steps_per_epoch: 1000 - validation_steps: 50 - -train: - max_epochs: 50 - optimizer: adamw - learning_rate: 1.0e-4 - weight_decay: 0.01 - lr_decay: 0.95 - early_stopping: null - best_metric: loss diff --git a/configs/delphes_pretraining.yaml b/configs/delphes_pretraining.yaml deleted file mode 100644 index cd9a7b2f0cd73112ba96d8caf1b3c00df42870d8..0000000000000000000000000000000000000000 --- a/configs/delphes_pretraining.yaml +++ /dev/null @@ -1,61 +0,0 @@ -# 12-class pretraining on the Delphes samples of Ho et al. (2026), read from -# the HWresearch/Delphes Hugging Face dataset. -name: delphes_pretraining - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: edge_network - hid_size: 64 - out_size: 12 - n_layers: 4 - n_proc_steps: 4 - -task: - name: multiclass - -train: - max_epochs: 100 - learning_rate: 1.0e-4 - lr_decay: 0.99 - early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0} - best_metric: auc diff --git a/configs/delphes_tabicl.yaml b/configs/delphes_tabicl.yaml deleted file mode 100644 index 7dc7293c7af09983beb0e90b20fff2574c59ffc7..0000000000000000000000000000000000000000 --- a/configs/delphes_tabicl.yaml +++ /dev/null @@ -1,89 +0,0 @@ -# Meta-training of the TabICL network (src/gnn4colliders/tabular.py; BSD-3, -# so meta-trained weights can be shared freely) on the classification tasks -# of configs/delphes_pfn.yaml, on the fixed-width event rows of -# configs/gamgam_tabicl.yaml. By default TabICL is frozen and its in-context -# blocks learn LoRA updates (model.lora_rank; 0 trains all 28M weights). -# task.prior_classes=[1,6,7,8,9,10,11] leaves out the five processes of the -# GamGam task. Evaluate in context with -# gnn4colliders evaluate configs/gamgam_tabicl.yaml split=validation \ -# checkpoint=outputs/delphes_tabicl/best.pt -# Needs `uv sync --extra tabicl`. -# -# gnn4colliders train configs/delphes_tabicl.yaml -name: delphes_tabicl - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - # Tasks use the labels, not the weights; the cache is that of - # configs/delphes_pretraining.yaml. - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - # Events per step (tasks of context and query events). TabICL (28M - # weights) recomputes activations in the backward pass. - batch_size: 4096 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: tabicl - tabicl_checkpoint: tabicl-classifier-v2-20260212.ckpt - # Leading jets, electrons, muons, photons and MET in a table row. - slots: [6, 2, 2, 3, 1] - # As in configs/gamgam_tabicl.yaml, which evaluates this model. - relative_phi: false - pair_objects: 0 - max_classes: 10 - lora_rank: 8 - -task: - name: pfn - # The prior: classes per task, context events (log-uniform) and queries. - prior_classes: null - task_classes: [2, 10] - context_events: [16, 2048] - query_events: 256 - energy_jitter: 0.2 - steps_per_epoch: 500 - validation_steps: 50 - -train: - max_epochs: 20 - optimizer: adamw - learning_rate: 1.0e-4 - weight_decay: 0.01 - lr_decay: 0.95 - early_stopping: null - best_metric: loss diff --git a/configs/delphes_tabpfn.yaml b/configs/delphes_tabpfn.yaml deleted file mode 100644 index efb9d9fb01e1f719c86a196681f748b378988505..0000000000000000000000000000000000000000 --- a/configs/delphes_tabpfn.yaml +++ /dev/null @@ -1,89 +0,0 @@ -# Meta-training of the TabPFN-3.5 network (src/gnn4colliders/tabular.py) on -# the classification tasks of configs/delphes_pfn.yaml, on the fixed-width -# event rows of configs/gamgam_tabpfn.yaml. By default TabPFN is frozen and -# its in-context blocks learn LoRA updates (model.lora_rank; 0 trains all -# weights). task.prior_classes=[1,6,7,8,9,10,11] leaves out the five -# processes of the GamGam task. Evaluate in context with -# gnn4colliders evaluate configs/gamgam_tabpfn.yaml split=validation \ -# checkpoint=outputs/delphes_tabpfn/best.pt -# Needs `uv sync --extra tabpfn` and the TabPFN-3.5 weights (docs/usage.md). -# -# gnn4colliders train configs/delphes_tabpfn.yaml -name: delphes_tabpfn - -data: - hf_repo: HWresearch/Delphes - hf_revision: 76a6c362bfba8e766ba7255a5c08a55257f78d0e - files: - - {path: samples/higgs/top-associated/tth/ttH_NLO_inc.root, label: 0} - - {path: samples/higgs/top-associated/thjb/tHjb_NLO_inc.root, label: 1} - - {path: samples/higgs/ggf/ggF_NLO_inc.root, label: 2} - - {path: samples/higgs/vbf/VBF_NLO_inc.root, label: 3} - - {path: samples/higgs/vh/wh/WH_NLO_inc.root, label: 4} - - {path: samples/higgs/vh/zh/ZH_NLO_inc.root, label: 5} - - {path: samples/photon/ttyy/ttyy.root, label: 6} - - {path: samples/top/multitop/tttt.root, label: 7} - - {path: samples/top/single-top/SingleT_schan.root, label: 8} - - {path: samples/top/ttbar/ttbar.root, label: 9} - - {path: samples/top/ttv/ttW.root, label: 10} - - {path: samples/top/multitop/ttt.root, label: 11} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. Pretraining trained - # on folds 2-9 and tested on 0-1; here fold 0 validates and fold 1 tests. - fold_var: Number - fold_modulus: 10 - train_folds: [2, 3, 4, 5, 6, 7, 8, 9] - validation_folds: [0] - test_folds: [1] - # Tasks use the labels, not the weights; the cache is that of - # configs/delphes_pretraining.yaml. - weight_var: weight - cache: outputs/delphes_pretraining/events.pt - # Events per step (tasks of context and query events): one task of the - # largest context. TabPFN recomputes activations in the backward pass. - batch_size: 2304 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: tabpfn - tabpfn_version: v3.5 - tabpfn_checkpoint: null - # Leading jets, electrons, muons, photons and MET in a table row. - slots: [6, 2, 2, 3, 1] - # As in configs/gamgam_tabpfn.yaml, which evaluates this model. - relative_phi: false - pair_objects: 0 - max_classes: 10 - lora_rank: 8 - -task: - name: pfn - # The prior: classes per task, context events (log-uniform) and queries. - prior_classes: null - task_classes: [2, 10] - context_events: [16, 2048] - query_events: 256 - energy_jitter: 0.2 - steps_per_epoch: 500 - validation_steps: 50 - -train: - max_epochs: 20 - optimizer: adamw - learning_rate: 1.0e-4 - weight_decay: 0.01 - lr_decay: 0.95 - early_stopping: null - best_metric: loss diff --git a/configs/gamgam.yaml b/configs/gamgam.yaml deleted file mode 100644 index fdc97ab367d82fc69e08ad687c2e851cf573038b..0000000000000000000000000000000000000000 --- a/configs/gamgam.yaml +++ /dev/null @@ -1,97 +0,0 @@ -# Higgs production-mode classification on ATLAS Open Data (13 TeV, 2020 -# release) GamGam MC: the 5-class from-scratch baseline of Ho et al., -# JINST 21 (2026) P08006, Tables 1 and 3. 2,473,335 events. -# -# gnn4colliders prepare configs/gamgam.yaml -# gnn4colliders train configs/gamgam.yaml [train.fast=true] [train.seed=2] -# gnn4colliders evaluate configs/gamgam.yaml -name: gamgam - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -# Baseline GNN of Ho et al. (2026), Table 1. -model: - name: edge_network - hid_size: 64 - out_size: 5 - n_layers: 4 - n_proc_steps: 4 - dropout: 0.0 - -task: - name: multiclass - -# Sections 3.2 and 3.5: Adam at 1e-4 decayed by 0.99 per epoch, stopped -# after 5 epochs without training-loss improvement; the epoch with the best -# validation AUC is kept (best.pt). train.fast=true adds torch.compile, CUDA -# graphs, bfloat16 and TF32 (about 3x faster on an A100; results agree to -# rounding, not bitwise). -train: - max_epochs: 200 - learning_rate: 1.0e-4 - lr_decay: 0.99 - early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0} - best_metric: auc diff --git a/configs/gamgam_deep_sets.yaml b/configs/gamgam_deep_sets.yaml deleted file mode 100644 index 33dda9b0150a8599e5b2f569ac30185d969d345b..0000000000000000000000000000000000000000 --- a/configs/gamgam_deep_sets.yaml +++ /dev/null @@ -1,88 +0,0 @@ -# The GamGam 5-class task with Deep Sets (a per-object MLP, summed): the -# baseline without object interactions (no paper reference). Uses the cache -# of configs/gamgam.yaml. -# -# gnn4colliders train configs/gamgam_deep_sets.yaml -name: gamgam_deep_sets - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: deep_sets - hid_size: 64 - out_size: 5 - n_layers: 3 - dropout: 0.0 - -task: - name: multiclass - -train: - max_epochs: 100 - learning_rate: 1.0e-4 - lr_decay: 0.99 - early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0} - best_metric: auc diff --git a/configs/gamgam_finetune.yaml b/configs/gamgam_finetune.yaml deleted file mode 100644 index 49c3e86e1b32ac1e669aa7b79db9d5cae7f090ff..0000000000000000000000000000000000000000 --- a/configs/gamgam_finetune.yaml +++ /dev/null @@ -1,97 +0,0 @@ -# The GamGam 5-class task fine-tuned from the multiclass pretrained model of -# Ho et al. (2026): its edge network with a new 5-class classifier, all -# weights trained. Uses the cache of configs/gamgam.yaml. -# -# gnn4colliders train configs/gamgam_finetune.yaml -name: gamgam_finetune - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -# The pretrained network's size (Pretrained_GNN/*/config.yaml). -model: - name: edge_network - pretrained: Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt - freeze_backbone: false - hid_size: 64 - out_size: 5 - n_layers: 4 - n_proc_steps: 4 - dropout: 0.0 - -task: - name: multiclass - -# Sections 3.2 and 3.5: Adam at 1e-4 decayed by 0.99 per epoch, stopped -# after 5 epochs without training-loss improvement; the epoch with the best -# validation AUC is kept (best.pt). train.fast=true adds torch.compile, CUDA -# graphs, bfloat16 and TF32 (about 3x faster on an A100; results agree to -# rounding, not bitwise). -train: - max_epochs: 200 - learning_rate: 1.0e-4 - lr_decay: 0.99 - early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0} - best_metric: auc diff --git a/configs/gamgam_jepa_finetune.yaml b/configs/gamgam_jepa_finetune.yaml deleted file mode 100644 index 43fa167159424eb342adc794dd1d72dd2f86195a..0000000000000000000000000000000000000000 --- a/configs/gamgam_jepa_finetune.yaml +++ /dev/null @@ -1,95 +0,0 @@ -# The GamGam 5-class task with the transformer encoder pretrained by JEPA on -# the Delphes samples (configs/delphes_jepa.yaml). model.pretrained=null -# trains the same architecture from scratch. Uses the cache of -# configs/gamgam.yaml. -# -# gnn4colliders train configs/gamgam_jepa_finetune.yaml -name: gamgam_jepa_finetune - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: transformer - out_size: 5 - hid_size: 128 - n_layers: 4 - n_heads: 8 - ffn_size: 512 - dropout: 0.0 - pair_bias: true - pretrained: outputs/delphes_jepa/last.pt - # true trains only the classifier (linear probe). - freeze_backbone: false - -task: - name: multiclass - -train: - max_epochs: 100 - learning_rate: 1.0e-4 - lr_decay: 0.99 - early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0} - best_metric: auc diff --git a/configs/gamgam_pfn.yaml b/configs/gamgam_pfn.yaml deleted file mode 100644 index d2c3892dd1cf3dee07c156fb8bd0260545f51091..0000000000000000000000000000000000000000 --- a/configs/gamgam_pfn.yaml +++ /dev/null @@ -1,90 +0,0 @@ -# The GamGam 5-class task classified in context by the PFN meta-trained on the -# Delphes samples (configs/delphes_pfn.yaml), with no training on GamGam: -# data.train_events labelled training events, an equal number per class, are -# the context. The model section must match the meta-training. Uses the -# cache of configs/gamgam.yaml. -# -# gnn4colliders evaluate configs/gamgam_pfn.yaml split=validation -name: gamgam_pfn -checkpoint: outputs/delphes_pfn/best.pt - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Labelled events (equal per class) that form the context. - train_events: 1000 - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: pfn - out_size: 5 - hid_size: 128 - n_layers: 4 - n_heads: 8 - ffn_size: 512 - pair_bias: true - max_classes: 10 - pfn_layers: 6 - -task: - name: multiclass diff --git a/configs/gamgam_tabdpt.yaml b/configs/gamgam_tabdpt.yaml deleted file mode 100644 index 9033da0b0ef77c9cb99ff1f26a578cb15e84d0fe..0000000000000000000000000000000000000000 --- a/configs/gamgam_tabdpt.yaml +++ /dev/null @@ -1,91 +0,0 @@ -# The GamGam 5-class task classified in context by TabDPT (Layer 6; -# Apache-2.0 weights from Hugging Face), with no training: data.train_events -# labelled training events, an equal number per class, are the context of -# every prediction (src/gnn4colliders/tabular.py). TabDPT is trained on real -# tables; with model.tabdpt_context_size each event's context is its nearest -# context events (retrieval), which allows contexts beyond what fits at once. -# Needs `uv sync --extra tabdpt`. Uses the cache of configs/gamgam.yaml. -# -# gnn4colliders evaluate configs/gamgam_tabdpt.yaml split=validation -name: gamgam_tabdpt - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Labelled events (equal per class) that form the context. - train_events: 1000 - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - # Events classified per TabDPT call. - batch_size: 8192 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: tabdpt - out_size: 5 - tabdpt_estimators: 8 - tabdpt_context_size: null - # Leading jets, electrons, muons, photons and MET in a table row. - slots: [6, 2, 2, 3, 1] - relative_phi: false - pair_objects: 0 - -task: - name: multiclass diff --git a/configs/gamgam_tabicl.yaml b/configs/gamgam_tabicl.yaml deleted file mode 100644 index 2ef050b7c20459e9c477b2317e3b2581d57f7c64..0000000000000000000000000000000000000000 --- a/configs/gamgam_tabicl.yaml +++ /dev/null @@ -1,95 +0,0 @@ -# The GamGam 5-class task classified in context by TabICLv2 (Inria; BSD-3 -# weights from Hugging Face), with no training: data.train_events labelled -# training events, an equal number per class, are the context of every -# prediction (src/gnn4colliders/tabular.py). Needs `uv sync --extra tabicl`. -# Uses the cache of configs/gamgam.yaml. The rows are those of -# configs/gamgam_tabpfn.yaml; model.relative_phi=true and -# model.pair_objects=4 add the physics features. -# -# gnn4colliders evaluate configs/gamgam_tabicl.yaml split=validation -# -# With checkpoint=outputs/delphes_tabicl/best.pt, the TabICL network -# meta-trained by configs/delphes_tabicl.yaml (model.lora_rank and -# model.max_classes must match it) classifies with its own forward pass. -name: gamgam_tabicl - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Labelled events (equal per class) that form the context. - train_events: 1000 - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - # Events classified per TabICL call. - batch_size: 8192 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: tabicl - out_size: 5 - tabicl_checkpoint: tabicl-classifier-v2-20260212.ckpt - tabicl_estimators: 8 - # Leading jets, electrons, muons, photons and MET in a table row. - slots: [6, 2, 2, 3, 1] - relative_phi: false - pair_objects: 0 - -task: - name: multiclass diff --git a/configs/gamgam_tabpfn.yaml b/configs/gamgam_tabpfn.yaml deleted file mode 100644 index 72537cb2c2397cbcc81358074a7b9b7856eb5fcf..0000000000000000000000000000000000000000 --- a/configs/gamgam_tabpfn.yaml +++ /dev/null @@ -1,96 +0,0 @@ -# The GamGam 5-class task classified in context by TabPFN (Prior Labs), with -# no training: data.train_events labelled training events, an equal number -# per class, are the context of every prediction (src/gnn4colliders/tabular.py). -# Needs `uv sync --extra tabpfn` and, for TabPFN-2.5 and later, a one-time -# login (a token in ~/.cache/tabpfn or TABPFN_TOKEN). Uses the cache of -# configs/gamgam.yaml. -# -# gnn4colliders evaluate configs/gamgam_tabpfn.yaml split=validation -# -# With checkpoint=outputs/delphes_tabpfn/best.pt, the TabPFN network -# meta-trained by configs/delphes_tabpfn.yaml (model.lora_rank and -# model.max_classes must match it) classifies with its own forward pass. -name: gamgam_tabpfn - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Labelled events (equal per class) that form the context. - train_events: 1000 - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - # Events classified per TabPFN call. - batch_size: 8192 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: tabpfn - out_size: 5 - tabpfn_version: v3.5 - tabpfn_estimators: 4 - # Leading jets, electrons, muons, photons and MET in a table row. - slots: [6, 2, 2, 3, 1] - # Physics features: azimuths from the leading object's, and the mass and - # Delta R of every pair of the leading pair_objects objects. - relative_phi: false - pair_objects: 0 - -task: - name: multiclass diff --git a/configs/gamgam_transformer.yaml b/configs/gamgam_transformer.yaml deleted file mode 100644 index 8590cb2b4f5d9b6df6a1657c123886fa73dedadb..0000000000000000000000000000000000000000 --- a/configs/gamgam_transformer.yaml +++ /dev/null @@ -1,92 +0,0 @@ -# The GamGam 5-class task with a transformer over the same objects (no -# paper reference). Uses the cache of configs/gamgam.yaml. -# -# gnn4colliders train configs/gamgam_transformer.yaml -name: gamgam_transformer - -data: - hf_repo: ho22joshua/atlas_opendata - hf_revision: d1bcc9ccaa304ea09e11f7d724342568f72a0d6e - files: - - path: GamGam/2020/MC/mc_341081.ttH125_gamgam.GamGam.root - label: 0 - size_bytes: 217819076 - sha256: c498974926f6c55143d17fc4d3d6fa69bd3718ff92cf7d5cad5c3d67a836c7b7 - - path: GamGam/2020/MC/mc_343981.ggH125_gamgam.GamGam.root - label: 1 - size_bytes: 218418347 - sha256: 05b3e6709e8836a7b862a7fde09086b3a9d4d494d62a6984d5df39d8daad13eb - - path: GamGam/2020/MC/mc_345041.VBFH125_gamgam.GamGam.root - label: 2 - size_bytes: 108616148 - sha256: 92334fce3bffc9eb6e0745dd6ef354d5ce33da764858511b8f9ca49b862b44e2 - - path: GamGam/2020/MC/mc_345318.WpH125J_Wincl_gamgam.GamGam.root - label: 3 - size_bytes: 29757932 - sha256: f255a9c3c6935f976f788b6f32d9a8f6bf576eb12ecd03a6f3fd0a5dea9d1248 - - path: GamGam/2020/MC/mc_345319.ZH125J_Zincl_gamgam.GamGam.root - label: 4 - size_bytes: 58701804 - sha256: c85623b04ac0176c64acdb2ecf8e992a63824cfc430a421df742319518ff6748 - class_names: [ttH, ggF, VBF, WH, ZH] - tree_name: mini - # Electrons and muons share the lep_* collection; b-tagging uses the - # MV2c10 70% efficiency working point of the Open Data documentation. - derived_branches: - ele_pt: {source: lep_pt, where: {lep_type: 11}} - ele_eta: {source: lep_eta, where: {lep_type: 11}} - ele_phi: {source: lep_phi, where: {lep_type: 11}} - ele_charge: {source: lep_charge, where: {lep_type: 11}} - mu_pt: {source: lep_pt, where: {lep_type: 13}} - mu_eta: {source: lep_eta, where: {lep_type: 13}} - mu_phi: {source: lep_phi, where: {lep_type: 13}} - mu_charge: {source: lep_charge, where: {lep_type: 13}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - # Object types (node types) as in pretraining: jet, electron, muon, photon, MET. - feature_branches: - - [jet_pt, ele_pt, mu_pt, photon_pt, met_et] - - [jet_eta, ele_eta, mu_eta, photon_eta, 0] - - [jet_phi, ele_phi, mu_phi, photon_phi, met_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Momenta are in MeV; 1e-4 gives the pretraining unit of 10 GeV. - scales: [1.0e-4, 1, 1, 1.0e-4, 1, 1, 1] - # eventNumber % 10 gives the paper's 80:10:10 train:validation:test split. - fold_var: eventNumber - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - # Unweighted: each event counts once and classes are balanced by the - # per-class mean loss (23% of ttH events have mcWeight = -1). - weight_var: null - cache: outputs/gamgam/events.pt - batch_size: 1024 - prepare_workers: 8 - loader_workers: 4 - pin_memory: true - -model: - name: transformer - hid_size: 64 - out_size: 5 - n_layers: 2 - n_heads: 4 - ffn_size: 256 - dropout: 0.1 - # Attention bias from object-pair (deta, dphi, dR, mass) features, as in the - # Particle Transformer; false gives a plain transformer. - pair_bias: true - -task: - name: multiclass - -train: - max_epochs: 100 - learning_rate: 1.0e-4 - lr_decay: 0.99 - early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0} - best_metric: auc diff --git a/configs/tth_cp_even_odd.yaml b/configs/tth_cp_even_odd.yaml deleted file mode 100644 index 367a9397bfc19ba070b8da0b10d227e2f5723644..0000000000000000000000000000000000000000 --- a/configs/tth_cp_even_odd.yaml +++ /dev/null @@ -1,49 +0,0 @@ -# CP-even versus CP-odd top-Yukawa ttH (H -> gamma gamma), Delphes samples. -# Point data.files at local copies of the two samples. -name: tth_cp_even_odd - -data: - files: - - {path: data/raw/ttH_NLO_download/samples/higgs/top-associated/tth/ttH_NLO.root, label: 0} - - {path: data/raw/ttH_cp_even_odd/samples/higgs/top-associated/tth/ttH_CPodd.root, label: 1} - tree_name: output - feature_branches: - - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] - - [jet_eta, ele_eta, mu_eta, ph_eta, 0] - - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] - - CALC_E - - [jet_btag, 0, 0, 0, 0] - - [0, ele_charge, mu_charge, 0, 0] - - NODE_TYPE - object_types: [vector, vector, vector, vector, single] - # Delphes momenta are in GeV; 0.1 gives the pretraining unit of 10 GeV. - scales: [0.1, 1, 1, 0.1, 1, 1, 1] - # Number is the event index; Number % 10 is the fold. - fold_var: Number - fold_modulus: 10 - train_folds: [0, 1, 2, 3, 4, 5, 6, 7] - validation_folds: [8] - test_folds: [9] - weight_var: weight - cache: outputs/tth_cp_even_odd/events.pt - batch_size: 1024 - -model: - name: edge_network - hid_size: 64 - out_size: 1 - n_layers: 4 - n_proc_steps: 4 - -# NLO samples have signed weights; |weight| keeps the loss denominator -# stable (metrics use events with positive weight). -task: - name: binary - absolute_weights: true - -train: - max_epochs: 100 - learning_rate: 1.0e-4 - lr_decay: 0.99 - early_stopping: {monitor: train_loss, patience: 5, min_delta: 0.0} - best_metric: auc diff --git a/docs/design.md b/docs/design.md deleted file mode 100644 index 0c0c9c9fbac98ee8be85cfa0f6c9bf43182d9913..0000000000000000000000000000000000000000 --- a/docs/design.md +++ /dev/null @@ -1,392 +0,0 @@ -# Design - -## Modules - -| Module | Contents | -| --- | --- | -| `config.py` | the experiment dataclasses, YAML loading and overrides | -| `data.py` | ROOT reading, Hugging Face download, derived branches, object features, `PackedEvents` and its cache | -| `loader.py` | `EventLoader`: deterministic shuffling, DDP sharding, collation workers | -| `graphs.py` | fully connected `TensorGraph` batches, edge features, CUDA-graph padding buffers | -| `sequences.py` | padded token batches for the transformer and Deep Sets, event augmentations | -| `models.py` | `EdgeNetwork`, `FineTunedEdgeNetwork`, `ObjectTransformer`, `TransformerClassifier`, `DeepSets`, checkpoint loading | -| `jepa.py` | JEPA pretraining of the transformer or edge network encoder (encoding or kinematic targets), event embeddings | -| `contrastive.py` | contrastive (SimCLR) pretraining of the same encoders | -| `probe.py` | k-NN and linear probes of frozen encoders, linear CKA between them | -| `pfn.py` | our prior-fitted network over events and the task prior of meta-training | -| `tabular.py` | event rows and encoder embeddings for TabPFN, TabICL and TabDPT; LoRA meta-training of TabPFN and TabICL | -| `tasks.py` | multiclass and binary losses and metrics, the pretraining and PFN losses | -| `train.py` | `Trainer`, early stopping, checkpoints | -| `cuda_graphs.py` | captured training steps | -| `distributed.py` | `torchrun` process group and collectives | -| `inference.py` | model outputs over a loader, prediction files | -| `export.py` | ONNX export | -| `cli.py` | the six commands | - -## Data flow - -`prepare` reads each ROOT file in chunks of 100,000 entries (in parallel with -`data.prepare_workers`) and stores every event's objects as rows of one -`[objects, 7]` tensor with offsets, plus per-event labels, folds, weights and -source file and entry. Nothing in the cache depends on the model: graphs and -token sequences are built per batch, in the loader workers. - -A batch of the edge network is a set of fully connected graphs in flat -edge-index form (`edge_src`, `edge_dst`, `node_batch`, `edge_batch`). Edges -are directed, source-major and without self-loops except for one-node graphs, -with features `[deta, dphi, dR]`. These conventions, the seven node columns, -the MLP layout and the parameter names are those of the original ROOT-GNN, so -the pretrained checkpoints in `Pretrained_GNN/` load without conversion. - -The transformer and Deep Sets take each event's objects as tokens -`[events, objects, 7]`, zero-padded to the largest event of the batch, with -a mask of real objects. A token is a linear map of the kinematic columns -(`phi` as its cosine and sine) plus a learned embedding of `node_type`. The -objects are stored by type (jets, electrons, muons, photons, MET), so a -position embedding would mostly relearn the type; without one, neither -model depends on object order, as with the edge network. - -* The transformer is a stack of pre-norm attention blocks, a mean over real - objects and a linear classifier. With `model.pair_bias` (as in the - Particle Transformer, Qu et al. 2022), an MLP of the pair features - `[deta, dphi, dR, log(1 + m^2)]` (the edge features plus the invariant - mass of the pair, objects taken as massless) adds a bias per head to the - attention logits, so the transformer sees the same pairwise angles as the - edge network. -* Deep Sets (Zaheer et al. 2017; the particle flow network of Komiske et - al. 2019) applies an MLP to each object, sums over real objects and - applies an MLP to the sum: the baseline without object interactions. - -### JEPA pretraining - -`model.name: jepa` pretrains an object encoder without labels, as in -I-JEPA (Assran et al. 2023), on the 12 Delphes samples of the supervised -pretraining. `model.backbone` selects the encoder: the transformer's -(`configs/delphes_jepa.yaml`) or the edge network, of the size of the -supervised pretrained models (`configs/delphes_jepa_edge.yaml`), whose -object encodings are its node states after the last processing step and -whose context is the fully connected graph of the visible objects only. - -* a random `model.mask_fraction` of each event's objects (rounded, at least - one, and at least one left visible) is hidden; -* the context encoder sees the visible objects; the target encoder, an - exponential moving average of it (momentum rising linearly from - `model.ema_momentum` to 1 over training) without gradients, sees all - objects; -* a predictor (`model.predictor_layers` attention blocks over the visible - objects' encodings and one query per hidden object) predicts the hidden - objects' layer-normalized target encodings, with a mean squared error. - A query is told only the hidden object's type and its index among the - objects of that type (their pt order), not its kinematics; pair features - are used only between visible objects. For the edge network, its decoded - global state (what its classifier reads) is one more predictor input, so - the whole network is pretrained. - -The history reports `target_spread`, the standard deviation of the target -encodings over a batch's objects averaged over dimensions, which falls -towards zero if the representation collapses. The validation loss compares -predictions with a moving target, so it does not select epochs: fine-tuning -starts from `last.pt` (`configs/gamgam_jepa_finetune.yaml` for the -transformer, `configs/gamgam_finetune.yaml` with `model.pretrained` for the -edge network; encoder sizes must match). With `model.freeze_backbone=true` only the classifier is -trained (a linear probe); `model.pretrained=null` trains the same -architecture from scratch. - -`model.variance_weight` adds the VICReg variance term (Bardes et al. 2022): -the mean shortfall below 1 of each dimension's standard deviation over the -batch's context encodings (and the edge network's global states). - -First checks (one A100, `train.fast=true`, 2 pretraining epochs on Delphes -fold 2, then 2 GamGam epochs, validation AUC, seed 1; not converged): - -| GamGam model | Validation AUC | -| --- | --- | -| transformer from scratch / from JEPA | 0.898 / 0.904 | -| edge network from scratch | 0.895 | -| edge network from the supervised pretrained model | 0.907 | -| edge network from JEPA, last node states only | 0.845 | -| edge network from JEPA, mean node states | 0.869 | -| the same with `variance_weight` 0.1 / 1 | 0.874 / 0.810 | - -The edge network's JEPA encodings collapse without the variance term -(`target_spread` 0.09 against about 1 with it), and none of the edge -network variants yet beats training from scratch. - -### Contrastive and masked-object pretraining, and probes - -Two more label-free objectives pretrain the same encoders on the same -samples, and their encoders load the same way (`context_encoder`): - -* `contrastive.py` (`configs/delphes_contrastive.yaml`): SimCLR, as JetCLR - and RS3L for jets. Two views of each event (random phi rotation, eta - flip, momentum scale within `model.scale_jitter`, each object but the - first dropped with `model.object_dropout`) are embedded and projected; - the NT-Xent loss at `model.temperature` contrasts them with the other - events of the batch. `embedding_spread` watches for collapse. -* `model.jepa_target: features` (`configs/delphes_masked.yaml`): masked - object modelling (Golling et al. 2024). JEPA's masking and predictor - predict the hidden objects' kinematics instead of their target - encodings, with no target encoder: the comparison with JEPA isolates - predicting in embedding space. - -`gnn4colliders probe` compares encoders without fine-tuning (as jBOT): -k-NN (cosine, `model.probe_neighbours`) and linear (L-BFGS logistic -regression, `model.probe_l2`) classifiers of frozen event embeddings, -trained on `data.train_events` events. A randomly initialized encoder -(`model.pretrained=null`) is the reference; supervised checkpoints probe -too. Results on GamGam: "Label efficiency on GamGam" below. - -`scripts/compare_encoders.py` compares encoders of one architecture by -linear CKA (Kornblith et al. 2019), as the paper does: of their event -embeddings (a matrix), and layer by layer against the first one (the mean -over objects after each transformer block, or the edge network's mean node -states after each step, then the event embedding). With a random encoder -first, it shows where each objective moves the representation. - -### In-context models: TabPFN, TabICL, TabDPT and a collider PFN - -Both classify a dataset's events given `data.train_events` labelled -training events (an equal number per class) as context, without training on -that dataset; `cli.in_context_forward` builds them for `evaluate` and -`predict`. - -* `tabular.py`: TabPFN (Prior Labs; `model.tabpfn_version`, default 3.5) - on fixed-width rows: the leading `model.slots` objects of each type (6 - jets, 2 electrons, 2 muons, 3 photons, MET) with their six kinematic - columns, NaN when absent, and the number of objects of each type. -* `pfn.py`: our own prior-fitted network. The object transformer (which - can start from the JEPA encoder) embeds each event; the context events get - a label embedding; `model.pfn_layers` attention blocks run over a task's - events, every event attending to the context events only, so queries are - independent and the context's keys and values are computed once. It is - meta-trained (`task.name: pfn`) on tasks drawn from the Delphes samples: - 2-10 classes of `task.prior_classes` in random order, Dirichlet class - proportions, 16-2048 context events (log-uniform), 256 queries, and per - event a random phi rotation and eta flip, per task a momentum and energy - scale within `task.energy_jitter`. The GamGam processes are among the - Delphes samples; `task.prior_classes=[1,6,7,8,9,10,11]` leaves them out. - -* `tabular.py` also runs TabICLv2 (Qu et al., ICML 2026; BSD-3, open - pretraining code) on the same rows, as a second tabular in-context model. - Rows can measure azimuths from the leading object's (`model.relative_phi`) - and end with the invariant mass and Delta R of every pair of the - `model.pair_objects` leading objects by pt (energies are `pt cosh(eta)`, - so objects are massless; a pair with MET gives a transverse-like mass). - Both are off by default, so the default rows are unchanged. -* `tabular.py` also runs TabDPT (Ma et al. 2024; Apache-2.0), trained on - real tables rather than a synthetic prior; with - `model.tabdpt_context_size` each event's context is its nearest context - events (faiss retrieval), beyond what fits in one pass. It computes in - float64 in places (on a Mac it runs on the CPU), and on macOS faiss and - torch crash together unless `OMP_NUM_THREADS=1`. -* Tabular models can see, instead of or beside the rows, the event - embeddings of a pretrained object encoder (`model.tabular_inputs: - embeddings` or `both`, the encoder of `model.backbone` and - `model.pretrained`): physics-aware embeddings under a generic tabular - prior. In meta-training the encoder trains with the LoRA updates unless - `model.freeze_backbone`. -* `tabular.py` also meta-trains the TabPFN-3.5 network itself on the - PFN's tasks (`model.name: tabpfn`, `task.name: pfn`), on the rows above - (built after the augmentation), starting from Prior Labs' weights. Calling - the network directly skips TabPFN's preprocessing and ensembling, which - on the GamGam fixtures cost nothing (validation AUC 0.852 against 0.843 - with `predict_proba`, 61 events). By default the network is frozen and - the four attention projections and two MLP layers of each of its 24 - in-context blocks get rank-8 LoRA updates (2.8M parameters, zero at the - start, so training starts from TabPFN as released); `model.lora_rank=0` - trains all weights. In context, the context's keys and values are cached - (the same logits as one pass over context and queries, to 2e-5). It - separates what Prior Labs' prior and what collider tasks contribute: - TabPFN as released, TabPFN meta-trained here, and our PFN. TabICL - (`configs/delphes_tabicl.yaml`; BSD-3, so meta-trained weights can be - shared) meta-trains the same way. Its attention reads fused projection - weights directly, so LoRA is a parametrization of the weights (12 blocks, - 0.59M of 28M parameters at rank 8), for TabPFN too. TabICL's network takes - no NaN: rows are standardized by the context's statistics, absent values - at the mean, clipped at 8 sigma. On the fixtures this gives AUC 0.827 - against 0.820 with its classifier, and the cached in-context forward - equals the training forward only in float32 (with mixed precision on the - CPU, AUC 0.59). - -### Label efficiency on GamGam (first study, 2026-10-22) - -GamGam validation AUC (%) against the number of labelled training events N -(`data.train_events`, equal per class), one seed, on one 4-GPU A100 node in -an interactive allocation (3.5 h). Each entry is evaluated on 20,000 random -validation events (`data.eval_events=20000`, about +-0.3); on all 246,922, -evaluated for the supervised runs and the leading TabPFN runs, each is 0.05 -to 0.41 lower, in the same order (TabPFN-3.5: 80.28, 86.92, 89.01, 0.08 to -0.13 lower, matching the M2 Max's 80.26 and 86.93). - -| model | N = 100 | 1000 | 10000 | 50000 | -|---|---|---|---|---| -| edge network, from scratch | 78.00 | 82.74 | 87.66 | | -| edge network, from `Pretrained_GNN` | 84.95 | 87.48 | 89.98 | | -| transformer, from scratch | 72.33 | 83.47 | 87.03 | | -| transformer, from JEPA | 73.25 | 85.10 | 88.60 | | -| transformer, from masked objects | 75.92 | 86.12 | 88.73 | | -| TabPFN-3.5 (rows) | 80.41 | 87.05 | 89.09 | 89.56 | -| TabPFN-3.5 (rows, `relative_phi`, `pair_objects: 4`) | 80.61 | 87.78 | 89.56 | | -| TabPFN-3.5 (rows + JEPA embeddings) | 80.97 | 87.43 | 89.44 | | -| TabPFN-3.5 (rows + masked-object embeddings) | 82.99 | 88.22 | 89.94 | | -| TabPFN-3.5 (rows + `Pretrained_GNN` embeddings) | 87.32 | 88.96 | 90.07 | | -| TabPFN-3.5, released network, direct forward | 80.74 | 86.64 | 89.03 | | -| TabPFN-3.5, meta-trained on Delphes (LoRA) | 80.52 | 84.77 | 88.34 | | -| TabICLv2 | 80.53 | 85.68 | 88.41 | 89.12 | -| TabDPT | 77.48 | 84.71 | 88.42 | | -| our PFN, meta-trained on Delphes | 78.13 | 76.18 | 76.51 | | - -Probes (k-NN / linear) of the frozen encoders, the same N: - -| encoder | N = 100 | 1000 | 10000 | -|---|---|---|---| -| random | 73.4 / 72.6 | 76.6 / 83.1 | 79.5 / 85.4 | -| JEPA | 73.0 / 73.4 | 78.9 / 83.4 | 81.8 / 87.9 | -| masked objects | 75.0 / 77.5 | 81.7 / 85.7 | 84.4 / 88.5 | -| contrastive | 61.6 / 67.9 | 70.0 / 80.3 | 76.4 / 85.9 | - -Setup (every run and its overrides: `scripts/label_efficiency_study.sh`): -the pretraining runs saw the same 4M class-balanced Delphes events for 8 -epochs (`data.train_events=4000000`, 50 to 70 min each); the PFN -meta-trained for 25 epochs of 1000 steps (25,000 steps, 95 min) and TabPFN -for 6 epochs of 500 (3,000 steps, 48 min). Supervised runs train for a -fixed number of steps with early stopping off (`train.early_stopping=null`; -batch size, epochs: 100, 1000 for N = 100; 128, 400 for N = 1000; 256, 100 -for N = 10000; 20 to 35 min at N = 100, 2 to 4 min at N = 10000) and keep -the epoch of best validation AUC, chosen on the same 20,000 events that are -reported. That choice uses validation labels the in-context models do not -get, and from N = 1000 most supervised runs overfit: their last epoch is far -below the kept one. On all validation events (92% of them not used for the -choice) the kept epochs lose 0.05 to 0.41, no more than the in-context runs -lose, so picking the maximum over noisy evaluations inflates them little: - -| supervised run (N = 100 / 1000 / 10000) | kept epoch, 20k | kept epoch, all | last epoch, 20k | -|---|---|---|---| -| edge network, from scratch | 78.00 / 82.74 / 87.66 | 77.80 / 82.69 / 87.42 | 68.88 / 78.88 / 87.65 | -| edge network, from `Pretrained_GNN` | 84.95 / 87.48 / 89.98 | 84.54 / 87.09 / 89.74 | 80.75 / 81.58 / 86.85 | -| transformer, from scratch | 72.33 / 83.47 / 87.03 | 71.99 / 83.35 / 86.95 | 72.33 / 83.44 / 87.03 | -| transformer, from JEPA | 73.25 / 85.10 / 88.60 | 73.07 / 84.96 / 88.45 | 72.34 / 78.82 / 81.46 | -| transformer, from masked objects | 75.92 / 86.12 / 88.73 | 75.73 / 85.84 / 88.60 | 75.92 / 79.73 / 85.62 | - -With only N labelled events there is no large labelled set to choose the -epoch on; without the choice (last epoch) every supervised model is below -TabPFN-3.5 as released at every N except `Pretrained_GNN` at N = 100. - -With the configs' settings instead (batch size 1024, at most 100 epochs, -early stopping after 5 epochs without training-loss improvement) N = 1000 -is one step per epoch, at most 100 steps: the transformer then reaches -70.04 at N = 1000 instead of 83.47. The in-context models take 10 to 30 s -for the 20,000 queries at N = 1000 and 10000 (TabPFN 109 s at N = 50000; -the first run of each, N = 100, also loads the weights). - -Findings: - -* TabPFN-3.5 as released, without training, comes within 0.4 (N = 1000) - and 0.9 (N = 10000) of fine-tuning the supervised `Pretrained_GNN` (0.2 - and 0.7 on all validation events), with its epoch chosen on validation - labels, and above it from N = 1000 without that choice. It beats every - model trained from scratch at every N. TabICL is close behind; TabDPT is - behind at small N. The physics row features add up to 0.7. -* Encoder embeddings help TabPFN most at small N. With the label-free - masked-object encoder: +2.6 at N = 100. With `Pretrained_GNN` (trained - with labels, on Delphes samples that include the GamGam processes): the - best result at every N run (100 to 10000; not run at 50000), +6.9 at - N = 100. -* Meta-training on Delphes tasks does not help. TabPFN with LoRA loses 1.9 - at N = 1000 against the same network with zero updates (its direct - forward costs at most 0.4 against the classifier's preprocessing and - ensembling). Our PFN is below TabPFN and gets worse with more context - (76 at N = 1000, inside the meta-training context sizes 16 to 2048): to - investigate before the held-out-processes run, which was not done. -* Of the label-free pretraining objectives, masked-object modelling is the - best, in probes and fine-tuned; JEPA's prediction in embedding space adds - nothing over it here. Both improve fine-tuning the transformer over - training it from scratch (+0.9 to +3.6) but stay below fine-tuning the - edge network from `Pretrained_GNN`. That gap mixes the architecture (from - scratch, the edge network is 5.7 above the transformer at N = 100) with - the objective (label-free or supervised); the like-for-like comparison, - JEPA of the edge network (`configs/delphes_jepa_edge.yaml`), was not run. -* Contrastive pretraining is below a random encoder in the k-NN probe at - every N and in the linear probe at N = 100 and 1000 (67.9 against 72.6, - 80.3 against 83.1); at N = 10000 its linear probe is 0.5 above (85.9 - against 85.4). Without the momentum-scale augmentation - (`model.scale_jitter=0`, 4 epochs) it is no better (72.4 / 84.3 at - N = 10000). - -Not done: - -* the held-out-processes meta-training (`task.prior_classes` without the - five GamGam processes) of the PFN, gated on it beating TabPFN, and of - TabPFN and TabICL; -* TabICL meta-training and `model.lora_rank=0`; -* fine-tuning from contrastive pretraining (only probed), and JEPA of the - edge network; -* more seeds, and the test split. - -The loss is a weighted mean per class present in the batch, averaged over -those classes, which balances the classes without explicit class weights. -Metrics (accuracy, one-vs-rest AUC averaged over classes, macro precision, -recall and F1) use events with positive weight. - -## Performance - -Graphs differ in size, which suits neither `torch.compile` nor CUDA graphs -out of the box: - -* `torch.compile(dynamic=True)` treats node and edge counts as symbolic - sizes, so one compiled graph serves all batches without padding. -* CUDA graph replay needs fixed shapes. Batches are copied into buffers - rounded up to multiples of 512 nodes and 4096 edges; the remainder forms 64 - padding graphs connected only to themselves, so the real graphs' outputs do - not change. A step is captured once per bucket, a few per epoch. The - optimizer runs in Adam's `capturable` form with device-tensor learning - rates, which the scheduler updates in place. -* bfloat16 autocast covers the forward pass; logits, loss, graph sums (which - divide by edge counts) and the optimizer stay float32. - -GamGam training step, 1024-event batches (about 5,600 nodes and 30,000 -edges), one A100-SXM4-80GB, `OMP_NUM_THREADS=8`, four loader workers (step -times measured before the move to flat modules, epoch times after): - -| Mode | ms per batch | Epoch with validation | -| --- | ---: | ---: | -| float32 eager | 19.2 | 37 s | -| CUDA graphs + TF32 | 8.5 | | -| `torch.compile` | 13.0 | | -| compile + CUDA graphs + TF32 | 6.1 | | -| compile + CUDA graphs + bfloat16 (`train.fast`) | 5.9 | 9 s (first: 59 s, compiling) | - -Eager training is bound by kernel launches (about a thousand small kernels -per step), not by arithmetic; bfloat16 alone is slower because of the extra -cast kernels, which `torch.compile` fuses. Results of the fast modes agree -with float32 to rounding, not bitwise (validation AUC 0.895 against 0.897 -after three epochs). Preparing the 2.5M GamGam events with 8 workers takes -about 9 s on a login node (40 s was seen on a compute node, likely a cold -file cache). - -PyTorch sizes its CPU thread pool to every visible core; on a 128-core node -that makes batch construction several times slower and CPU reductions -non-reproducible, so set `OMP_NUM_THREADS`. Single-threaded CPU training is -bitwise reproducible. - -## ATLAS Open Data GamGam versus Ho et al. (2026) - -The main benchmark reproduces the from-scratch baseline of Table 3 (5 -classes, 2.5M events, 80:10:10 split, five seeds). Choices the paper does not -specify: - -* electrons and muons are split from `lep_*` by `lep_type` (11, 13); -* b-tagging is `jet_MV2c10 > 0.8244273`, the 70% working point; -* momenta and energies are scaled by 1e-4 (MeV to the pretraining unit of - 10 GeV); -* the fold is `eventNumber % 10`: folds 0-7 train, 8 validates, 9 tests; -* training is unweighted (23% of ttH events have `mcWeight = -1`). - -The paper's stopping rule (5 epochs without training-loss improvement) never -fires within 200 epochs on this task; validation AUC peaks earlier (epoch 67 -for seed 1) and `best.pt` keeps that epoch. - -| Test split | Accuracy | AUC | F1 | Precision | Recall | -| --- | ---: | ---: | ---: | ---: | ---: | -| Paper (5 seeds) | 71.8 ± 0.2 | 91.2 | 63.0 | 62.2 | 66.3 | -| This package, seed 1 (CUDA graphs + TF32) | 71.1 | 91.3 | 62.6 | 62.4 | 66.7 | diff --git a/docs/usage.md b/docs/usage.md deleted file mode 100644 index eeca9291cbcef5cc04aa7645e16c03a8054d78b3..0000000000000000000000000000000000000000 --- a/docs/usage.md +++ /dev/null @@ -1,255 +0,0 @@ -# Using GNN4Colliders - -GNN4Colliders trains event classifiers on reconstructed collider objects: - -```text -ROOT files --prepare--> event cache (one row per object) --> graph or token batches --> model -``` - -## Install - -Python 3.14 and PyTorch 2.14 (CUDA 13), managed by [uv](https://docs.astral.sh/uv/): - -```bash -uv sync --dev # includes pytest, ruff and the ONNX packages -uv sync --dev --extra tabpfn # also TabPFN (configs/gamgam_tabpfn.yaml) -uv sync --dev --extra tabicl # also TabICL (configs/gamgam_tabicl.yaml) -uv sync --dev --extra tabdpt # also TabDPT (configs/gamgam_tabdpt.yaml) -uv sync --dev --extra plots # matplotlib, for scripts/plot_label_efficiency.py -uv run pytest # about 2 minutes; GPU tests run when CUDA is available -uv run pytest --device mps # the fixture runs on an Apple GPU -``` - -## Commands - -Every command takes an experiment file and optional `dotted.key=value` -overrides (values are YAML): - -```bash -uv run gnn4colliders prepare configs/gamgam.yaml # ROOT -> data.cache -uv run gnn4colliders train configs/gamgam.yaml # -> run_dir/{best,last}.pt -uv run gnn4colliders evaluate configs/gamgam.yaml # -> run_dir/test_metrics.json -uv run gnn4colliders predict configs/gamgam.yaml # -> run_dir/test_predictions.npz -uv run gnn4colliders export configs/gamgam.yaml # -> run_dir/model.onnx -uv run gnn4colliders probe configs/gamgam_jepa_finetune.yaml \ - model.pretrained=outputs/delphes_jepa/last.pt data.train_events=1000 # -> run_dir/test_probe.json -``` - -`train` writes, in `run_dir` (default `outputs/`): - -* `config.yaml`: the resolved configuration; -* `history.jsonl`: one line per epoch with the loss and metrics of the - training and validation splits, the learning rate and the epoch time; -* `best.pt`: the epoch with the best validation `train.best_metric` (AUC by - default), `last.pt`: the latest epoch (resume with - `train.resume=/last.pt`), and `epoch_NNNN.pt` every - `train.save_every` epochs. - -`probe` fits k-NN and linear classifiers to the frozen event embeddings of -the encoder of `model.pretrained` (JEPA, contrastive or supervised; random -without) on `data.train_events` training events and reports their metrics -on `split`, to compare pretraining without fine-tuning. - -`data.eval_events=20000` evaluates (and validates during training, -predicts and probes) only that many random validation or test events, -seeded by `data.seed`: AUCs to about 0.003, in minutes, for quick runs. - -`evaluate`, `predict` and `export` use `checkpoint` (default `best.pt`, else -`last.pt`) on `split` (default `test`); `output` overrides the output file. -Predictions hold the logits, scores, predicted class, label, weight, fold, -source file and entry of every event; `gnn4colliders.inference.write_root_scores` -adds scores to a copy of a ROOT tree. - -## Experiments - -| File | Task | -| --- | --- | -| `configs/gamgam.yaml` | ATLAS Open Data GamGam, 5 Higgs production modes, from scratch (the paper baseline and main benchmark) | -| `configs/gamgam_finetune.yaml` | the same task fine-tuned from the pretrained multiclass model | -| `configs/gamgam_transformer.yaml` | the same task with a transformer with pair attention bias (`model.pair_bias=false` for a plain transformer; same event cache) | -| `configs/gamgam_deep_sets.yaml` | the same task with Deep Sets, the baseline without object interactions (same event cache) | -| `configs/tth_cp_even_odd.yaml` | CP-even versus CP-odd ttH (Delphes, binary; local files) | -| `configs/delphes_pretraining.yaml` | 12-class pretraining on the Delphes samples | -| `configs/delphes_jepa.yaml` | JEPA pretraining of the transformer encoder on the same samples, without labels (same event cache) | -| `configs/delphes_jepa_edge.yaml` | the same for the edge network (fine-tune with `configs/gamgam_finetune.yaml model.pretrained=outputs/delphes_jepa_edge/last.pt`) | -| `configs/delphes_contrastive.yaml` | contrastive (SimCLR) pretraining of the transformer encoder on the same samples (`model.backbone=edge_network` for the edge network) | -| `configs/delphes_masked.yaml` | masked object modelling: JEPA predicting hidden objects' kinematics | -| `configs/gamgam_jepa_finetune.yaml` | GamGam with the JEPA-pretrained encoder (`model.pretrained=null`: from scratch); also takes contrastive and masked-object encoders, and `probe` | -| `configs/gamgam_tabpfn.yaml` | GamGam classified in context by TabPFN-3.5, no training (`evaluate` only) | -| `configs/gamgam_tabicl.yaml` | the same by TabICLv2 (BSD-3 weights, downloaded from Hugging Face on first use) | -| `configs/gamgam_tabdpt.yaml` | the same by TabDPT (Apache-2.0 weights; `model.tabdpt_context_size` retrieves each event's context) | -| `configs/delphes_pfn.yaml` | meta-training of our own PFN on tasks drawn from the Delphes samples (same event cache) | -| `configs/gamgam_pfn.yaml` | GamGam classified in context by that PFN, no training (`evaluate` only) | -| `configs/delphes_tabpfn.yaml` | meta-training of TabPFN-3.5 (LoRA by default) on the same tasks; evaluate with `configs/gamgam_tabpfn.yaml checkpoint=outputs/delphes_tabpfn/best.pt` | -| `configs/delphes_tabicl.yaml` | the same for TabICL; evaluate with `configs/gamgam_tabicl.yaml checkpoint=outputs/delphes_tabicl/best.pt` | - -Each file documents its keys; `src/gnn4colliders/config.py` lists every key -with its default. Common overrides: - -```bash -train.fast=true # torch.compile + CUDA graphs (edge network) + bfloat16 + TF32 -train.seed=3 data.seed=3 # initialization and shuffling seeds -train.max_epochs=100 -model.pretrained=Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt -model.freeze_backbone=true # train only the new classifier -data.files=[{path: my.root, label: 0}] data.hf_repo=null # local inputs -``` - -### Data - -Large inputs belong on scratch storage, not in the repository: set -`HF_HOME` (where Hugging Face downloads go) and pass `data.cache=...`, e.g. -on Perlmutter `HF_HOME=$SCRATCH/.cache/huggingface` and -`data.cache=$SCRATCH/gnn4colliders/delphes/events.pt` (the 105M Delphes -events take 52 GB of ROOT files and a 21 GB cache). - -`data.files` lists the ROOT files and their class labels (or a label branch -name). With `data.hf_repo` the paths are inside that Hugging Face dataset at -`data.hf_revision` and are downloaded once and checked against `sha256` and -`size_bytes`. - -Each object becomes a row of `pt, eta, phi, energy, btag, charge, node_type` -built from `feature_branches` (one branch or constant per object type, -`CALC_E` = pt cosh(eta), `NODE_TYPE` = object-type index) and multiplied by -`scales`. `derived_branches` select objects or threshold tagger scores before -that (the GamGam file splits `lep_*` into electrons and muons). The fold of an -event is `fold_var`, modulo `fold_modulus` when set, and -`train_folds`/`validation_folds`/`test_folds` select the splits. Events are -unweighted unless `weight_var` names a weight branch. - -### In-context models and label efficiency - -`data.train_events=N` keeps N training events, an equal number per class -(smaller N give subsets of larger ones). Supervised models train on them; -the in-context models (`model.name: tabpfn` or `pfn`) use them as the -labelled context and are only evaluated: - -```bash -uv run gnn4colliders evaluate configs/gamgam_tabpfn.yaml split=validation data.train_events=1000 -uv run gnn4colliders train configs/delphes_pfn.yaml # meta-training -uv run gnn4colliders evaluate configs/gamgam_pfn.yaml split=validation data.train_events=1000 -uv run gnn4colliders train configs/gamgam.yaml data.train_events=1000 # supervised reference -``` - -`scripts/label_efficiency.py` runs such a sweep over several `N` and -configs (restartable) and prints validation AUC and seconds per run; the -metrics JSON records `seconds` (setup, inference) and, on CUDA, -`peak_memory_gib`: - -```bash -uv run python scripts/label_efficiency.py configs/gamgam_tabpfn.yaml configs/gamgam.yaml \ - --events 100 1000 10000 50000 -``` - -TabPFN-2.5 and later weights are non-commercial (research and internal -benchmarking are allowed). tabpfn downloads them after a one-time login at -https://ux.priorlabs.ai (token in `~/.cache/tabpfn/auth_token` or -`TABPFN_TOKEN`), but a checkpoint already in its cache directory loads -without the login, so on compute nodes download it from Hugging Face: - -```bash -export TABPFN_MODEL_CACHE_DIR=$SCRATCH/.cache/tabpfn # default ~/Library/Caches/tabpfn on macOS -uvx --from huggingface_hub hf download Prior-Labs/tabpfn_3_5 config.json \ - tabpfn-v3.5-20260909.safetensors tabpfn-v3.5-20260909_multiclass.safetensors \ - --local-dir $TABPFN_MODEL_CACHE_DIR -``` - -`model.tabpfn_checkpoint=tabpfn-v3.5-20260909_multiclass.safetensors` -selects the multiclass checkpoint. `model.tabpfn_version=v2` needs no login. - -TabDPT retrieval (`model.tabdpt_context_size`) on macOS needs -`OMP_NUM_THREADS=1`: faiss and torch load separate OpenMP runtimes. - -Tabular models can use a pretrained encoder's event embeddings instead of -or beside the rows: `model.tabular_inputs=embeddings` (or `both`) with -`model.pretrained=outputs/delphes_jepa/last.pt` and the encoder's sizes -(those of `configs/delphes_jepa.yaml`: `model.hid_size=128 model.n_layers=4 -model.n_heads=8 model.ffn_size=512 model.pair_bias=true`). - -TabPFN and TabICL classify the same table rows (`model.slots` objects of -each type). Two options add physics they cannot learn from a few events: -`model.relative_phi=true` measures every azimuth from the leading object's, -and `model.pair_objects=4` appends the invariant mass and Delta R of every -pair of the four leading objects (12 columns). Compare with and without, -one run name per variant: - -```bash -uv run python scripts/label_efficiency.py configs/gamgam_tabicl.yaml --events 100 1000 \ - --override name=gamgam_tabicl_physics model.relative_phi=true model.pair_objects=4 -``` - -TabPFN-3.5 can also be meta-trained on the Delphes tasks, like our PFN. -By default only LoRA updates of its in-context blocks train -(`model.lora_rank`, 8; `model.lora_rank=0` trains all 219M weights), and -checkpoints hold only those (33 MB with the optimizer state, against -2.6 GB for all weights). With a `checkpoint`, TabPFN is -evaluated with its own forward pass, not the classifier: - -```bash -uv run gnn4colliders train configs/delphes_tabpfn.yaml -uv run gnn4colliders evaluate configs/gamgam_tabpfn.yaml split=validation \ - data.train_events=1000 checkpoint=outputs/delphes_tabpfn/best.pt -``` - -Activations are recomputed in the backward pass. On a 32 GB Apple M2 Max -one task of 2048 context events and 256 queries takes about 4.5 s and -16 GiB per step with all weights trained, so the Mac is for development; -meta-train on a GPU node. - -### Fine-tuning - -`Pretrained_GNN/` holds the pretrained models of the paper, the final -epoch of each: `multiclass_pretrained_model_12/model_epoch_71.pt` and -`multilabel_pretrained_model_41/model_epoch_44.pt` (with their training -configuration and log). With `model.pretrained` the edge network is loaded from -that checkpoint and gets a new `model.out_size` classifier; `model.hid_size`, -`n_layers` and `n_proc_steps` must match it (64, 4, 4). A JEPA checkpoint -(`configs/delphes_jepa*.yaml`) loads the same way: its context encoder -initializes the edge network or the transformer encoder. - -## GPUs, Perlmutter and DDP - -`train.device: auto` uses CUDA, else an Apple GPU (MPS), else the CPU. MPS -is for development on fixtures and small runs (no float64: metrics are -computed on the CPU; no DDP, CUDA graphs or TF32). - -On a single GPU, `train.fast=true` is about 3x faster than float32 for the -GamGam model (see [design.md](design.md)). Set `OMP_NUM_THREADS` to the CPU -allocation and, on Perlmutter compute nodes, `UV_CACHE_DIR=$SCRATCH/.cache/uv` -(uv cannot lock its home-directory cache there). - -```bash -# One training per seed, each followed by test-split evaluation: -for seed in 1 2 3 4 5; do - sbatch -A --export=ALL,SEED=$seed scripts/train.sh train.fast=true -done -uv run python scripts/summarize_ensemble.py outputs/gamgam/seed* # mean +- std -``` - -The study of in-context models and pretraining has one wrapper per kind of -job, each on one GPU (the header of each script lists its variables): - -```bash -# Label-efficiency sweep of released in-context models and supervised references: -sbatch -A --export=ALL,CONFIGS="configs/gamgam_tabpfn.yaml configs/gamgam_tabicl.yaml configs/gamgam.yaml" \ - scripts/label_efficiency.sh -# Meta-training, then GamGam in context with the meta-trained model: -sbatch -A --export=ALL,CONFIG=configs/delphes_tabicl.yaml,EVAL_CONFIG=configs/gamgam_tabicl.yaml \ - scripts/meta_train.sh task.prior_classes=[1,6,7,8,9,10,11] -# Label-free pretraining, then k-NN and linear probes on GamGam: -sbatch -A --export=ALL,CONFIG=configs/delphes_contrastive.yaml scripts/pretrain.sh -# Linear CKA between the pretrained encoders, layer by layer (a login node is enough): -uv run python scripts/compare_encoders.py configs/gamgam_jepa_finetune.yaml --events 5000 \ - --encoders random=random jepa=outputs/delphes_jepa/last.pt \ - contrastive=outputs/delphes_contrastive/last.pt masked=outputs/delphes_masked/last.pt -# Curves and tables of every finished run: -uv run --extra plots python scripts/plot_label_efficiency.py --sweeps outputs/label_efficiency \ - --probes outputs/delphes_jepa outputs/delphes_contrastive outputs/delphes_masked -``` - -Under `torchrun` with more than one process, training is data parallel: -`data.batch_size` is per process, training events are sharded over ranks -(padded to equal length) and metrics are computed over all ranks' events. -CUDA graphs need a single process (so `train.fast=true` does too; DDP is -tested in float32 eager). `scripts/train_ddp.sh` launches all GPUs of -one or more nodes. diff --git a/plots/cka_layerwise.svg b/plots/cka_layerwise.svg deleted file mode 100644 index 2ec27e33555a77f8e566045062b86db033ae39b2..0000000000000000000000000000000000000000 --- a/plots/cka_layerwise.svg +++ /dev/null @@ -1,17500 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/plots/hyperparameter_optimization.svg b/plots/hyperparameter_optimization.svg deleted file mode 100644 index 233da7654f8ed96d9d4aac9a0311459419812bf1..0000000000000000000000000000000000000000 --- a/plots/hyperparameter_optimization.svg +++ /dev/null @@ -1,954 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/plots/trainingtime.svg b/plots/trainingtime.svg deleted file mode 100644 index 0ecf607f82c7a117e577f78afe2ddd6b01663c00..0000000000000000000000000000000000000000 --- a/plots/trainingtime.svg +++ /dev/null @@ -1,4622 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/pyproject.toml b/pyproject.toml deleted file mode 100644 index 360b57bc5cbed0db5d63fbe0bd08d27e3da9712d..0000000000000000000000000000000000000000 --- a/pyproject.toml +++ /dev/null @@ -1,80 +0,0 @@ -[build-system] -requires = ["setuptools>=68"] -build-backend = "setuptools.build_meta" - -[project] -name = "gnn4colliders" -dynamic = ["version"] -description = "Machine-learning tools for collider event analysis." -readme = "README.md" -requires-python = ">=3.14,<3.15" -dependencies = [ - "awkward>=2.0", - "huggingface-hub>=0.36,<1", - "numpy>=2.1", - "pyyaml>=6", - "torch==2.14.1", - "uproot>=5.0", -] - -[project.optional-dependencies] -onnx = [ - "onnx>=1.16,<2", - "onnxruntime>=1.18,<2", -] -tabpfn = [ - "tabpfn>=9.2,<10", -] -tabicl = [ - "tabicl>=2.2,<3", -] -tabdpt = [ - "tabdpt>=1.3.1,<2", -] -plots = [ - "matplotlib>=3.9,<4", -] - -[project.scripts] -gnn4colliders = "gnn4colliders.cli:main" - -[dependency-groups] -dev = [ - "pytest>=8", - "ruff>=0.6", - "onnx>=1.16,<2", - "onnxruntime>=1.18,<2", -] - -[tool.setuptools.packages.find] -where = ["src"] - -[tool.setuptools.dynamic] -version = { attr = "gnn4colliders.__version__" } - -[tool.uv.sources] -torch = [ - { index = "pytorch-cu130", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" }, -] - -[[tool.uv.index]] -name = "pytorch-cu130" -url = "https://download.pytorch.org/whl/cu130" -explicit = true - -[tool.pytest.ini_options] -testpaths = ["tests"] -addopts = "-ra" -markers = [ - "gpu: tests requiring CUDA", -] - -[tool.ruff] -line-length = 88 -target-version = "py314" -src = ["src"] - -[tool.ruff.lint] -select = ["E", "F", "I"] - -[tool.ruff.format] diff --git a/root_gnn_dgl/README.md b/root_gnn_dgl/README.md new file mode 100644 index 0000000000000000000000000000000000000000..7337e9ed4e171ccfae13100f48c23a4afbe3ac89 --- /dev/null +++ b/root_gnn_dgl/README.md @@ -0,0 +1,137 @@ +# root_gnn_dgl + +## Data Directory (for Hackathon) +`/global/cfs/projectdirs/trn007/lbl_atlas/data/` + +* `stats_all`: full statistics sample, ~10M events per process +* `stats_100K`: reduced statistics sample, 100K events per process +* `processed_graphs`: graphs that have already been processed +* `scores`: a copy of the samples along with the GNN scores for each event + +## Environment Setup + +The environment dependencies for this project are listed in `setup/environment.yml`. Follow the steps below to set up the environment: + +### Step 1: Install Conda +If you don’t already have Conda installed, install either Miniconda (lightweight) or Anaconda (full version): + +- **Miniconda**: Download and install from [https://docs.conda.io/en/latest/miniconda.html](https://docs.conda.io/en/latest/miniconda.html). +- **Anaconda**: Download and install from [https://www.anaconda.com/products/distribution](https://www.anaconda.com/products/distribution). + +### Step 2: Clone the Repository +Clone this repository to your local machine: +```bash +git init +git lfs install +git clone https://huggingface.co/HWresearch/GNN4Colliders +``` +If you want to clone without large files - just their pointers +```bash +GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/HWresearch/GNN4Colliders +``` + +### Step 3: Create the Conda Environment +Use the `environment.yml` file to create the Conda environment: +```bash +conda env create -f setup/environment.yml -n +``` + +### Step 4: Activate the Environment +Activate the newly created environment: +```bash +conda activate +``` +Replace with the name of the environment specified in Step 4. + +### Step 5: Test the Environment +Run the `setup/test_setup.py` script to confirm that all packages needed for training are properly set up. +```bash +python setup/test_setup.py +``` +## Running the Demo +The demo training is an example of our ML workflow, consisting of training a pretrained model, then finetuning it for an analysis task, while also training a model for the analysis task from scratch. The config files for the demo are located in the directory `configs/stats_100K/`. The demo can be run on a login node on Perlmutter (if enough GPU memory is availble). + +To check login node GPU memory availability, use the command `nvidia-smi`. If there is not enough memory available, you can switch to another login node with the command `ssh login**`, where `**` is a number between 0 and 39. + +For better performance, it is recommended to run the training and inference of the demo on a shared interactive node, where you have access to one exclusive GPU. An interactive node can be requested using the shell script in `jobs/interactive.sh`. + +The pretraining for the demo is a multiclass classification training on 12 datasets corresponding to 12 distinct physics processes, containing 100,000 simulated collision events each. The pretraining is then fintuned on a binary classification task between two datasets containing 100,000 simulated collision events each for two different processes, called ttH CP Even and ttH CP Odd. + +The entire demo can be ran with the command +```bash +source run_demo.sh +``` + +This shell script can also be used as an example to run the entire workflow. + +## Data Preparation +The first step in the process is to convert the events stored in ROOT files into DGL graph objects. This conversion is handled automatically by the Dataset objects during their creation, provided the graph data has not already been saved to disk. To accomplish this, a simple script is used to initialize the relevant Dataset object and then exit. This script needs to be executed for each data chunk in each dataset being used for training. + +Below is an example of how to use the `scripts/prep_data.py` script: + +```bash +datasets=("ttH" "tHjb" "ggF" "VBF" "WH" "ZH" "ttyy" "tttt" "SingleT_schan" "ttbar" "ttW" "ttt") +chunks=3 + +for data in "${datasets[@]}"; do + python scripts/prep_data.py --config configs/demo/pretraining_multiclass.yaml --dataset "$data" --shuffle_mode --chunk 0 + for ((i=0; i python jobs/inference/run_inference.py`. diff --git a/root_gnn_dgl/configs/attention/ttH_CP_even_vs_odd.yaml b/root_gnn_dgl/configs/attention/ttH_CP_even_vs_odd.yaml new file mode 100755 index 0000000000000000000000000000000000000000..dc124be431b366eed524c31ac6fc11c2325f4f42 --- /dev/null +++ b/root_gnn_dgl/configs/attention/ttH_CP_even_vs_odd.yaml @@ -0,0 +1,58 @@ +Training_Name: ttH_CP_even_vs_odd +Training_Directory: trainings/attention/ttH_CP_even_vs_odd +Model: + module: models.GCN + class: Attention_Edge_Network + args: + hid_size: 64 + in_size: 7 + out_size: 1 + n_layers: 4 + n_proc_steps: 4 + dropout: 0 + num_heads: 2 +Training: + epochs: 500 + batch_size: 1024 + learning_rate: 0.0001 + gamma: 0.99 +Datasets: + ttH_CP_even: &dataset_defn + module: root_gnn_base.dataset + class: LazyDataset + shuffle_chunks: 3 + batch_size: 1024 + padding_mode: NODE + args: &dataset_args + name: ttH_CP_even + label: 0 + # weight_var: weight + chunks: 3 + buffer_size: 2 + file_names: ttH_NLO.root + tree_name: output + fold_var: Number + raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/ + save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/attention/ttH_CP_even_vs_odd/ + node_branch_names: + - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] + - [jet_eta, ele_eta, mu_eta, ph_eta, 0] + - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] + - CALC_E + - [jet_btag, 0, 0, 0, 0] + - [0, ele_charge, mu_charge, 0, 0] + - NODE_TYPE + node_branch_types: [vector, vector, vector, vector, single] + node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1] + folding: + n_folds: 4 + test: [0] + # validation: 1 + train: [1, 2, 3] + ttH_CP_odd: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttH_CP_odd + label: 1 + file_names: ttH_CPodd.root diff --git a/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml b/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml new file mode 100755 index 0000000000000000000000000000000000000000..ae56b014189afbba4886f421dda9aad12f6d467c --- /dev/null +++ b/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml @@ -0,0 +1,67 @@ +Training_Name: finetuning_ttH_CP_even_vs_odd +Training_Directory: trainings/stats_100K/finetuning_ttH_CP_even_vs_odd +Model: + module: models.GCN + class: Transferred_Learning_Finetuning + args: + pretraining_path: trainings/stats_100K/pretraining_multiclass/model_epoch_100.pt # update to the last epoch of the pretraining + pretraining_model: + module: models.GCN + class: Edge_Network + args: + hid_size: 64 + in_size: 7 + out_size: 12 + n_layers: 4 + n_proc_steps: 4 + hid_size: 64 + in_size: 7 + out_size: 1 + n_layers: 4 + n_proc_steps: 4 + dropout: 0 +Training: + epochs: 500 + batch_size: 1024 + learning_rate: 0.0001 + gamma: 0.99 +Datasets: + ttH_CP_even: &dataset_defn + module: root_gnn_base.dataset + class: LazyDataset + shuffle_chunks: 3 + batch_size: 1024 + padding_mode: NONE #one of STEPS, FIXED, or NONE + args: &dataset_args + name: ttH_CP_even + label: 0 + # weight_var: weight + chunks: 3 + buffer_size: 2 + file_names: ttH_NLO.root + tree_name: output + fold_var: Number + raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/ + save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/ + node_branch_names: + - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] + - [jet_eta, ele_eta, mu_eta, ph_eta, 0] + - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] + - CALC_E + - [jet_btag, 0, 0, 0, 0] + - [0, ele_charge, mu_charge, 0, 0] + - NODE_TYPE + node_branch_types: [vector, vector, vector, vector, single] + node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1] + folding: + n_folds: 4 + test: [0] + # validation: 1 + train: [1, 2, 3] + ttH_CP_odd: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttH_CP_odd + label: 1 + file_names: ttH_CPodd.root diff --git a/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml b/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c182c8588431ebd1016178fc7f8b9c1dda543a0b --- /dev/null +++ b/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml @@ -0,0 +1,134 @@ +Training_Name: pretraining_multiclass +Training_Directory: trainings/stats_100K/pretraining_multiclass/ +Model: + module: models.GCN + class: Edge_Network + args: + hid_size: 64 + in_size: 7 + out_size: 12 + n_layers: 4 + n_proc_steps: 4 + dropout: 0 +Loss: + module: torch.nn + class: CrossEntropyLoss + args: {} + finish: + module: torch.nn + class: Softmax + args: {dim: 1} +Training: + epochs: 500 + batch_size: 1024 + learning_rate: 0.0001 + gamma: 0.99 +Datasets: + ttH: &dataset_defn + module: root_gnn_base.dataset + class: LazyDataset + shuffle_chunks: 3 + padding_mode: NONE #one of STEPS, FIXED, or NONE + args: &dataset_args + name: ttH + label: 0 + # weight_var: weight + chunks: 3 + buffer_size: 2 + file_names: ttH_NLO_inc.root + tree_name: output + fold_var: Number + raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/ + save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/pretraining_multiclass/ + node_branch_names: + - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] + - [jet_eta, ele_eta, mu_eta, ph_eta, 0] + - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] + - CALC_E + - [jet_btag, 0, 0, 0, 0] + - [0, ele_charge, mu_charge, 0, 0] + - NODE_TYPE + node_branch_types: [vector, vector, vector, vector, single] + node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1] + folding: + n_folds: 4 + test: [0] + # validation: [1] + train: [1, 2, 3] + tHjb: + <<: *dataset_defn + args: + <<: *dataset_args + name: tHjb + label: 1 + file_names: tHjb_NLO_inc.root + ggF: + <<: *dataset_defn + args: + <<: *dataset_args + name: ggF + label: 2 + file_names: ggF_NLO_inc.root + VBF: + <<: *dataset_defn + args: + <<: *dataset_args + name: VBF + label: 3 + file_names: VBF_NLO_inc.root + WH: + <<: *dataset_defn + args: + <<: *dataset_args + name: WH + label: 4 + file_names: WH_NLO_inc.root + ZH: + <<: *dataset_defn + args: + <<: *dataset_args + name: ZH + label: 5 + file_names: ZH_NLO_inc.root + ttyy: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttyy_ch + label: 6 + file_names: 'ttyy.root' + tttt: + <<: *dataset_defn + args: + <<: *dataset_args + name: tttt + label: 7 + file_names: 'tttt.root' + SingleT_schan: + <<: *dataset_defn + args: + <<: *dataset_args + name: SingleT_schan + label: 8 + file_names: 'SingleT_schan.root' + ttbar: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttbar + label: 9 + file_names: 'ttbar.root' + ttW: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttW + label: 10 + file_names: 'ttW.root' + ttt: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttt + label: 11 + file_names: 'ttt.root' \ No newline at end of file diff --git a/root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml b/root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml new file mode 100755 index 0000000000000000000000000000000000000000..5fcb903b08a68a8640bf65ee4aab0c81a9d91d66 --- /dev/null +++ b/root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml @@ -0,0 +1,57 @@ +Training_Name: ttH_CP_even_vs_odd +Training_Directory: trainings/stats_100K/ttH_CP_even_vs_odd +Model: + module: models.GCN + class: Edge_Network + args: + hid_size: 64 + in_size: 7 + out_size: 1 + n_layers: 4 + n_proc_steps: 4 + dropout: 0 +Training: + epochs: 500 + batch_size: 1024 + learning_rate: 0.0001 + gamma: 0.99 +Datasets: + ttH_CP_even: &dataset_defn + module: root_gnn_base.dataset + class: LazyDataset + shuffle_chunks: 3 + batch_size: 1024 + padding_mode: NONE #one of STEPS, FIXED, or NONE + args: &dataset_args + name: ttH_CP_even + label: 0 + # weight_var: weight + chunks: 3 + buffer_size: 2 + file_names: ttH_NLO.root + tree_name: output + fold_var: Number + raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_100K/ + save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_100K/ttH_CP_even_vs_odd/ + node_branch_names: + - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] + - [jet_eta, ele_eta, mu_eta, ph_eta, 0] + - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] + - CALC_E + - [jet_btag, 0, 0, 0, 0] + - [0, ele_charge, mu_charge, 0, 0] + - NODE_TYPE + node_branch_types: [vector, vector, vector, vector, single] + node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1] + folding: + n_folds: 4 + test: [0] + # validation: 1 + train: [1, 2, 3] + ttH_CP_odd: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttH_CP_odd + label: 1 + file_names: ttH_CPodd.root diff --git a/root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml b/root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml new file mode 100755 index 0000000000000000000000000000000000000000..4ea8a0f4e8f4c1fb2467574a7f3730a7cf0768e9 --- /dev/null +++ b/root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml @@ -0,0 +1,62 @@ +Training_Name: finetuning_ttH_CP_even_vs_odd +Training_Directory: trainings/stats_all/finetuning_ttH_CP_even_vs_odd +Model: + module: models.GCN + class: Transferred_Learning_Finetuning + args: + pretraining_path: trainings/stats_100K/pretraining_multiclass/model_epoch_100.pt # update to the last epoch of the pretraining + pretraining_model: + module: models.GCN + class: Edge_Network + args: + hid_size: 64 + in_size: 7 + out_size: 12 + n_layers: 4 + n_proc_steps: 4 + hid_size: 64 + in_size: 7 + out_size: 1 + n_layers: 4 + n_proc_steps: 4 + dropout: 0 +Datasets: + ttH_CP_even: &dataset_defn + module: root_gnn_base.dataset + class: LazyDataset + shuffle_chunks: 10 + batch_size: 1024 + padding_mode: NONE #one of STEPS, FIXED, or NONE + args: &dataset_args + name: ttH_CP_even + label: 0 + # weight_var: weight + chunks: 10 + buffer_size: 3 + file_names: ttH_NLO.root + tree_name: output + fold_var: Number + raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/ + save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/ttH_CP_even_vs_odd/ + node_branch_names: + - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] + - [jet_eta, ele_eta, mu_eta, ph_eta, 0] + - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] + - CALC_E + - [jet_btag, 0, 0, 0, 0] + - [0, ele_charge, mu_charge, 0, 0] + - NODE_TYPE + node_branch_types: [vector, vector, vector, vector, single] + node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1] + folding: + n_folds: 4 + test: [0] + # validation: 1 + train: [1, 2, 3] + ttH_CP_odd: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttH_CP_odd + label: 1 + file_names: ttH_CPodd.root diff --git a/root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml b/root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml new file mode 100644 index 0000000000000000000000000000000000000000..8cc480f9f3956ee05dcc033f89ed05d217658de2 --- /dev/null +++ b/root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml @@ -0,0 +1,134 @@ +Training_Name: pretraining_multiclass +Training_Directory: trainings/stats_all/pretraining_multiclass/ +Model: + module: models.GCN + class: Edge_Network + args: + hid_size: 64 + in_size: 7 + out_size: 12 + n_layers: 4 + n_proc_steps: 4 + dropout: 0 +Loss: + module: torch.nn + class: CrossEntropyLoss + args: {} + finish: + module: torch.nn + class: Softmax + args: {dim: 1} +Training: + epochs: 500 + batch_size: 1024 + learning_rate: 0.0001 + gamma: 0.99 +Datasets: + ttH: &dataset_defn + module: root_gnn_base.dataset + class: LazyDataset + shuffle_chunks: 10 + padding_mode: NONE #one of STEPS, FIXED, or NONE + args: &dataset_args + name: ttH + label: 0 + # weight_var: weight + chunks: 10 + buffer_size: 3 + file_names: ttH_NLO_inc.root + tree_name: output + fold_var: Number + raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/ + save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/pretraining_multiclass/ + node_branch_names: + - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] + - [jet_eta, ele_eta, mu_eta, ph_eta, 0] + - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] + - CALC_E + - [jet_btag, 0, 0, 0, 0] + - [0, ele_charge, mu_charge, 0, 0] + - NODE_TYPE + node_branch_types: [vector, vector, vector, vector, single] + node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1] + folding: + n_folds: 4 + test: [0] + # validation: [1] + train: [1, 2, 3] + tHjb: + <<: *dataset_defn + args: + <<: *dataset_args + name: tHjb + label: 1 + file_names: tHjb_NLO_inc.root + ggF: + <<: *dataset_defn + args: + <<: *dataset_args + name: ggF + label: 2 + file_names: ggF_NLO_inc.root + VBF: + <<: *dataset_defn + args: + <<: *dataset_args + name: VBF + label: 3 + file_names: VBF_NLO_inc.root + WH: + <<: *dataset_defn + args: + <<: *dataset_args + name: WH + label: 4 + file_names: WH_NLO_inc.root + ZH: + <<: *dataset_defn + args: + <<: *dataset_args + name: ZH + label: 5 + file_names: ZH_NLO_inc.root + ttyy: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttyy_ch + label: 6 + file_names: 'ttyy.root' + tttt: + <<: *dataset_defn + args: + <<: *dataset_args + name: tttt + label: 7 + file_names: 'tttt.root' + SingleT_schan: + <<: *dataset_defn + args: + <<: *dataset_args + name: SingleT_schan + label: 8 + file_names: 'SingleT_schan.root' + ttbar: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttbar + label: 9 + file_names: 'ttbar.root' + ttW: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttW + label: 10 + file_names: 'ttW.root' + ttt: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttt + label: 11 + file_names: 'ttt.root' \ No newline at end of file diff --git a/root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml b/root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml new file mode 100755 index 0000000000000000000000000000000000000000..0d074acab903084ce8895a47ba231f28732f9094 --- /dev/null +++ b/root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml @@ -0,0 +1,57 @@ +Training_Name: ttH_CP_even_vs_odd +Training_Directory: trainings/stats_all/ttH_CP_even_vs_odd +Model: + module: models.GCN + class: Edge_Network + args: + hid_size: 64 + in_size: 7 + out_size: 1 + n_layers: 4 + n_proc_steps: 4 + dropout: 0 +Training: + epochs: 500 + batch_size: 1024 + learning_rate: 0.0001 + gamma: 0.99 +Datasets: + ttH_CP_even: &dataset_defn + module: root_gnn_base.dataset + class: LazyDataset + shuffle_chunks: 10 + batch_size: 1024 + padding_mode: NONE #one of STEPS, FIXED, or NONE + args: &dataset_args + name: ttH_CP_even + label: 0 + # weight_var: weight + chunks: 10 + buffer_size: 3 + file_names: ttH_NLO.root + tree_name: output + fold_var: Number + raw_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/stats_all/ + save_dir: /global/cfs/projectdirs/trn007/lbl_atlas/data/processed_graphs/stats_all/ttH_CP_even_vs_odd/ + node_branch_names: + - [jet_pt, ele_pt, mu_pt, ph_pt, MET_met] + - [jet_eta, ele_eta, mu_eta, ph_eta, 0] + - [jet_phi, ele_phi, mu_phi, ph_phi, MET_phi] + - CALC_E + - [jet_btag, 0, 0, 0, 0] + - [0, ele_charge, mu_charge, 0, 0] + - NODE_TYPE + node_branch_types: [vector, vector, vector, vector, single] + node_feature_scales: [1e-1, 1, 1, 1e-1, 1, 1, 1] + folding: + n_folds: 4 + test: [0] + # validation: 1 + train: [1, 2, 3] + ttH_CP_odd: + <<: *dataset_defn + args: + <<: *dataset_args + name: ttH_CP_odd + label: 1 + file_names: ttH_CPodd.root diff --git a/root_gnn_dgl/jobs/cpu.sh b/root_gnn_dgl/jobs/cpu.sh new file mode 100644 index 0000000000000000000000000000000000000000..ae17a41d66445d21ba5ccfb636ccda113b921eb6 --- /dev/null +++ b/root_gnn_dgl/jobs/cpu.sh @@ -0,0 +1 @@ +salloc --nodes=1 --ntasks=64 --cpus-per-task=1 --qos=interactive --time=04:00:00 --constraint=cpu --account=trn007 diff --git a/root_gnn_dgl/jobs/inference/run_inference.py b/root_gnn_dgl/jobs/inference/run_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..b393af6d25073ad269af1d82d0e912d93587ad7d --- /dev/null +++ b/root_gnn_dgl/jobs/inference/run_inference.py @@ -0,0 +1,214 @@ +import os +import subprocess +from mpi4py import MPI + +# Define the number of GPUs per node +GPUS_PER_NODE = 4 + +# Define your files, configs, and branch namesfiles = [ +files = [ + "merged_aMCPy8_tHjb125_CP_0_AF3_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_0_AF3_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_0_FullSim_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_0_FullSim_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_15_AF3_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_15_AF3_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_30_AF3_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_30_AF3_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_45_AF3_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_45_AF3_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_45_FullSim_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_45_FullSim_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_60_AF3_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_60_AF3_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_75_AF3_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_75_AF3_lep_scaled.root", + "merged_aMCPy8_tHjb125_CP_90_AF3_had_scaled.root", + "merged_aMCPy8_tHjb125_CP_90_AF3_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_0_AF3_had_scaled.root", + "merged_aMCPy8_ttH125_CP_0_AF3_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_0_FullSim_had_scaled.root", + "merged_aMCPy8_ttH125_CP_0_FullSim_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_15_AF3_had_scaled.root", + "merged_aMCPy8_ttH125_CP_15_AF3_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_30_AF3_had_scaled.root", + "merged_aMCPy8_ttH125_CP_30_AF3_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_45_AF3_had_scaled.root", + "merged_aMCPy8_ttH125_CP_45_AF3_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_45_FullSim_had_scaled.root", + "merged_aMCPy8_ttH125_CP_45_FullSim_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_60_AF3_had_scaled.root", + "merged_aMCPy8_ttH125_CP_60_AF3_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_75_AF3_had_scaled.root", + "merged_aMCPy8_ttH125_CP_75_AF3_lep_scaled.root", + "merged_aMCPy8_ttH125_CP_90_AF3_had_scaled.root", + "merged_aMCPy8_ttH125_CP_90_AF3_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_0_AF3_had_scaled.root", + "merged_aMCPy8_tWH125_CP_0_AF3_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_0_FullSim_had_scaled.root", + "merged_aMCPy8_tWH125_CP_0_FullSim_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_15_AF3_had_scaled.root", + "merged_aMCPy8_tWH125_CP_15_AF3_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_30_AF3_had_scaled.root", + "merged_aMCPy8_tWH125_CP_30_AF3_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_45_AF3_had_scaled.root", + "merged_aMCPy8_tWH125_CP_45_AF3_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_45_FullSim_had_scaled.root", + "merged_aMCPy8_tWH125_CP_45_FullSim_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_60_AF3_had_scaled.root", + "merged_aMCPy8_tWH125_CP_60_AF3_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_75_AF3_had_scaled.root", + "merged_aMCPy8_tWH125_CP_75_AF3_lep_scaled.root", + "merged_aMCPy8_tWH125_CP_90_AF3_had_scaled.root", + "merged_aMCPy8_tWH125_CP_90_AF3_lep_scaled.root", + "merged_hadnti_data.root", + "merged_hadti_data.root", + "merged_lepnti_data.root", + "merged_lepti_data.root", + "merged_MGPy8_ttyy_combined_FullSim_had_scaled.root", + "merged_MGPy8_ttyy_combined_FullSim_lep_scaled.root", + "merged_PhPy8_4FS_A14NNPDF_bbH125_FullSim_had_scaled.root", + "merged_PhPy8_4FS_A14NNPDF_bbH125_FullSim_lep_scaled.root", + "merged_PhPy8_ggH125_NNLOPS_FullSim_had_scaled.root", + "merged_PhPy8_ggH125_NNLOPS_FullSim_lep_scaled.root", + "merged_PhPy8_ggZH125_FullSim_had_scaled.root", + "merged_PhPy8_ggZH125_FullSim_lep_scaled.root", + "merged_PhPy8_VBFH125_FullSim_had_scaled.root", + "merged_PhPy8_VBFH125_FullSim_lep_scaled.root", + "merged_PhPy8_WmH125J_MINLO_FullSim_had_scaled.root", + "merged_PhPy8_WmH125J_MINLO_FullSim_lep_scaled.root", + "merged_PhPy8_WpH125J_MINLO_FullSim_had_scaled.root", + "merged_PhPy8_WpH125J_MINLO_FullSim_lep_scaled.root", + "merged_PhPy8_ZH125J_MINLO_FullSim_had_scaled.root", + "merged_PhPy8_ZH125J_MINLO_FullSim_lep_scaled.root", +] +# Configurations for each type of file +config = [ + "configs/model_configs/run_3_ttH/v06/sb_yukawa_no_weights.yaml", + "configs/model_configs/run_3_ttH/v06/sb_yukawa_weights.yaml", + "configs/model_configs/run_3_ttH/v06/tHjb_CP_no_weights.yaml", + "configs/model_configs/run_3_ttH/v06/tHjb_CP_weights.yaml", + "configs/model_configs/run_3_ttH/v06/ttH_CP_no_weights.yaml", + "configs/model_configs/run_3_ttH/v06/ttH_CP_weights.yaml", + "configs/model_configs/run_3_ttH/v06/ttH_tHjb_tWH_CP_no_weights.yaml", + "configs/model_configs/run_3_ttH/v06/ttH_tHjb_tWH_CP_weights.yaml", + "configs/model_configs/run_3_ttH/v06/tWH_CP_no_weights.yaml", + "configs/model_configs/run_3_ttH/v06/tWH_CP_weights.yaml", + +] + +# Branch names for each type of file +branch_name = [ + "sb_yukawa_no_weights_score", + "sb_yukawa_weights_score", + "tHjb_CP_no_weights_score", + "tHjb_CP_weights_score", + "ttH_CP_no_weights_score", + "ttH_CP_weights_score", + "ttH_tHjb_tWH_CP_no_weights_score", + "ttH_tHjb_tWH_CP_weights_score", + "tWH_CP_no_weights_score", + "tWH_CP_weights_score", +] + +# MPI Initialization +comm = MPI.COMM_WORLD +rank = comm.Get_rank() # Process rank (unique ID for each process) +size = comm.Get_size() # Total number of processes across all nodes + +# Master process manages the task queue +def master(): + num_workers = size - 1 # Exclude the master process + task_queue = files[:] # Copy the list of files + workers_busy = [False] * num_workers # Track if workers are busy + + print(f"Master process starting with {num_workers} workers.") + + # Distribute initial tasks to workers + for worker_rank in range(1, size): + if task_queue: + # Send a file to the worker + file = task_queue.pop(0) + comm.send(file, dest=worker_rank, tag=0) + workers_busy[worker_rank - 1] = True + else: + # No tasks left, send termination signal + comm.send(None, dest=worker_rank, tag=1) + + # Process results and assign new tasks dynamically + while any(workers_busy): + # Receive results from workers + status = MPI.Status() + result = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status) + worker_rank = status.Get_source() + workers_busy[worker_rank - 1] = False # Worker is now free + + # Assign a new task if available + if task_queue: + file = task_queue.pop(0) + comm.send(file, dest=worker_rank, tag=0) + workers_busy[worker_rank - 1] = True + else: + # No tasks left, send termination signal + comm.send(None, dest=worker_rank, tag=1) + + print("Master process: All tasks completed.") + +# Worker processes handle tasks dynamically +def worker(): + gpu_id = (rank - 1) % GPUS_PER_NODE # Assign GPU ID based on rank + print(f"Worker {rank} starting, GPU ID: {gpu_id}", flush=True) + + while True: + # Receive a task from the master + file = comm.recv(source=0, tag=MPI.ANY_TAG) + if file is None: + # Termination signal received + print(f"Worker {rank}: Termination signal received.") + break + + # Initialize config and branch_name for each file + local_config = config[:] # Create a copy of the global config + local_branch_name = branch_name[:] # Create a copy of the global branch_name + + # Determine file type and append appropriate configurations + if "had" in file: + local_config += had_config # Append had-specific configurations + local_branch_name += had_branch_name # Append had-specific branch names + elif "lep" in file: + local_config += lep_config # Append lep-specific configurations + local_branch_name += lep_branch_name # Append lep-specific branch names + + # Process the file + print(f"Worker {rank}: Processing file {file} on GPU {gpu_id}", flush=True) + print(f"branches = {local_branch_name}") + target = f"/global/cfs/projectdirs/atlas/joshua/ttHCP/ntuples/v02/preselection/merged/{file}" + destination = f"/global/cfs/projectdirs/atlas/joshua/ttHCP/ntuples/scores/gnn_v06/{file}" + + # Set the GPU for this process + os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) + + # Build the command + command = [ + "python", "scripts/inference.py", + "--target", target, + "--destination", destination, + "--config", *local_config, + "--branch_name", *local_branch_name, + "--chunks", "1", + "--chunkno", "0", + "--write" + ] + + # Execute the command + subprocess.run(command) + + # Notify the master that the task is complete + comm.send(f"Worker {rank}: Finished {file}", dest=0, tag=0) + +# Main function +if __name__ == "__main__": + if rank == 0: + master() + else: + worker() \ No newline at end of file diff --git a/root_gnn_dgl/jobs/interactive.sh b/root_gnn_dgl/jobs/interactive.sh new file mode 100644 index 0000000000000000000000000000000000000000..8197dd547c8fd7292172146af10d734a68ddaadb --- /dev/null +++ b/root_gnn_dgl/jobs/interactive.sh @@ -0,0 +1 @@ +salloc --nodes 1 --qos shared_interactive --time 04:00:00 --constraint gpu --account=trn007 --gres=gpu:1 diff --git a/root_gnn_dgl/jobs/prep_data/prep_data.sh b/root_gnn_dgl/jobs/prep_data/prep_data.sh new file mode 100755 index 0000000000000000000000000000000000000000..16919d0003d40d2ff9fcd53f2662251448373d6a --- /dev/null +++ b/root_gnn_dgl/jobs/prep_data/prep_data.sh @@ -0,0 +1,20 @@ +#!/bin/bash + +# Arguments +config=$1 # YAML configuration file +dataset=$2 # Dataset name +CHUNKS=$3 # Number of chunks + +args=$4 # Additional arguments to pass to prep_data.py + +# Prepare the initial command for chunk 0 +command="scripts/prep_data.py --dataset $dataset --shuffle_mode --chunk 0 --config $config $args" +echo "$command" +python $command + +# Loop through all chunks +for ((i=0; i<$CHUNKS; i++)); do + command="scripts/prep_data.py --dataset $dataset --shuffle_mode --chunk $i --config $config $args" + echo "$command" + python $command +done \ No newline at end of file diff --git a/root_gnn_dgl/jobs/prep_data/run_processing.py b/root_gnn_dgl/jobs/prep_data/run_processing.py new file mode 100644 index 0000000000000000000000000000000000000000..3cf70ffaf38af5bc5484dae3a537ea70be771d67 --- /dev/null +++ b/root_gnn_dgl/jobs/prep_data/run_processing.py @@ -0,0 +1,103 @@ +from concurrent.futures import ThreadPoolExecutor, as_completed +import yaml +import os +import subprocess + +def extract_dataset_keys_and_chunks(yaml_file): + """ + Extract keys and shuffle_chunks under the 'Datasets' section from a YAML file. + """ + try: + # Open and parse the YAML file + with open(yaml_file, 'r') as file: + config = yaml.safe_load(file) + + # Check if 'Datasets' exists in the YAML file + if 'Datasets' in config: + datasets = config['Datasets'] + results = [] + for key, value in datasets.items(): + # Extract shuffle_chunks if available + shuffle_chunks = value.get('shuffle_chunks', None) + results.append((key, shuffle_chunks)) + return results + else: + print(f"No 'Datasets' section found in {yaml_file}.") + return [] + except Exception as e: + print(f"Error reading {yaml_file}: {e}") + return [] + +def call_bash_script(config, dataset_key, shuffle_chunks, bash_script): + """ + Call a bash script with dataset_key and shuffle_chunks as arguments. + """ + try: + # Construct the command to call the bash script + command = f"{bash_script} {config} {dataset_key} {shuffle_chunks}" + + print(f"Executing command: {command}") # Debugging: Print the command being executed + + # Run the command + result = subprocess.run(command, shell=True, check=True) + return f"Success: {dataset_key}" + except subprocess.CalledProcessError as e: + return f"Error executing bash script for dataset {dataset_key}: {e.stderr.decode()}" + +def process_yaml_file(config, base_directory, bash_script): + """ + Process a single YAML file by extracting datasets and calling the bash script in parallel. + """ + yaml_file = os.path.join(base_directory, config) # Construct full path + if os.path.exists(yaml_file): + print(f"Processing file: {config}") + datasets = extract_dataset_keys_and_chunks(yaml_file) + + # Use ThreadPoolExecutor to parallelize bash script calls for datasets + max_workers = min(len(datasets), os.cpu_count()) # Limit workers to number of datasets or CPU cores + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(call_bash_script, config, key, shuffle_chunks, bash_script) + for key, shuffle_chunks in datasets + ] + + # Collect and print results as they complete + results = [] + for future in as_completed(futures): + results.append(future.result()) + return results + else: + return [f"File not found: {yaml_file}"] + +def main(): + # Base directory containing the YAML files + base_directory = os.getcwd() + "/" + + # List of YAML configuration files (relative to the base directory) + configs = [ + # "configs/stats_100K/pretraining_multiclass.yaml", + # "configs/stats_100K/ttH_CP_even_vs_odd.yaml", + # "configs/stats_all/pretraining_multiclass.yaml", + # "configs/stats_all/ttH_CP_even_vs_odd.yaml", + "configs/attention/ttH_CP_even_vs_odd.yaml", + ] + + # Path to the bash script to be called + bash_script = base_directory + "jobs/prep_data/prep_data.sh" + + # Use ThreadPoolExecutor to process YAML files concurrently + max_workers = os.cpu_count() # Use all available CPU cores + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(process_yaml_file, config, base_directory, bash_script) + for config in configs + ] + + # Collect and print results as they complete + for future in as_completed(futures): + results = future.result() + for result in results: + print(result) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/root_gnn_dgl/jobs/salloc.sh b/root_gnn_dgl/jobs/salloc.sh new file mode 100644 index 0000000000000000000000000000000000000000..5c7f378bda998338b844af9fe35493993432518d --- /dev/null +++ b/root_gnn_dgl/jobs/salloc.sh @@ -0,0 +1 @@ +salloc --nodes 4 --qos interactive --time 04:00:00 --constraint gpu --account=trn007 --gres=gpu:4 diff --git a/root_gnn_dgl/jobs/training/multinode/run_multinode_1.sh b/root_gnn_dgl/jobs/training/multinode/run_multinode_1.sh new file mode 100755 index 0000000000000000000000000000000000000000..f15bbf9051b700e77344bffaad4670c856538092 --- /dev/null +++ b/root_gnn_dgl/jobs/training/multinode/run_multinode_1.sh @@ -0,0 +1,51 @@ +#!/bin/bash +#SBATCH -C gpu +#SBATCH -N 4 +#SBATCH --gres=gpu:4 +#SBATCH -q regular +#SBATCH --mail-user=ho22joshua@berkeley.edu +#SBATCH --mail-type=ALL +#SBATCH -t 05:00:00 +#SBATCH -A atlas +#SBATCH -o /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/jobs/slurm/%j.out # STDOUT + + +CONFIG="$*" + +echo "Executing command: $CONFIG" + +cd /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/ + +eval "$(conda shell.bash hook)" +conda init bash +conda activate /global/homes/j/joshuaho/.conda/envs/dgl + +# Run the Python script and capture the output +MASTER_PORT=$(python /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/scripts/find_free_port.py) + +# export MASTER_ADDR=$(hostname) +export MASTER_PORT=$MASTER_PORT + +# Dynamically get the hostname of the first node in the allocation to use as MASTER_ADDR +export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) + +# Debugging: Print the master address and port +echo "Master Address: $MASTER_ADDR" +echo "Master Port: $MASTER_PORT" + +TORCHRUN_ARGUMENTS="--nnodes=$SLURM_NNODES \ + --nproc-per-node=$SLURM_GPUS_ON_NODE \ + --rdzv-id=$SLURM_JOB_ID \ + --rdzv-backend=c10d \ + --rdzv-endpoint=$MASTER_ADDR:$MASTER_PORT" + +GPUS=$(( $SLURM_NNODES * 4 )) + +echo GPUS: $GPUS + +srun --gpus=$GPUS \ + --ntasks-per-node=1 \ + /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/launch_image.sh \ + "--entrypoint /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/jobs/run_multinode_2.sh" \ + $TORCHRUN_ARGUMENTS \ + $CONFIG \ No newline at end of file diff --git a/root_gnn_dgl/jobs/training/multinode/run_multinode_2.sh b/root_gnn_dgl/jobs/training/multinode/run_multinode_2.sh new file mode 100755 index 0000000000000000000000000000000000000000..2b890388629f5c2b7d2b7a4498cddca827986bfd --- /dev/null +++ b/root_gnn_dgl/jobs/training/multinode/run_multinode_2.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +TORCHRUN_ARGUMENTS="$1 $2 $3 $4 $5" +shift 5 +TRAIN_ARGUMENTS="$@" + +# Print the entire argument string +echo "TORCHRUN_ARGUMENTS: $TORCHRUN_ARGUMENTS" +echo "TRAIN_ARGUMENTS: $TRAIN_ARGUMENTS" + +cd /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/ + +DIRECTORY="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/configs/model_configs/" +COMMAND="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/scripts/training_script.py --preshuffle --nocompile --lazy --config $DIRECTORY$TRAIN_ARGUMENTS" + +eval "$(conda shell.bash hook)" +conda init bash +conda activate /opt/conda/envs/dgl + +echo $COMMAND + +torchrun \ + $TORCHRUN_ARGUMENTS \ + $COMMAND \ No newline at end of file diff --git a/root_gnn_dgl/jobs/training/multinode/run_multinode_3.sh b/root_gnn_dgl/jobs/training/multinode/run_multinode_3.sh new file mode 100755 index 0000000000000000000000000000000000000000..1c9c31abb71f95201b8e9c3c9c16d44269968f1e --- /dev/null +++ b/root_gnn_dgl/jobs/training/multinode/run_multinode_3.sh @@ -0,0 +1,24 @@ +#!/bin/bash + +CONFIG=$1 +shift +ARGUEMENTS="$*" + + +DIRECTORY="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/configs/model_configs/" +BASE_COMMAND="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/scripts/training_script.py $ARGUEMENTS --preshuffle --nocompile --lazy --config $DIRECTORY" + +echo "launched image" +cd /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/ + +COMMAND="$BASE_COMMAND$CONFIG" + +eval "$(conda shell.bash hook)" +conda init bash +conda activate /opt/conda/envs/dgl + +echo "Running my script now" +python $COMMAND + + +echo "Done" \ No newline at end of file diff --git a/root_gnn_dgl/jobs/training/multinode/submit.sh b/root_gnn_dgl/jobs/training/multinode/submit.sh new file mode 100644 index 0000000000000000000000000000000000000000..e0a24828805c5e7585e5f90b0c7a1b0e979ca6cf --- /dev/null +++ b/root_gnn_dgl/jobs/training/multinode/submit.sh @@ -0,0 +1,21 @@ +date + +DIRECTORY="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/configs/model_configs/" + +configs=( + 'pretraining_multilabel/multilabel_5_particle_counting.yaml --restart --multinode' + 'pretraining_multilabel/multilabel_17_higgs_kinematics.yaml --restart --multinode' + 'pretraining_multilabel/multilabel_29_top_kinematics.yaml --restart --multinode' + 'pretraining_multilabel/multilabel_41_higgs_tops_all_kinematics.yaml --restart --multinode' +) + +counter=0 + +for job in "${configs[@]}" +do + sbatch --job-name="$job" \ + /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/jobs/run_multinode_1.sh "$job" + ((counter++)) +done + +echo "Total jobs submitted: $counter" \ No newline at end of file diff --git a/root_gnn_dgl/jobs/training/singlegpu/run_job.sh b/root_gnn_dgl/jobs/training/singlegpu/run_job.sh new file mode 100755 index 0000000000000000000000000000000000000000..cf4317b17682f4e6f16b4f406eeb4d19dcc2d541 --- /dev/null +++ b/root_gnn_dgl/jobs/training/singlegpu/run_job.sh @@ -0,0 +1,15 @@ +#!/bin/bash +#SBATCH -N 1 +#SBATCH -C gpu +#SBATCH -q shared +#SBATCH --mail-user=ho22joshua@berkeley.edu +#SBATCH --mail-type=ALL +#SBATCH -t 15:00:00 +#SBATCH -A atlas +#SBATCH -o /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/jobs/slurm/%j.out # STDOUT + +ARGUEMENTS="$*" + +echo "Arguements: $ARGUEMENTS" +echo "launching image" +source launch_image.sh "--entrypoint /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/jobs/run_job_image.sh" $ARGUEMENTS \ No newline at end of file diff --git a/root_gnn_dgl/jobs/training/singlegpu/run_job_image.sh b/root_gnn_dgl/jobs/training/singlegpu/run_job_image.sh new file mode 100755 index 0000000000000000000000000000000000000000..5f213b5e478703af57bcd897815abc45bde63793 --- /dev/null +++ b/root_gnn_dgl/jobs/training/singlegpu/run_job_image.sh @@ -0,0 +1,22 @@ +#!/bin/bash + +CONFIG=$1 +shift +ARGUEMENTS="$*" + +DIRECTORY="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/configs/model_configs/" +BASE_COMMAND="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/scripts/training_script.py $ARGUEMENTS --preshuffle --nocompile --lazy --config $DIRECTORY" + +echo "launched image" +cd /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/ + +COMMAND="$BASE_COMMAND$CONFIG" + +eval "$(conda shell.bash hook)" +conda init bash +conda activate /opt/conda/envs/dgl + +echo "Running my script now" +echo $COMMAND +python -u $COMMAND +echo "Done" \ No newline at end of file diff --git a/root_gnn_dgl/jobs/training/singlegpu/submit.sh b/root_gnn_dgl/jobs/training/singlegpu/submit.sh new file mode 100644 index 0000000000000000000000000000000000000000..474bcd0b91cb0a56e0fa9f14c4ab4c8557334343 --- /dev/null +++ b/root_gnn_dgl/jobs/training/singlegpu/submit.sh @@ -0,0 +1,20 @@ +date + +DIRECTORY="/global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/configs/model_configs/" + +configs=( + "run_3_ttH/v05/sb_yukawa_cp_abs_weights.yaml --abs" +) + +counter=0 + +hours=12 +time="${hours}:00:00" + +for job in "${configs[@]}" +do + sbatch --job-name="$job" --time="$time" /global/cfs/projectdirs/atlas/joshua/root_gnn/root_gnn_dgl/jobs/run_job.sh "$job" + ((counter++)) +done + +echo "Total jobs submitted: $counter" \ No newline at end of file diff --git a/root_gnn_dgl/models/GCN.py b/root_gnn_dgl/models/GCN.py new file mode 100755 index 0000000000000000000000000000000000000000..b4fd1b47bc0d0115c9a41dd52cece8e85dc40fc9 --- /dev/null +++ b/root_gnn_dgl/models/GCN.py @@ -0,0 +1,1945 @@ +import dgl +import dgl.nn as dglnn + +import torch +import torch.nn as nn +import torch.nn.functional as F + +import sys +import os +file_path = os.getcwd() +sys.path.append(file_path) + +import root_gnn_base.dataset as datasets +from root_gnn_base import utils + +import gc + +def Make_SLP(in_size, out_size, activation = nn.ReLU, dropout = 0): + layers = [] + layers.append(nn.Linear(in_size, out_size)) + layers.append(activation()) + layers.append(nn.Dropout(dropout)) + return layers + +def Make_MLP(in_size, hid_size, out_size, n_layers, activation = nn.ReLU, dropout = 0): + layers = [] + if n_layers > 1: + layers += Make_SLP(in_size, hid_size, activation, dropout) + for i in range(n_layers-2): + layers += Make_SLP(hid_size, hid_size, activation, dropout) + layers += Make_SLP(hid_size, out_size, activation, dropout) + else: + layers += Make_SLP(in_size, out_size, activation, dropout) + layers.append(torch.nn.LayerNorm(out_size)) + return nn.Sequential(*layers) + +class MLP(nn.Module): + def __init__(self, in_size, hid_size, out_size, n_layers, activation = nn.ReLU, dropout = 0, **kwargs): + super().__init__() + print(f'Unused args while creating MLP: {kwargs}') + self.layers = Make_MLP(in_size, hid_size, hid_size, n_layers-1, activation, dropout) + self.linear = nn.Linear(hid_size, out_size) + + def forward(self, x): + return self.linear(self.layers(x)) + +def broadcast_global_to_nodes(g, globals): + boundaries = g.batch_num_nodes() + return torch.repeat_interleave(globals, boundaries, dim=0) + +def broadcast_global_to_edges(g, globals): + boundaries = g.batch_num_edges() + return torch.repeat_interleave(globals, boundaries, dim=0) + +def copy_v(edges): + return {'m_v': edges.dst['h']} + +def partial_reset(model : nn.Module): + in_size = len(model.classify.weight[0]) + out_size = len(model.classify.weight) + device = next(model.classify.parameters()).device + torch.manual_seed(2) + model.classify = nn.Linear(in_size, out_size) + model.classify.to(device) + print(model.classify.weight) + +def print_model(model: nn.Module): + print(model) + +def print_mlp(layer): + for l in layer.children(): + if isinstance(l, nn.Linear): + print(l.state_dict()) + else: + print(l) + + +# This function returns a model with the whole GNN completely reset +def full_reset(model : nn.Module): + mlp_list = [model.node_encoder, model.edge_encoder, model.global_encoder, + model.node_update, model.edge_update, model.global_update, + model.global_decoder] + + for mlp in mlp_list: + for layer in mlp.children(): + if hasattr(layer, 'reset_parameters'): + layer.reset_parameters() + partial_reset(model) + + +class GCN(nn.Module): + def __init__(self, in_size, hid_size, out_size, n_layers, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.layers = nn.ModuleList() + + # two-layer GCN + self.layers.extend( + [nn.Linear(in_size, hid_size),] + + [nn.Linear(hid_size, hid_size) for i in range(n_layers)] + + [dglnn.GraphConv(hid_size, hid_size) for i in range(n_layers)] + + [nn.Linear(hid_size, hid_size) for i in range(n_layers)] + ) + self.classify = nn.Linear(hid_size, out_size) + #self.dropout = nn.Dropout(0.05) + + def forward(self, g): + h = g.ndata['features'] + for i, layer in enumerate(self.layers): + if i >= self.n_layers + 1 and i < self.n_layers * 2 + 1: + h = layer(g, h) + else: + h = layer(h) + h = F.relu(h) + with g.local_scope(): + g.ndata['h'] = h + # Calculate graph representation by average readout. + hg = dgl.mean_nodes(g, 'h') + return self.classify(hg) + +class GCN_global(nn.Module): + def __init__(self, in_size, hid_size=4, out_size=1, n_layers=1, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + + #encoder + self.node_encoder = Make_MLP(in_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(1, hid_size, hid_size, n_layers, dropout=dropout) + + #GCN + self.node_update = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.conv = dglnn.GraphConv(hid_size, hid_size) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + def forward(self, g): + h = self.node_encoder(g.ndata['features']) + h_global = self.global_encoder(g.batch_num_nodes()[:, None].to(torch.float)) + for i in range(self.n_layers): + h = self.node_update(h) + h = self.conv(g, h) + g.ndata['h'] = h + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h')), dim = 1)) + h_global = self.global_decoder(h_global) + return self.classify(h_global) + +class GCN_global_2way(nn.Module): + def __init__(self, in_size, hid_size=4, out_size=1, n_layers=1, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + + #encoder + self.node_encoder = Make_MLP(in_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(1, hid_size, hid_size, n_layers, dropout=dropout) + + #GCN + self.node_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.conv = dglnn.GraphConv(hid_size, hid_size) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + def forward(self, g): + h = self.node_encoder(g.ndata['features']) + h_global = self.global_encoder(g.batch_num_nodes()[:, None].to(torch.float)) + for i in range(self.n_layers): + h = self.node_update(torch.cat((h, broadcast_global_to_nodes(g, h_global)), dim = 1)) + h = self.conv(g, h) + g.ndata['h'] = h + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h')), dim = 1)) + h_global = self.global_decoder(h_global) + return self.classify(h_global) + +class Edge_Network(nn.Module): + def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + if (len(sample_global) == 0): + self.has_global = False + else: + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + def forward(self, g, global_feats): + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + + batch_num_nodes = None + sum_weights = None + if "w" in g.ndata: + batch_indices = g.batch_num_nodes() + # Find non-zero rows (non-padded nodes) + non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1) + # Split the mask according to the batch indices + batch_num_nodes = [] + start_idx = 0 + for num_nodes in batch_indices: + end_idx = start_idx + num_nodes + non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item() + batch_num_nodes.append(non_padded_count) + start_idx = end_idx + batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device) + sum_weights = batch_num_nodes[:, None].repeat(1, 64) + global_feats = batch_num_nodes[:, None].to(torch.float) + + h_global = self.global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + if "w" in g.ndata: + mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights + h_global = self.global_update(torch.cat((h_global, mean_nodes, dgl.mean_edges(g, 'e')), dim = 1)) + else: + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.global_decoder(h_global) + return self.classify(h_global) + + def representation(self, g, global_feats): + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + + batch_num_nodes = None + sum_weights = None + if "w" in g.ndata: + batch_indices = g.batch_num_nodes() + # Find non-zero rows (non-padded nodes) + non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1) + # Split the mask according to the batch indices + batch_num_nodes = [] + start_idx = 0 + for num_nodes in batch_indices: + end_idx = start_idx + num_nodes + non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item() + batch_num_nodes.append(non_padded_count) + start_idx = end_idx + batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device) + sum_weights = batch_num_nodes[:, None].repeat(1, 64) + global_feats = batch_num_nodes[:, None].to(torch.float) + + h_global = self.global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + if "w" in g.ndata: + mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights + h_global = self.global_update(torch.cat((h_global, mean_nodes, dgl.mean_edges(g, 'e')), dim = 1)) + else: + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + before_global_decoder = h_global + after_global_decoder = self.global_decoder(before_global_decoder) + after_classify = self.classify(after_global_decoder) + return before_global_decoder, after_global_decoder, after_classify + + def __str__(self): + layer_names = ["node_encoder", "edge_encoder", "global_encoder", + "node_update", "edge_update", "global_update", "global_decoder"] + + layers = [self.node_encoder, self.edge_encoder, self.global_encoder, + self.node_update, self.edge_update, self.global_update, self.global_decoder] + + for i in range(len(layers)): + print(layer_names[i]) + for layer in layers[i].children(): + if isinstance(layer, nn.Linear): + print(layer.state_dict()) + + print("classify") + print(self.classify.weight) + return "" + +class Transferred_Learning(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + + if (len(sample_global) == 0): + self.has_global = False + else: + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + # Freeze Weights + for param in self.pretrained_model.parameters(): + param.requires_grad = False # Freeze all layers + + self.global_decoder = Make_MLP(pretraining_model['args']['hid_size'], hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def forward(self, g, global_feats): + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.TL_global_decoder(h_global) + return self.classify(self.global_decoder(h_global)) + +class Transferred_Learning_Graph(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, additional_proc_steps=1, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + + if (len(sample_global) == 0): + self.has_global = False + else: + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + self.additional_proc_steps = additional_proc_steps + + # Freeze Weights + for param in self.pretrained_model.parameters(): + param.requires_grad = False # Freeze all layers + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def forward(self, g, global_feats): + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + for j in range(self.additional_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + + h_global = self.global_decoder(h_global) + return self.classify(h_global) + +class Transferred_Learning_Parallel(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + # Freeze Weights + for param in self.pretrained_model.parameters(): + param.requires_grad = False # Freeze all layers + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size + pretraining_model['args']['hid_size'], out_size) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def Pretrained_Output(self, g): + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.TL_global_decoder(h_global) + return h_global + + def forward(self, g, global_feats): + pretrained_global = self.Pretrained_Output(g.clone()) + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.global_decoder(h_global) + + return self.classify(torch.cat((pretrained_global, h_global), dim = 1)) + +class Transferred_Learning_Sequential(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + # Freeze Weights + for param in self.pretrained_model.parameters(): + param.requires_grad = False # Freeze all layers + + #encoder + self.mlp = Make_MLP(pretraining_model['args']['hid_size'], hid_size, hid_size, n_layers, dropout=dropout) + + self.classify = nn.Linear(hid_size, out_size) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def Pretrained_Output(self, g): + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.TL_global_decoder(h_global) + return h_global + + def forward(self, g, global_feats): + pretrained_global = self.Pretrained_Output(g.clone()) + global_features = self.mlp(pretrained_global) + return self.classify(global_features) + + +class Transferred_Learning_Message_Passing(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + # Freeze Weights + for param in self.pretrained_model.parameters(): + param.requires_grad = False # Freeze all layers + + #encoder + self.mlp = Make_MLP(pretraining_model['args']['hid_size']*pretraining_model['args']['n_proc_steps'], hid_size, hid_size, n_layers, dropout=dropout) + + self.classify = nn.Linear(hid_size, out_size) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def Pretrained_Output(self, g): + message_passing = None + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + if (message_passing is None): + message_passing = h_global.clone() + else: + message_passing = torch.cat((message_passing, h_global.clone()), dim=1) + h_global = self.TL_global_decoder(h_global) + return message_passing + + def forward(self, g, global_feats): + pretrained_global = self.Pretrained_Output(g.clone()) + #print(f"message_passing layers have size = {pretrained_global.shape}") + #print(pretrained_global) + global_features = self.mlp(pretrained_global) + return self.classify(global_features) + +class Transferred_Learning_Message_Passing_Parallel(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + # Freeze Weights + for param in self.pretrained_model.parameters(): + param.requires_grad = False # Freeze all layers + + self.classify = nn.Linear(pretraining_model['args']['hid_size']*pretraining_model['args']['n_proc_steps'] + hid_size, out_size) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def Pretrained_Output(self, g): + message_passing = None + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + if (message_passing is None): + message_passing = h_global.clone() + else: + message_passing = torch.cat((message_passing, h_global.clone()), dim=1) + h_global = self.TL_global_decoder(h_global) + return message_passing + + def forward(self, g, global_feats): + pretrained_message = self.Pretrained_Output(g.clone()) + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.global_decoder(h_global) + return self.classify(torch.cat((pretrained_message, h_global), dim = 1)) + +class Transferred_Learning_Finetuning(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, frozen_pretraining=False, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + + if (len(sample_global) == 0): + self.has_global = False + else: + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + print(f"Freeze Pretraining = {frozen_pretraining}") + if (frozen_pretraining): + for param in self.pretrained_model.parameters(): + param.requires_grad = False # Freeze all layers + for param in self.pretrained_model[7]: + param.requires_grad = True + + torch.manual_seed(2) + self.classify = nn.Linear(pretraining_model['args']['hid_size'], out_size) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def Pretrained_Output(self, g): + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.TL_global_decoder(h_global) + return h_global + + def forward(self, g, global_feats): + h_global = self.Pretrained_Output(g.clone()) + return self.classify(h_global) + + def representation(self, g, global_feats): + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + + before_global_decoder = h_global + after_global_decoder = self.TL_global_decoder(before_global_decoder) + after_classify = self.classify(after_global_decoder) + return before_global_decoder, after_global_decoder, after_classify + + def __str__(self): + layer_names = ["node_encoder", "edge_encoder", "global_encoder", + "node_update", "edge_update", "global_update", "global_decoder"] + + layers = [self.pretrained_model[1], self.pretrained_model[2], self.pretrained_model[3], + self.pretrained_model[4], self.pretrained_model[5], self.pretrained_model[6], + self.pretrained_model[7]] + + for i in range(len(layers)): + print(layer_names[i]) + for layer in layers[i].children(): + if isinstance(layer, nn.Linear): + print(layer.state_dict()) + + print("classify") + print(self.classify.weight) + return "" + + +class Transferred_Learning_Parallel_Finetuning(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, learning_rate=0.0001, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + + self.learning_rate = learning_rate + + self.parallel_params = [] + self.finetuning_params = [] + + + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + self.finetuning_params.append(self.pretrained_model) + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size + pretraining_model['args']['hid_size'], out_size) + + self.parallel_params.append(self.node_encoder) + self.parallel_params.append(self.edge_encoder) + self.parallel_params.append(self.global_encoder) + self.parallel_params.append(self.node_update) + self.parallel_params.append(self.edge_update) + self.parallel_params.append(self.global_update) + self.parallel_params.append(self.global_decoder) + self.parallel_params.append(self.classify) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x): + for layer in self.pretrained_model[2]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_node_update(self, x): + for layer in self.pretrained_model[4]: + x = layer(x) + return x + + def TL_edge_update(self, x): + for layer in self.pretrained_model[5]: + x = layer(x) + return x + + def TL_global_update(self, x): + for layer in self.pretrained_model[6]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def Pretrained_Output(self, g): + h = self.TL_node_encoder(g.ndata['features']) + e = self.TL_edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.TL_global_decoder(h_global) + return h_global + + def forward(self, g, global_feats): + pretrained_global = self.Pretrained_Output(g.clone()) + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.global_decoder(h_global) + + return self.classify(torch.cat((pretrained_global, h_global), dim = 1)) + + def parameters(self, recurse: bool = True): + params = [] + for model_section in self.parallel_params: + if (type(self.learning_rate) == dict and self.learning_rate["trainable_lr"]): + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["trainable_lr"]}) + else: + params.append({'params': model_section.parameters(), 'lr': 0.0001}) + for model_section in self.finetuning_params: + if (type(self.learning_rate) == dict and self.learning_rate["finetuning_lr"]): + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["finetuning_lr"]}) + else: + params.append({'params': model_section.parameters(), 'lr': 0.0001}) + return params + +class Attention(nn.Module): + def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, num_heads = 1, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + #attention + self.multihead_attn = nn.MultiheadAttention(hid_size, num_heads, dropout=dropout, batch_first=True) + self.queries = nn.Linear(hid_size, hid_size) + self.keys = nn.Linear(hid_size, hid_size) + self.values = nn.Linear(hid_size, hid_size) + + def forward(self, g, global_feats): + h = self.node_encoder(g.ndata['features']) + g.ndata['h'] = h + + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + + batch_num_nodes = None + sum_weights = None + if "w" in g.ndata: + batch_indices = g.batch_num_nodes() + # Find non-zero rows (non-padded nodes) + non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1) + # Split the mask according to the batch indices + batch_num_nodes = [] + start_idx = 0 + for num_nodes in batch_indices: + end_idx = start_idx + num_nodes + non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item() + batch_num_nodes.append(non_padded_count) + start_idx = end_idx + batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device) + sum_weights = batch_num_nodes[:, None].repeat(1, 64) + global_feats = batch_num_nodes[:, None].to(torch.float) + + h_global = self.global_encoder(global_feats) + + h_original_shape = h.shape + num_graphs = len(dgl.unbatch(g)) + num_nodes = g.batch_num_nodes()[0].item() + padding_mask = g.ndata['padding_mask'] > 0 + padding_mask = torch.reshape(padding_mask, (num_graphs, num_nodes)) + + h = g.ndata['h'] + query = self.queries(h) + key = self.keys(h) + value = self.values(h) + query = torch.reshape(query, (num_graphs, num_nodes, h_original_shape[1])) + key = torch.reshape(key, (num_graphs, num_nodes, h_original_shape[1])) + value = torch.reshape(value, (num_graphs, num_nodes, h_original_shape[1])) + h, _ = self.multihead_attn(query, key, value, key_padding_mask=padding_mask) + h = torch.reshape(h, h_original_shape) + + h = self.node_update(torch.cat((h, broadcast_global_to_nodes(g, h_global)), dim = 1)) + g.ndata['h'] = h + mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights + h_global = self.global_update(torch.cat((h_global, mean_nodes), dim = 1)) + h_global = self.global_decoder(h_global) + return self.classify(h_global) + +class Attention_Edge_Network(nn.Module): + def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, num_heads = 1, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + + #attention + self.multihead_attn = nn.MultiheadAttention(hid_size, num_heads, dropout=dropout, batch_first=True) + self.queries = nn.Linear(hid_size, hid_size) + self.keys = nn.Linear(hid_size, hid_size) + self.values = nn.Linear(hid_size, hid_size) + + def forward(self, g, global_feats): + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.global_encoder(global_feats) + + h = g.ndata['h'] + h_original_shape = h.shape + num_graphs = len(dgl.unbatch(g)) + num_nodes = g.batch_num_nodes()[0].item() + padding_mask = g.ndata['padding_mask'].bool() + + padding_mask = torch.reshape(padding_mask, (num_graphs, num_nodes)) + + for i in range(self.n_proc_steps): + + h = g.ndata['h'] + query = self.queries(h) + key = self.keys(h) + value = self.values(h) + query = torch.reshape(query, (num_graphs, num_nodes, h_original_shape[1])) + key = torch.reshape(key, (num_graphs, num_nodes, h_original_shape[1])) + value = torch.reshape(value, (num_graphs, num_nodes, h_original_shape[1])) + h, _ = self.multihead_attn(query, key, value, key_padding_mask=padding_mask) + h = torch.reshape(h, h_original_shape) + + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h', 'w'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.global_decoder(h_global) + return self.classify(h_global) + +class Attention_Unbatched(nn.Module): + def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, num_heads = 1, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + + #attention + self.multihead_attn = nn.MultiheadAttention(hid_size, 1, dropout=dropout) + self.queries = nn.Linear(hid_size, hid_size) + self.keys = nn.Linear(hid_size, hid_size) + self.values = nn.Linear(hid_size, hid_size) + + + + def forward(self, g, global_feats): + + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.global_encoder(global_feats) + + for i in range(self.n_proc_steps): + + unbatched_g = dgl.unbatch(g) + for graph in unbatched_g: + h = graph.ndata['h'] + h, _ = self.multihead_attn(self.queries(h), self.keys(h), self.values(h)) + graph.ndata['h'] = h + g = dgl.batch(unbatched_g) + + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.global_decoder(h_global) + return self.classify(h_global) + +class Transferred_Learning_Attention(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, num_heads, dropout=0, learning_rate=0.0001, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.learning_rate = learning_rate + + self.pretraining_params = [] + self.attention_params = [] + + self.pretrained_model = utils.buildFromConfig(pretraining_model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + + checkpoint = torch.load(pretraining_path) + self.pretrained_model.load_state_dict(checkpoint['model_state_dict']) + pretrained_layers = list(self.pretrained_model.children()) + pretrained_layers = pretrained_layers[:-1] + self.pretrained_model = nn.Sequential(*pretrained_layers) + + self.pretraining_params.append(self.pretrained_model[1]) + self.pretraining_params.append(self.pretrained_model[3]) + self.pretraining_params.append(self.pretrained_model[7]) + + #attention + self.multihead_attn = nn.MultiheadAttention(hid_size, num_heads, dropout=dropout, batch_first=True) + self.queries = nn.Linear(hid_size, hid_size) + self.keys = nn.Linear(hid_size, hid_size) + self.values = nn.Linear(hid_size, hid_size) + + self.node_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(2*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + self.classify = nn.Linear(pretraining_model['args']['hid_size'], out_size) + + self.attention_params.append(self.multihead_attn) + + self.attention_params.append(self.queries) + self.attention_params.append(self.keys) + self.attention_params.append(self.values) + self.attention_params.append(self.classify) + self.attention_params.append(self.node_update) + self.attention_params.append(self.global_update) + + def TL_node_encoder(self, x): + for layer in self.pretrained_model[1]: + x = layer(x) + return x + + def TL_global_encoder(self, x): + for layer in self.pretrained_model[3]: + x = layer(x) + return x + + def TL_global_decoder(self, x): + for layer in self.pretrained_model[7]: + x = layer(x) + return x + + def forward(self, g, global_feats): + h = self.TL_node_encoder(g.ndata['features']) + g.ndata['h'] = h + + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + + batch_num_nodes = None + sum_weights = None + if "w" in g.ndata: + batch_indices = g.batch_num_nodes() + # Find non-zero rows (non-padded nodes) + non_padded_nodes_mask = torch.any(g.ndata['features'] != 0, dim=1) + # Split the mask according to the batch indices + batch_num_nodes = [] + start_idx = 0 + for num_nodes in batch_indices: + end_idx = start_idx + num_nodes + non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item() + batch_num_nodes.append(non_padded_count) + start_idx = end_idx + batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata['features'].device) + sum_weights = batch_num_nodes[:, None].repeat(1, 64) + global_feats = batch_num_nodes[:, None].to(torch.float) + + h_global = self.TL_global_encoder(global_feats) + + h_original_shape = h.shape + num_graphs = len(dgl.unbatch(g)) + num_nodes = g.batch_num_nodes()[0].item() + padding_mask = g.ndata['padding_mask'] > 0 + padding_mask = torch.reshape(padding_mask, (num_graphs, num_nodes)) + + h = g.ndata['h'] + query = self.queries(h) + key = self.keys(h) + value = self.values(h) + query = torch.reshape(query, (num_graphs, num_nodes, h_original_shape[1])) + key = torch.reshape(key, (num_graphs, num_nodes, h_original_shape[1])) + value = torch.reshape(value, (num_graphs, num_nodes, h_original_shape[1])) + h, _ = self.multihead_attn(query, key, value, key_padding_mask=padding_mask) + h = torch.reshape(h, h_original_shape) + + h = self.node_update(torch.cat((h, broadcast_global_to_nodes(g, h_global)), dim = 1)) + g.ndata['h'] = h + mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights + h_global = self.global_update(torch.cat((h_global, mean_nodes), dim = 1)) + h_global = self.TL_global_decoder(h_global) + return self.classify(h_global) + + def parameters(self, recurse: bool = True): + params = [] + for model_section in self.pretraining_params: + if (type(self.learning_rate) == dict and self.learning_rate["pretraining_lr"]): + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["pretraining_lr"]}) + else: + params.append({'params': model_section.parameters(), 'lr': 0.0001}) + for model_section in self.attention_params: + if (type(self.learning_rate) == dict and self.learning_rate["attention_lr"]): + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["attention_lr"]}) + else: + params.append({'params': model_section.parameters(), 'lr': 0.0001}) + return params + +class Multimodel_Transferred_Learning(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, frozen_pretraining=True, learning_rate=None, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.learning_rate = learning_rate + input_size = 0 + + self.pretraining_params = [] + self.model_params = [] + + self.pretrained_models = [] + for model, path in zip(pretraining_model, pretraining_path): + input_size += model['args']['hid_size'] + model = utils.buildFromConfig(model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + + checkpoint = torch.load(path)['model_state_dict'] + new_state_dict = {} + for k, v in checkpoint.items(): + new_key = k.replace('module.', '') + new_state_dict[new_key] = v + model.load_state_dict(new_state_dict) + pretrained_layers = list(model.children()) + pretrained_layers = pretrained_layers[:-1] + + model = nn.Sequential(*pretrained_layers) + + # Freeze Weights + print(f"Freeze Pretraining = {frozen_pretraining}") + if (frozen_pretraining): + for param in model.parameters(): + param.requires_grad = False # Freeze all layers + self.pretraining_params.append(model) + self.pretrained_models.append(model) + + print(f"len(pretrained_models) = {len(self.pretrained_models)}") + print(f"input size = {input_size}") + + self.final_mlp = Make_MLP(input_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + self.model_params.append(self.final_mlp) + self.model_params.append(self.classify) + + def TL_node_encoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][1]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][2]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][2]: + x = layer(x) + return x + + def TL_global_encoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][3]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][3]: + x = layer(x) + return x + + def TL_node_update(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][4]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][4]: + x = layer(x) + return x + + def TL_edge_update(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][5]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][5]: + x = layer(x) + return x + + def TL_global_update(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][6]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][6]: + x = layer(x) + return x + + def TL_global_decoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][7]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][7]: + x = layer(x) + return x + + def Pretrained_Output(self, g, model_idx): + h = self.TL_node_encoder(g.ndata['features'], model_idx) + e = self.TL_edge_encoder(g.edata['features'], model_idx) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats, model_idx) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1), model_idx) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1), model_idx) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1), model_idx) + # h_global = self.TL_global_decoder(h_global, model_idx) + return h_global + + def forward(self, g, global_feats): + h_global = [] + for i in range(len(self.pretrained_models)): + h_global.append(self.Pretrained_Output(g.clone(), i)) + h_global = torch.concatenate(h_global, dim=1) + return self.classify(self.final_mlp(h_global)) + + def to(self, device): + for i in range(len(self.pretrained_models)): + self.pretrained_models[i].to(device) + self.classify.to(device) + self.final_mlp.to(device) + return self + + def parameters(self, recurse: bool = True): + params = [] + for model_section in self.pretraining_params: + if (type(self.learning_rate) == dict and self.learning_rate["pretraining_lr"]): + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["pretraining_lr"]}) + else: + params.append({'params': model_section.parameters(), 'lr': 0.00001}) + for model_section in self.model_params: + if (type(self.learning_rate) == dict and self.learning_rate["model_lr"]): + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["model_lr"]}) + else: + params.append({'params': model_section.parameters(), 'lr': 0.0001}) + return params + + +class MultiModel(nn.Module): + def __init__(self, pretraining_path, pretraining_model, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, frozen_pretraining=True, learning_rate=None, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + self.learning_rate = learning_rate + input_size = 0 + + self.model_params = [] + self.pretraining_params = [] + + self.pretrained_models = [] + for model, path in zip(pretraining_model, pretraining_path): + input_size += model['args']['hid_size'] + model = utils.buildFromConfig(model, {'sample_graph': sample_graph, 'sample_global': sample_global}) + + checkpoint = torch.load(path)['model_state_dict'] + new_state_dict = {} + for k, v in checkpoint.items(): + new_key = k.replace('module.', '') + new_state_dict[new_key] = v + model.load_state_dict(new_state_dict) + pretrained_layers = list(model.children()) + pretrained_layers = pretrained_layers[:-1] + + model = nn.Sequential(*pretrained_layers) + + # Freeze Weights + print(f"Freeze Pretraining = {frozen_pretraining}") + if (frozen_pretraining): + for param in model.parameters(): + param.requires_grad = False # Freeze all layers + self.pretraining_params.append(model) + self.pretrained_models.append(model) + + print(f"len(pretrained_models) = {len(self.pretrained_models)}") + print(f"input size = {input_size}") + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + self.final_mlp = Make_MLP(input_size + hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.classify = nn.Linear(hid_size, out_size) + + self.model_params.append(self.final_mlp) + self.model_params.append(self.classify) + + def TL_node_encoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][1]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][1]: + x = layer(x) + return x + + def TL_edge_encoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][2]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][2]: + x = layer(x) + return x + + def TL_global_encoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][3]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][3]: + x = layer(x) + return x + + def TL_node_update(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][4]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][4]: + x = layer(x) + return x + + def TL_edge_update(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][5]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][5]: + x = layer(x) + return x + + def TL_global_update(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][6]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][6]: + x = layer(x) + return x + + def TL_global_decoder(self, x, model_idx): + try: + for layer in self.pretrained_models[model_idx][7]: + x = layer(x) + return x + except (NotImplementedError, IndexError): + for layer in self.pretrained_models[model_idx][1][7]: + x = layer(x) + return x + + def Pretrained_Output(self, g, model_idx): + h = self.TL_node_encoder(g.ndata['features'], model_idx) + e = self.TL_edge_encoder(g.edata['features'], model_idx) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.TL_global_encoder(global_feats, model_idx) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.TL_edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1), model_idx) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.TL_node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1), model_idx) + h_global = self.TL_global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1), model_idx) + # h_global = self.TL_global_decoder(h_global, model_idx) + return h_global + + def forward(self, g, global_feats): + h = self.node_encoder(g.ndata['features']) + e = self.edge_encoder(g.edata['features']) + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + h_global = self.global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = [h_global] + for i in range(len(self.pretrained_models)): + h_global.append(self.Pretrained_Output(g.clone(), i)) + h_global = torch.concatenate(h_global, dim=1) + return self.classify(self.final_mlp(h_global)) + + def to(self, device): + for i in range(len(self.pretrained_models)): + self.pretrained_models[i].to(device) + self.classify.to(device) + self.final_mlp.to(device) + self.node_encoder.to(device) + self.edge_encoder.to(device) + self.global_encoder.to(device) + + self.node_update.to(device) + self.edge_update.to(device) + self.global_update.to(device) + return self + + def parameters(self, recurse: bool = True): + params = [] + for i, model_section in enumerate(self.pretraining_params): + if (type(self.learning_rate) == dict and self.learning_rate["pretraining_lr"]): + print(f"Pretraining LR = {self.learning_rate['pretraining_lr'][i]}") + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["pretraining_lr"][i]}) + else: + print(f"Pretraining LR = 0.00001") + params.append({'params': model_section.parameters(), 'lr': 0.00001}) + for model_section in self.model_params: + if (type(self.learning_rate) == dict and self.learning_rate["model_lr"]): + print(f"Model LR = {self.learning_rate['model_lr']}") + params.append({'params': model_section.parameters(), 'lr': self.learning_rate["model_lr"]}) + else: + print(f"Model LR = 0.0001") + params.append({'params': model_section.parameters(), 'lr': 0.0001}) + return params + + +class Clustering(nn.Module): + def __init__(self, sample_graph, sample_global, hid_size, out_size, n_layers, n_proc_steps, dropout=0, **kwargs): + super().__init__() + print(f'Unused args while creating GCN: {kwargs}') + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.layers = nn.ModuleList() + if (len(sample_global) == 0): + self.has_global = False + else: + self.has_global = sample_global.shape[1] != 0 + gl_size = sample_global.shape[1] if self.has_global else 1 + + #encoder + self.node_encoder = Make_MLP(sample_graph.ndata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.edge_encoder = Make_MLP(sample_graph.edata['features'].shape[1], hid_size, hid_size, n_layers, dropout=dropout) + self.global_encoder = Make_MLP(gl_size, hid_size, hid_size, n_layers, dropout=dropout) + + #GNN + self.node_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.edge_update = Make_MLP(4*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + self.global_update = Make_MLP(3*hid_size, hid_size, hid_size, n_layers, dropout=dropout) + + #decoder + self.global_decoder = Make_MLP(hid_size, hid_size, out_size, n_layers, dropout=dropout) + + def model_forward(self, g, global_feats, features = 'features'): + h = self.node_encoder(g.ndata[features]) + e = self.edge_encoder(g.edata[features]) + + g.ndata['h'] = h + g.edata['e'] = e + if not self.has_global: + global_feats = g.batch_num_nodes()[:, None].to(torch.float) + + batch_num_nodes = None + sum_weights = None + if "w" in g.ndata: + batch_indices = g.batch_num_nodes() + # Find non-zero rows (non-padded nodes) + non_padded_nodes_mask = torch.any(g.ndata[features] != 0, dim=1) + # Split the mask according to the batch indices + batch_num_nodes = [] + start_idx = 0 + for num_nodes in batch_indices: + end_idx = start_idx + num_nodes + non_padded_count = non_padded_nodes_mask[start_idx:end_idx].sum().item() + batch_num_nodes.append(non_padded_count) + start_idx = end_idx + batch_num_nodes = torch.tensor(batch_num_nodes, device = g.ndata[features].device) + sum_weights = batch_num_nodes[:, None].repeat(1, 64) + global_feats = batch_num_nodes[:, None].to(torch.float) + + h_global = self.global_encoder(global_feats) + for i in range(self.n_proc_steps): + g.apply_edges(dgl.function.copy_u('h', 'm_u')) + g.apply_edges(copy_v) + g.edata['e'] = self.edge_update(torch.cat((g.edata['e'], g.edata['m_u'], g.edata['m_v'], broadcast_global_to_edges(g, h_global)), dim = 1)) + g.update_all(dgl.function.copy_e('e', 'm'), dgl.function.sum('m', 'h_e')) + g.ndata['h'] = self.node_update(torch.cat((g.ndata['h'], g.ndata['h_e'], broadcast_global_to_nodes(g, h_global)), dim = 1)) + if "w" in g.ndata: + mean_nodes = dgl.sum_nodes(g, 'h', 'w') / sum_weights + h_global = self.global_update(torch.cat((h_global, mean_nodes, dgl.mean_edges(g, 'e')), dim = 1)) + else: + h_global = self.global_update(torch.cat((h_global, dgl.mean_nodes(g, 'h'), dgl.mean_edges(g, 'e')), dim = 1)) + h_global = self.global_decoder(h_global) + return h_global + + def forward(self, g, global_feats): + h_global = self.model_forward(g, global_feats, 'features') + h_global_augmented = self.model_forward(g, global_feats, 'augmented_features') + return torch.cat((h_global, h_global_augmented), dim=1) + + def representation(self, g, global_feats): + h_global = self.model_forward(g, global_feats, 'features') + h_global_augmented = self.model_forward(g, global_feats, 'augmented_features') + return h_global, h_global_augmented, torch.cat((h_global, h_global_augmented), dim=1) + + def __str__(self): + layer_names = ["node_encoder", "edge_encoder", "global_encoder", + "node_update", "edge_update", "global_update", "global_decoder"] + + layers = [self.node_encoder, self.edge_encoder, self.global_encoder, + self.node_update, self.edge_update, self.global_update, self.global_decoder] + + for i in range(len(layers)): + print(layer_names[i]) + for layer in layers[i].children(): + if isinstance(layer, nn.Linear): + print(layer.state_dict()) + + print("classify") + print(self.classify.weight) + return "" \ No newline at end of file diff --git a/root_gnn_dgl/models/loss.py b/root_gnn_dgl/models/loss.py new file mode 100755 index 0000000000000000000000000000000000000000..c846861ca33f1020a0e8099a7e432fd0ab3528df --- /dev/null +++ b/root_gnn_dgl/models/loss.py @@ -0,0 +1,311 @@ +from torch import nn +import torch +from root_gnn_base import utils +import numpy as np + +class MaskedLoss(): + def __init__(self, mask = []): + self.mask = mask + + def make_mask(self, targets): + mask = torch.ones_like(targets[:,0]) + for m in self.mask: + if m['op'] == 'eq': + mask[targets[:,m['idx']] == m['val']] = 0 + elif m['op'] == 'gt': + mask[targets[:,m['idx']] > m['val']] = 0 + elif m['op'] == 'lt': + mask[targets[:,m['idx']] < m['val']] = 0 + elif m['op'] == 'ge': + mask[targets[:,m['idx']] >= m['val']] = 0 + elif m['op'] == 'le': + mask[targets[:,m['idx']] <= m['val']] = 0 + elif m['op'] == 'ne': + mask[targets[:,m['idx']] != m['val']] = 0 + else: + raise ValueError(f'Unknown mask op {m["op"]}') + return mask == 1 + +class MaskedL1Loss(MaskedLoss): + def __init__(self, mask = [], index = 0): + super().__init__(mask) + self.index = index + self.loss = nn.L1Loss() + + def __call__(self, logits, targets): + mask = self.make_mask(targets) + return self.loss(logits[mask], targets[mask][:,self.index]) + +class BCEWithLogitsLoss(): + def __init__(self, weight=None, reduction='mean'): + self.loss = nn.BCEWithLogitsLoss(weight=weight, reduction=reduction) + + def __call__(self, logits, targets): + return self.loss(logits[:,0], targets.float()) + +class MultiScore(): + def __init__(self, scores): + self. score_fcns = [] + self.start_idx = [] + self.end_idx = [] + for score in scores: + self.score_fcns.append(utils.buildFromConfig(score)) + self.start_idx.append(score['start_idx']) + self.end_idx.append(score['end_idx']) + + def __call__(self, last_layer): + scores = [] + for i in range(len(self.score_fcns)): + scores.append(self.score_fcns[i](last_layer[:, self.start_idx[i]:self.end_idx[i]])) + return torch.cat(scores, dim=1) + +class MultiLoss(): + def __init__(self, losses): + self.loss_fcns = [] + self.label_start_idx = [] + self.label_end_idx = [] + self.output_start_idx = [] + self.output_end_idx = [] + self.weights = [] + self.label_types = [] + for loss in losses: + self.loss_fcns.append(utils.buildFromConfig(loss)) + self.label_start_idx.append(loss['label_start_idx']) + self.label_end_idx.append(loss['label_end_idx']) + self.output_start_idx.append(loss['output_start_idx']) + self.output_end_idx.append(loss['output_end_idx']) + self.weights.append(loss.get('weight', 1.0)) + self.label_types.append(loss.get('label_type', 'float')) + + def __call__(self, logits, targets): + loss = 0 + # print(logits.shape, targets.shape) + for i in range(len(self.loss_fcns)): + if self.label_types[i] == 'int': + # print('loss', i, self.label_start_idx[i], self.label_end_idx[i], self.output_start_idx[i], self.output_end_idx[i]) + # print(logits[:, self.output_start_idx[i]:self.output_end_idx[i]].shape, targets[:, self.label_start_idx[i]].shape) + loss += self.weights[i] * self.loss_fcns[i](logits[:, self.output_start_idx[i]:self.output_end_idx[i]], targets[:, self.label_start_idx[i]].to(int)) + elif self.label_end_idx[i] - self.label_start_idx[i] == 1: + loss += self.weights[i] * self.loss_fcns[i](logits[:, self.output_start_idx[i]:self.output_end_idx[i]], targets[:, self.label_start_idx[i]]) + else: + # print('loos', i, self.label_start_idx[i], self.label_end_idx[i], self.output_start_idx[i], self.output_end_idx[i]) + # print(logits[:, self.output_start_idx[i]:self.output_end_idx[i]].shape, targets[:, self.label_start_idx[i]:self.label_end_idx[i]].shape) + loss += self.weights[i] * self.loss_fcns[i](logits[:, self.output_start_idx[i]:self.output_end_idx[i]], targets[:, self.label_start_idx[i]:self.label_end_idx[i]]) + return loss + +class AdvLoss(): + def __init__(self, loss, adv_loss, adv_weight=1.0): + self.loss_fcn = utils.buildFromConfig(loss) + self.adv_loss_fcn = utils.buildFromConfig(adv_loss) + self.adv_weight = adv_weight + + def __call__(self, logits, targets): + mask = targets[:,0] == 0 + loss = self.loss_fcn(logits[:,0], targets[:,0]) + adv_loss = self.adv_loss_fcn(logits[mask][:,1], targets[mask]) + return loss - self.adv_weight * adv_loss + +class MassWindowAdvLoss(AdvLoss): + def __call__(self, logits, targets): + mask = (targets[:,0] == 0) & (targets[:,1] > 5) & (targets[:,1] < 25) + print(mask, mask.shape, mask.sum()) + loss = self.loss_fcn(logits[:,0], targets[:,0]) + print(loss) + adv_loss = self.adv_loss_fcn(logits[mask][:,1], targets[mask][:,1]) + print(adv_loss) + return loss - self.adv_weight * adv_loss + +class KDELoss(MaskedLoss): + def __init__(self, mask = [], index = 0): + self.index = index + super().__init__(mask) + + def __call__(self, logits, targets): + mask = self.make_mask(targets) + logits = logits[mask] + targets = targets[mask][:,self.index] + N = logits.shape[0] + masses = targets / torch.sqrt(torch.mean(targets**2)) + scores = logits[:,0] / torch.sqrt(torch.mean(logits**2)) + + factor_2d = (1.0*N) ** (-2/6) + covs = (factor_2d * torch.var(masses), factor_2d * torch.var(scores)) + + m_diffs = torch.unsqueeze(masses, 1) - torch.unsqueeze(masses, 0) + s_diffs = torch.unsqueeze(scores, 1) - torch.unsqueeze(scores, 0) + + ymm = torch.exp(- (m_diffs**2) / (4 * covs[0])) + yss = torch.exp(- (s_diffs**2) / (4 * covs[1])) + + integral_rho_2d_rho_2d = torch.einsum('ij,ij->', ymm, yss) + integral_rho_1d_rho_1d = torch.einsum('ij,kl->', ymm, yss) + integral_rho_2d_rho_1d = torch.einsum('ij,ik->', ymm, yss) + raw_integral = integral_rho_2d_rho_2d - 2 * integral_rho_2d_rho_1d / N + integral_rho_1d_rho_1d / N**2 + return raw_integral / (4 * torch.pi * N**2) + +class MultiLabelLoss(): + def __init__(self, label_names, label_types, label_weights = None): + self.loss_fcn = [] + if (label_weights): + self.weights = torch.tensor(label_weights) + else: + self.weights = torch.ones(len(label_types)) + for type in label_types: + if (type == "r"): + self.loss_fcn.append(torch.nn.MSELoss(reduce=False)) + elif (type == "c"): + self.loss_fcn.append(torch.nn.BCEWithLogitsLoss()) + print(f"self.weights = {self.weights}") + + def __call__(self, logits, targets): + targets = targets.float() + loss = torch.zeros(len(logits[:, 0]), device = logits.get_device()) + for i in range(len(self.loss_fcn)): + loss += self.weights[i] * self.loss_fcn[i](logits[:, i], targets[:, i]) + return torch.mean(loss) + + +class MultiLabelFinish(): + def __init__(self, label_names, label_types): + self.finish_fcn = [] + for type in label_types: + if (type == "r"): + self.finish_fcn.append(None) + elif (type == "c"): + self.finish_fcn.append(torch.special.expit) + + def __call__(self, logits): + for i in range(len(self.finish_fcn)): + if (self.finish_fcn[i]): + logits[:, i] = self.finish_fcn[i](logits[:, i].to(torch.long)) + return logits + +class ContrastiveClusterLoss(): + def __init__(self, k=10, temperature=1, alpha=1): + self.k = k + self.temperature = temperature + self.alpha = alpha + + def __call__(self, logits, targets): + targets = targets.float() + logits_combined = logits.float() + + hid_size = int(len(logits[0]) / 2) + + logits = normalize_embeddings(logits_combined[:, :hid_size]) + logits_augmented = normalize_embeddings(logits_combined[:, hid_size:]) + + contrastive = contrastive_loss(logits, logits_augmented, self.temperature) + clustering, _ = clustering_loss(logits, self.k) + + variance_loss = variance_regularization(logits) + variance_regularization(logits_augmented) + + return torch.mean(contrastive + clustering + self.alpha * variance_loss) + +class ContrastiveClusterFinish(): + def __init__(self, k = 10, temperature = 1, max_cluster_iterations = 10): + self.k = k + self.temperature = temperature + self.max_cluster_iterations = max_cluster_iterations + + print(f"ContrastiveClusterFinish: k = {k}, temperature = {temperature}") + + def __call__(self, logits): + logits_combined = logits.float() + + hid_size = int(len(logits[0]) / 2) + + logits = logits_combined[:, :hid_size] + logits_augmented = logits_combined[:, hid_size:] + + contrastive = contrastive_loss(logits, logits_augmented, self.temperature) + clustering, _ = clustering_loss(logits, self.k, self.max_cluster_iterations) + variance = variance_regularization(logits) + variance_regularization(logits_augmented) + + return contrastive, clustering, variance + +def s(z_i, z_j): + z_i = torch.tensor(z_i) if not isinstance(z_i, torch.Tensor) else z_i + z_j = torch.tensor(z_j) if not isinstance(z_j, torch.Tensor) else z_j + + return torch.cdist(z_i, z_j, p=2) + # dot_product = torch.dot(z_i, z_j) + # norm_i = torch.linalg.norm(z_i) + # norm_j = torch.linalg.norm(z_j) + + # return dot_product / (norm_i * norm_j) + +def contrastive_loss(logits, logits_augmented, temperature=1, margin=1.0): + logits = torch.tensor(logits) if not isinstance(logits, torch.Tensor) else logits + logits_augmented = torch.tensor(logits_augmented) if not isinstance(logits_augmented, torch.Tensor) else logits_augmented + + z = torch.cat((logits, logits_augmented), dim=0) + similarity_matrix = torch.mm(z, z.t()) / temperature + norms = torch.linalg.norm(z, dim=1) + norm_matrix = torch.ger(norms, norms) + similarity_matrix = similarity_matrix / norm_matrix + mask = torch.eye(similarity_matrix.size(0), dtype=torch.bool) + + loss = 0 + for k in range(len(logits)): + numerator = torch.exp(similarity_matrix[k, k + len(logits)]) + denominator = torch.sum(torch.exp(similarity_matrix[k, ~mask[k]])) + + loss += -torch.log(numerator / denominator) + + return loss + + +def clustering_loss(logits, k=10, max_iterations=10): + # Step 1: Initialize cluster means + indices = torch.randperm(logits.size(0))[:k] + cluster_means = logits[indices] + + prev_assignments = None + assignment_history = [] + iteration = 0 + + while iteration < max_iterations: + iteration += 1 + + # Step 2: Assign each data point to the nearest cluster mean + distances = torch.cdist(logits, cluster_means, p=2) # Compute distances between logits and cluster means + cluster_assignments = torch.argmin(distances, dim=1) # Assign each point to the nearest cluster mean + + # Check for convergence: if assignments do not change, break the loop + if prev_assignments is not None and torch.equal(cluster_assignments, prev_assignments): + break + + # Check for cycles: if assignments have been seen before, break the loop + if any(torch.equal(cluster_assignments, prev) for prev in assignment_history): + break + + assignment_history.append(cluster_assignments.clone()) + prev_assignments = cluster_assignments.clone() + + # Step 3: Update cluster means based on assignments + new_cluster_means = torch.zeros_like(cluster_means) + for i in range(k): + assigned_points = logits[cluster_assignments == i] + if assigned_points.size(0) > 0: + new_cluster_means[i] = assigned_points.mean(dim=0) + else: + # If no points are assigned to the cluster, reinitialize the mean randomly + new_cluster_means[i] = logits[torch.randint(0, logits.size(0), (1,)).item()] + cluster_means = new_cluster_means + + # Step 4: Compute the clustering loss + distances = torch.cdist(logits, cluster_means, p=2) + min_distances = torch.min(distances, dim=1)[0] + loss = torch.sum(min_distances ** 2) + + return loss, cluster_means + +def normalize_embeddings(embeddings): + return embeddings / embeddings.norm(dim=1, keepdim=True) + +def variance_regularization(embeddings): + mean_embedding = embeddings.mean(dim=0) + variance = ((embeddings - mean_embedding) ** 2).mean() + return variance + diff --git a/root_gnn_dgl/root_gnn_base/batched_dataset.py b/root_gnn_dgl/root_gnn_base/batched_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..101ada4a0a7ae4c2fba01068579d190332fbe444 --- /dev/null +++ b/root_gnn_dgl/root_gnn_base/batched_dataset.py @@ -0,0 +1,190 @@ +from dgl.dataloading import GraphDataLoader +from torch.utils.data.sampler import SubsetRandomSampler +from torch.utils.data.sampler import SequentialSampler +from dgl.data import DGLDataset +import torch +import time +import os +import dgl +from root_gnn_base import utils + +def GetBatchedLoader(dataset, batch_size, mask_fn = None, drop_last=True, **kwargs): + if mask_fn == None: + mask_fn = lambda x: torch.ones(len(x), dtype=torch.bool) + dloader = GraphDataLoader(dataset, sampler=SubsetRandomSampler(torch.arange(len(dataset))[mask_fn(dataset)]), batch_size=batch_size, drop_last=drop_last, num_workers = 0) + return dloader + +#Dataset which contains prebatched shuffled graphs. Cannot be saved to disk, else batching info is lost. +class PreBatchedDataset(DGLDataset): + def __init__(self, start_dataset, batch_size, mask_fn = None, drop_last=True, save_to_disk = True, suffix = '', chunks = 1, chunkno = -1, shuffle = True, padding_mode = 'NONE', **kwargs): + print(f'Unused kwargs: {kwargs}') + self.start_dataset = start_dataset + self.start_dataset.load() + + self.batch_size = batch_size + self.chunks = chunks + self.chunkno = chunkno + self.mask_fn = mask_fn + self.drop_last = drop_last + self.graphs = [] + self.label = [] + self.padding_mode = padding_mode + self.save_to_disk = save_to_disk + self.shuffle = shuffle + self.suffix = suffix + self.current_chunk = None + self.current_chunk_idx = -1 + super().__init__(name = start_dataset.name + '_prebatched_padded', save_dir=start_dataset.save_dir) + + def process(self): + first = 0 + last = len(self.start_dataset) + if self.chunks > 1 and self.chunkno >= 0: + first = int(self.chunkno / self.chunks * len(self.start_dataset)) + last = int((self.chunkno + 1) / self.chunks * len(self.start_dataset)) + print(f'Processing chunk {self.chunkno} of {self.chunks} from {first} to {last} of {len(self.start_dataset)}') + mask = torch.logical_and(torch.logical_and(self.mask_fn(self.start_dataset), torch.arange(len(self.start_dataset)) >= first), torch.arange(len(self.start_dataset)) < last) + if self.shuffle: + dloader = GraphDataLoader(self.start_dataset, sampler=SubsetRandomSampler(torch.arange(len(self.start_dataset))[mask]), batch_size=self.batch_size, drop_last=self.drop_last) + else: #Only don't shuffle if we're doing inference. Then we want all of the events anyways? + dloader = GraphDataLoader(self.start_dataset, sampler=SequentialSampler(self.start_dataset), batch_size=self.batch_size, drop_last=self.drop_last) + self.graphs = [] + self.labels = [] + self.tracking = [] + self.globals = [] + self.batch_num_nodes = [] + self.batch_num_edges = [] + max_edges = 0 + max_nodes = 0 + load_batch_start = time.time() + for batch, label, tracking, global_feat in dloader: + if batch.num_edges() > max_edges: + max_edges = batch.num_edges() + if batch.num_nodes() > max_nodes: + max_nodes = batch.num_nodes() + self.graphs.append(batch) + self.labels.append(label) + self.tracking.append(tracking) + self.globals.append(global_feat) + load_batch_end = time.time() + print(f'Loaded {len(self.graphs)} batches in {load_batch_end - load_batch_start} seconds') + if self.padding_mode == 'STEPS': + pad_node, pad_edge = utils.pad_size(self.batch_size, max_edges, max_nodes) + elif self.padding_mode == 'FIXED': + print('Padding to fixed size. This is currently hardcoded.') + pad_node = 16000 + pad_edge = 104000 + elif self.padding_mode == 'NONE': + pad_node = 0 + pad_edge = 0 + else: + pad_node = 0 + pad_edge = 0 + print(f'Max edges: {max_edges}, Max nodes: {max_nodes}, Padding to {pad_edge} edges and {pad_node} nodes') + pad_start = time.time() + if self.padding_mode == 'NODE': + for i in range(len(self.graphs)): + unbatched_g = dgl.unbatch(self.graphs[i]) + max_num_nodes = max(g.number_of_nodes() for g in unbatched_g) + self.graphs[i] = utils.pad_batch_num_nodes(self.graphs[i], max_num_nodes) + self.batch_num_nodes.append(self.graphs[i].batch_num_nodes()) + self.batch_num_edges.append(self.graphs[i].batch_num_edges()) + else: + for i in range(len(self.graphs)): + self.graphs[i] = utils.pad_batch(self.graphs[i], pad_edge, pad_node) + self.batch_num_nodes.append(self.graphs[i].batch_num_nodes()) + self.batch_num_edges.append(self.graphs[i].batch_num_edges()) + pad_end = time.time() + print(f'Padded {len(self.graphs)} batches in {pad_end - pad_start} seconds') + + def save(self): + if not self.save_to_disk: + return + graph_path = os.path.join(self.save_dir, f'{self.name}_{self.chunkno}_{self.suffix}.bin') + print(f'Saving dataset to {graph_path}') + if len(self.graphs) == 0: + return + dgl.save_graphs(str(graph_path), self.graphs, {'labels': torch.stack(self.labels), 'batch_num_nodes': torch.stack(self.batch_num_nodes), 'batch_num_edges': torch.stack(self.batch_num_edges), 'tracking': torch.stack(self.tracking), 'globals': torch.stack(self.globals)}) + + def has_cache(self): + if not self.save_to_disk: + return False + for ch in range(self.chunks): + graph_path = os.path.join(self.save_dir, f'{self.name}_{ch}_{self.suffix}.bin') + if not os.path.exists(graph_path): + print(f'Cache file {graph_path} does not exist, not loading from cache.') + return False + return True + + def load(self): + if not self.save_to_disk: + return + self.graphs = [] + label_chunks = [] + tracking_chunks = [] + global_chunks = [] + for ch in range(self.chunks): + graph_path = os.path.join(self.save_dir, f'{self.name}_{ch}_{self.suffix}.bin') + print(f'Loading dataset from {graph_path}') + graphs, label_dict = dgl.load_graphs(graph_path) + label_chunks.append(label_dict['labels']) + tracking_chunks.append(label_dict['tracking']) + global_chunks.append(label_dict['globals']) + for g, bnn, bne in zip(graphs, label_dict['batch_num_nodes'], label_dict['batch_num_edges']): + g.set_batch_num_nodes(bnn) + g.set_batch_num_edges(bne) + self.graphs.extend(graphs) + self.labels = torch.cat(label_chunks) + self.tracking = torch.cat(tracking_chunks) + self.globals = torch.cat(global_chunks) + + def __getitem__(self, idx): + return self.graphs[idx], self.labels[idx], self.tracking[idx], self.globals[idx] + + def __len__(self): + return len(self.graphs) + +#Dataset which contains prebatched shuffled graphs. Cannot be saved to disk, else batching info is lost. +class LazyPreBatchedDataset(PreBatchedDataset): + def __init__(self, **kwargs): + # print(f'Unused kwargs: {kwargs}') + self.current_chunk = None + self.current_chunk_idx = -10 + self.label_chunks = [] + super().__init__(**kwargs) + + def load(self): + if not self.save_to_disk: + return + self.label_chunks = [] + for ch in range(self.chunks): + graph_path = os.path.join(self.save_dir, f'{self.name}_{ch}_{self.suffix}.bin') + print(f'Loading dataset from {graph_path}') + label_dict = dgl.data.graph_serialize.load_labels_v2(graph_path) + self.label_chunks.append(label_dict) + + def __getitem__(self, idx): + chunk_idx = -1 + sum = 0 + ev_idx = -999 + for i in range(len(self.label_chunks)): + count = len(self.label_chunks[i]['labels']) + if idx < sum + count: + chunk_idx = i + ev_idx = idx - sum + break + sum += count + if chunk_idx != self.current_chunk_idx: + # print(f"rank {self.rank} getting data from {self.name}_{chunk_idx}_{self.suffix}.bin") + self.current_chunk, _ = dgl.load_graphs(os.path.join(self.save_dir, f'{self.name}_{chunk_idx}_{self.suffix}.bin')) + self.current_chunk_idx = chunk_idx + g = self.current_chunk[ev_idx] + g.set_batch_num_nodes(self.label_chunks[chunk_idx]['batch_num_nodes'][ev_idx]) + g.set_batch_num_edges(self.label_chunks[chunk_idx]['batch_num_edges'][ev_idx]) + return g, self.label_chunks[chunk_idx]['labels'][ev_idx], self.label_chunks[chunk_idx]['tracking'][ev_idx], self.label_chunks[chunk_idx]['globals'][ev_idx] + + def __len__(self): + l = 0 + for chunk in self.label_chunks: + l += len(chunk['labels']) + return l diff --git a/root_gnn_dgl/root_gnn_base/custom_scheduler.py b/root_gnn_dgl/root_gnn_base/custom_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..b437115e0712665a6b8512e067a4a2df5d7be99a --- /dev/null +++ b/root_gnn_dgl/root_gnn_base/custom_scheduler.py @@ -0,0 +1,565 @@ +import types +import math +import torch +from torch import inf +from functools import wraps, partial +import warnings +import weakref +from collections import Counter +from bisect import bisect_right + +from models import GCN + + + + +### Code from: https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#ReduceLROnPlateau + +Optimizer = torch.optim.Optimizer + +__all__ = ['LambdaLR', 'MultiplicativeLR', 'StepLR', 'MultiStepLR', 'ConstantLR', 'LinearLR', + 'ExponentialLR', 'SequentialLR', 'CosineAnnealingLR', 'ChainedScheduler', 'ReduceLROnPlateau', + 'CyclicLR', 'CosineAnnealingWarmRestarts', 'OneCycleLR', 'PolynomialLR', 'LRScheduler'] + +EPOCH_DEPRECATION_WARNING = ( + "The epoch parameter in `scheduler.step()` was not necessary and is being " + "deprecated where possible. Please use `scheduler.step()` to step the " + "scheduler. During the deprecation, if epoch is different from None, the " + "closed form is used instead of the new chainable form, where available. " + "Please open an issue if you are unable to replicate your use case: " + "https://github.com/pytorch/pytorch/issues/new/choose." +) + + +def update_LR(opt, lr): + for param_group in opt.param_groups: + param_group['lr'] = lr + +def print_LR(opt): + for param_group in opt.param_groups: + print(f"LR = {param_group['lr']}") + +def _check_verbose_deprecated_warning(verbose): + """Raises a warning when verbose is not the default value.""" + if verbose != "deprecated": + warnings.warn("The verbose parameter is deprecated. Please use get_last_lr() " + "to access the learning rate.", UserWarning) + return verbose + return False + +class LRScheduler: + + def __init__(self, optimizer, last_epoch=-1, verbose="deprecated"): + + # Attach optimizer + if not isinstance(optimizer, Optimizer): + raise TypeError(f'{type(optimizer).__name__} is not an Optimizer') + self.optimizer = optimizer + + # Initialize epoch and base learning rates + if last_epoch == -1: + for group in optimizer.param_groups: + group.setdefault('initial_lr', group['lr']) + else: + for i, group in enumerate(optimizer.param_groups): + if 'initial_lr' not in group: + raise KeyError("param 'initial_lr' is not specified " + f"in param_groups[{i}] when resuming an optimizer") + self.base_lrs = [group['initial_lr'] for group in optimizer.param_groups] + self.last_epoch = last_epoch + + # Following https://github.com/pytorch/pytorch/issues/20124 + # We would like to ensure that `lr_scheduler.step()` is called after + # `optimizer.step()` + def with_counter(method): + if getattr(method, '_with_counter', False): + # `optimizer.step()` has already been replaced, return. + return method + + # Keep a weak reference to the optimizer instance to prevent + # cyclic references. + instance_ref = weakref.ref(method.__self__) + # Get the unbound method for the same purpose. + func = method.__func__ + cls = instance_ref().__class__ + del method + + @wraps(func) + def wrapper(*args, **kwargs): + instance = instance_ref() + instance._step_count += 1 + wrapped = func.__get__(instance, cls) + return wrapped(*args, **kwargs) + + # Note that the returned function here is no longer a bound method, + # so attributes like `__func__` and `__self__` no longer exist. + wrapper._with_counter = True + return wrapper + + self.optimizer.step = with_counter(self.optimizer.step) + self.verbose = _check_verbose_deprecated_warning(verbose) + + self._initial_step() + + def _initial_step(self): + """Initialize step counts and performs a step""" + self.optimizer._step_count = 0 + self._step_count = 0 + self.step() + + def state_dict(self): + """Returns the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in self.__dict__ which + is not the optimizer. + """ + return {key: value for key, value in self.__dict__.items() if key != 'optimizer'} + + def load_state_dict(self, state_dict): + """Loads the schedulers state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + self.__dict__.update(state_dict) + + def get_last_lr(self): + """ Return last computed learning rate by current scheduler. + """ + return self._last_lr + + def get_lr(self): + # Compute learning rate using chainable form of the scheduler + raise NotImplementedError + + def print_lr(self, is_verbose, group, lr, epoch=None): + """Display the current learning rate. + """ + if is_verbose: + if epoch is None: + print(f'Adjusting learning rate of group {group} to {lr:.4e}.') + else: + epoch_str = ("%.2f" if isinstance(epoch, float) else + "%.5d") % epoch + print(f'Epoch {epoch_str}: adjusting learning rate of group {group} to {lr:.4e}.') + + + def step(self, epoch=None): + # Raise a warning if old pattern is detected + # https://github.com/pytorch/pytorch/issues/20124 + if self._step_count == 1: + if not hasattr(self.optimizer.step, "_with_counter"): + warnings.warn("Seems like `optimizer.step()` has been overridden after learning rate scheduler " + "initialization. Please, make sure to call `optimizer.step()` before " + "`lr_scheduler.step()`. See more details at " + "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", UserWarning) + + # Just check if there were two first lr_scheduler.step() calls before optimizer.step() + elif self.optimizer._step_count < 1: + warnings.warn("Detected call of `lr_scheduler.step()` before `optimizer.step()`. " + "In PyTorch 1.1.0 and later, you should call them in the opposite order: " + "`optimizer.step()` before `lr_scheduler.step()`. Failure to do this " + "will result in PyTorch skipping the first value of the learning rate schedule. " + "See more details at " + "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", UserWarning) + self._step_count += 1 + + with _enable_get_lr_call(self): + if epoch is None: + self.last_epoch += 1 + values = self.get_lr() + else: + warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning) + self.last_epoch = epoch + if hasattr(self, "_get_closed_form_lr"): + values = self._get_closed_form_lr() + else: + values = self.get_lr() + + for i, data in enumerate(zip(self.optimizer.param_groups, values)): + param_group, lr = data + param_group['lr'] = lr + + self._last_lr = [group['lr'] for group in self.optimizer.param_groups] + + +# Including _LRScheduler for backwards compatibility +# Subclass instead of assign because we want __name__ of _LRScheduler to be _LRScheduler (assigning would make it LRScheduler). +class _LRScheduler(LRScheduler): + pass + + +class _enable_get_lr_call: + + def __init__(self, o): + self.o = o + + def __enter__(self): + self.o._get_lr_called_within_step = True + return self + + def __exit__(self, type, value, traceback): + self.o._get_lr_called_within_step = False + + +class Dynamic_LR(LRScheduler): + """Reduce learning rate when a metric has stopped improving. + Models often benefit from reducing the learning rate by a factor + of 2-10 once learning stagnates. This scheduler reads a metrics + quantity and if no improvement is seen for a 'patience' number + of epochs, the learning rate is reduced. + + Args: + optimizer (Optimizer): Wrapped optimizer. + mode (str): One of `min`, `max`. In `min` mode, lr will + be reduced when the quantity monitored has stopped + decreasing; in `max` mode it will be reduced when the + quantity monitored has stopped increasing. Default: 'min'. + factor (float): Factor by which the learning rate will be + reduced. new_lr = lr * factor. Default: 0.1. + patience (int): Number of epochs with no improvement after + which learning rate will be reduced. For example, if + `patience = 2`, then we will ignore the first 2 epochs + with no improvement, and will only decrease the LR after the + 3rd epoch if the loss still hasn't improved then. + Default: 10. + threshold (float): Threshold for measuring the new optimum, + to only focus on significant changes. Default: 1e-4. + threshold_mode (str): One of `rel`, `abs`. In `rel` mode, + dynamic_threshold = best * ( 1 + threshold ) in 'max' + mode or best * ( 1 - threshold ) in `min` mode. + In `abs` mode, dynamic_threshold = best + threshold in + `max` mode or best - threshold in `min` mode. Default: 'rel'. + cooldown (int): Number of epochs to wait before resuming + normal operation after lr has been reduced. Default: 0. + min_lr (float or list): A scalar or a list of scalars. A + lower bound on the learning rate of all param groups + or each group respectively. Default: 0. + eps (float): Minimal decay applied to lr. If the difference + between new and old lr is smaller than eps, the update is + ignored. Default: 1e-8. + verbose (bool): If ``True``, prints a message to stdout for + each update. Default: ``False``. + + .. deprecated:: 2.2 + ``verbose`` is deprecated. Please use ``get_last_lr()`` to access the + learning rate. + + Example: + >>> # xdoctest: +SKIP + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = ReduceLROnPlateau(optimizer, 'min') + >>> for epoch in range(10): + >>> train(...) + >>> val_loss = validate(...) + >>> # Note that step should be called after validate() + >>> scheduler.step(val_loss) + """ + + def __init__(self, optimizer, mode = 'max', factor=0.1, patience=10, + plateau_var = "test_auc", + threshold=1e-4, threshold_mode='rel', cooldown=0, + min_lr=0, max_lr=1e-4, eps=1e-8, verbose=False): + + """ + if factor >= 1.0: + raise ValueError('Factor should be < 1.0.') + """ + self.factor = factor + + # Attach optimizer + if not isinstance(optimizer, Optimizer): + raise TypeError(f'{type(optimizer).__name__} is not an Optimizer') + self.optimizer = optimizer + + if isinstance(min_lr, (list, tuple)): + if len(min_lr) != len(optimizer.param_groups): + raise ValueError(f"expected {len(optimizer.param_groups)} min_lrs, got {len(min_lr)}") + self.min_lrs = list(min_lr) + self.max_lrs = list(max_lr) + else: + self.min_lrs = [min_lr] * len(optimizer.param_groups) + self.max_lrs = [max_lr] * len(optimizer.param_groups) + + self.patience = patience + self.plateau_var = plateau_var + + self.verbose = verbose + self.cooldown = cooldown + self.cooldown_counter = 0 + self.mode = mode + self.threshold = threshold + self.threshold_mode = threshold_mode + self.best = None + self.num_bad_epochs = None + self.mode_worse = None # the worse value for the chosen mode + self.eps = eps + self.last_epoch = 0 + self._last_lr = [group['lr'] for group in self.optimizer.param_groups] + self._init_is_better(mode=mode, threshold=threshold, + threshold_mode=threshold_mode) + self._reset() + + def _reset(self): + """Resets num_bad_epochs counter and cooldown counter.""" + self.best = self.mode_worse + self.cooldown_counter = 0 + self.num_bad_epochs = 0 + + def step(self, model, metrics, epoch=None): + # convert `metrics` to float, in case it's a zero-dim Tensor + current = float(metrics[self.plateau_var]) + if epoch is None: + epoch = self.last_epoch + 1 + else: + warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning) + self.last_epoch = epoch + + if self.is_better(current, self.best): + if(self.verbose): + print("Model is improving!") + self.best = current + self.num_bad_epochs = 0 + else: + if(self.verbose): + print(f"Model is not improving :( best = {self.best}, current = {current}") + self.num_bad_epochs += 1 + + if self.in_cooldown: + self.cooldown_counter -= 1 + self.num_bad_epochs = 0 # ignore any bad epochs in cooldown + + if self.num_bad_epochs > self.patience: + self._reduce_lr(epoch) + self.cooldown_counter = self.cooldown + self.num_bad_epochs = 0 + + self._last_lr = [group['lr'] for group in self.optimizer.param_groups] + + def _reduce_lr(self, epoch): + print("Adjusting Learning Rate") + self._reset() + for i, param_group in enumerate(self.optimizer.param_groups): + old_lr = float(param_group['lr']) + new_lr = max(old_lr * self.factor, self.min_lrs[i]) + new_lr = min(new_lr, self.max_lrs[i]) + if abs(old_lr - new_lr) > self.eps: + param_group['lr'] = new_lr + + def get_last_lr(self): + return self._last_lr + @property + def in_cooldown(self): + return self.cooldown_counter > 0 + + def is_better(self, a, best): + if self.mode == 'min' and self.threshold_mode == 'rel': + rel_epsilon = 1. - self.threshold + return a < best * rel_epsilon + + elif self.mode == 'min' and self.threshold_mode == 'abs': + return a < best - self.threshold + + elif self.mode == 'max' and self.threshold_mode == 'rel': + rel_epsilon = self.threshold + 1. + return a > best * rel_epsilon + + else: # mode == 'max' and epsilon_mode == 'abs': + return a > best + self.threshold + + def _init_is_better(self, mode, threshold, threshold_mode): + if mode not in {'min', 'max'}: + raise ValueError('mode ' + mode + ' is unknown!') + if threshold_mode not in {'rel', 'abs'}: + raise ValueError('threshold mode ' + threshold_mode + ' is unknown!') + + if mode == 'min': + self.mode_worse = inf + else: # mode == 'max': + self.mode_worse = -inf + + self.mode = mode + self.threshold = threshold + self.threshold_mode = threshold_mode + + def state_dict(self): + return {key: value for key, value in self.__dict__.items() if key != 'optimizer'} + + def load_state_dict(self, state_dict): + self.__dict__.update(state_dict) + self._init_is_better(mode=self.mode, threshold=self.threshold, threshold_mode=self.threshold_mode) + +class Action_On_Plateau(): + + def __init__(self, mode = 'max', patience=10, + plateau_var = "test_auc", + threshold=1e-4, threshold_mode='rel', cooldown=0, + eps=1e-8, verbose=False): + + self.patience = patience + self.plateau_var = plateau_var + + self.verbose = verbose + self.cooldown = cooldown + self.cooldown_counter = 0 + self.mode = mode + self.threshold = threshold + self.threshold_mode = threshold_mode + self.best = None + self.num_bad_epochs = None + self.mode_worse = None # the worse value for the chosen mode + self.eps = eps + self.last_epoch = 0 + self._init_is_better(mode=mode, threshold=threshold, + threshold_mode=threshold_mode) + self._reset() + + def _reset(self): + """Resets num_bad_epochs counter and cooldown counter.""" + self.best = self.mode_worse + self.cooldown_counter = 0 + self.num_bad_epochs = 0 + + def step(self, model, metrics, epoch=None): + # convert `metrics` to float, in case it's a zero-dim Tensor + current = float(metrics[self.plateau_var]) + if epoch is None: + epoch = self.last_epoch + 1 + else: + warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning) + self.last_epoch = epoch + + if self.is_better(current, self.best): + if(self.verbose): + print("Model is improving!") + self.best = current + self.num_bad_epochs = 0 + else: + if(self.verbose): + print(f"Model is not improving :( best = {self.best}, current = {current}") + self.num_bad_epochs += 1 + + if self.in_cooldown: + self.cooldown_counter -= 1 + self.num_bad_epochs = 0 # ignore any bad epochs in cooldown + + if self.num_bad_epochs > self.patience: + self.action(model, metrics, epoch) + + def action(self, model, metrics, epoch=None): + if(self.verbose): + print("Doing my action") + + @property + def in_cooldown(self): + return self.cooldown_counter > 0 + + def is_better(self, a, best): + if self.mode == 'min' and self.threshold_mode == 'rel': + rel_epsilon = 1. - self.threshold + return a < best * rel_epsilon + + elif self.mode == 'min' and self.threshold_mode == 'abs': + return a < best - self.threshold + + elif self.mode == 'max' and self.threshold_mode == 'rel': + rel_epsilon = self.threshold + 1. + return a > best * rel_epsilon + + else: # mode == 'max' and epsilon_mode == 'abs': + return a > best + self.threshold + + def _init_is_better(self, mode, threshold, threshold_mode): + if mode not in {'min', 'max'}: + raise ValueError('mode ' + mode + ' is unknown!') + if threshold_mode not in {'rel', 'abs'}: + raise ValueError('threshold mode ' + threshold_mode + ' is unknown!') + + if mode == 'min': + self.mode_worse = inf + else: # mode == 'max': + self.mode_worse = -inf + + self.mode = mode + self.threshold = threshold + self.threshold_mode = threshold_mode + +class Partial_Reset(Action_On_Plateau): + + def __init__(self, mode='max', patience=10, plateau_var="test_auc", + threshold=0.0001, threshold_mode='rel', cooldown=0, + eps=1e-8, verbose=False): + + super().__init__(mode, patience, plateau_var, threshold, + threshold_mode, cooldown, eps, verbose) + + def action(self, model, metrics, epoch=None): + print("Partial Reset!!") + GCN.partial_reset(model) + self._reset() + self.cooldown_counter = self.cooldown + self.num_bad_epochs = 0 + + +class Full_Reset(Action_On_Plateau): + + def __init__(self, mode='max', patience=10, plateau_var="test_auc", + threshold=0.0001, threshold_mode='rel', cooldown=0, + eps=1e-8, verbose=False): + + super().__init__(mode, patience, plateau_var, threshold, + threshold_mode, cooldown, eps, verbose) + + def action(self, model, metrics, epoch=None): + print("Full Reset!!") + GCN.full_reset(model) + self._reset() + self.cooldown_counter = self.cooldown + self.num_bad_epochs = 0 + +class Dynamic_LR_AND_Partial_Reset(): + def __init__(self, optimizer, mode = 'max', factor=0.1, patience=10, + plateau_var = "test_auc", reset_patience=None, reset_plateau_var=None, + threshold=1e-4, threshold_mode='rel', cooldown=0, + min_lr=0, max_lr=1e-4, eps=1e-8, verbose=False): + + if (reset_patience == None): + reset_patience = patience + if(reset_plateau_var == None): + reset_plateau_var = plateau_var + + self.dynamic_lr = Dynamic_LR(optimizer, mode=mode, factor=factor, patience = patience, + plateau_var=plateau_var, threshold=threshold, threshold_mode =threshold_mode, + cooldown=cooldown, min_lr=min_lr, max_lr=max_lr, eps=eps, verbose=verbose) + + self.partial_reset = Partial_Reset(mode=mode, patience=reset_patience, plateau_var=reset_plateau_var, + threshold=threshold, threshold_mode=threshold_mode, cooldown=cooldown, + eps=eps) + + def step(self, model, metrics, epoch=None): + self.dynamic_lr.step(model=model, metrics=metrics, epoch=epoch) + self.partial_reset.step(model=model, metrics=metrics, epoch=epoch) + +class Dynamic_LR_AND_Full_Reset(): + def __init__(self, optimizer, mode = 'max', factor=0.1, patience=10, + plateau_var = "test_auc", reset_patience=None, reset_plateau_var=None, + threshold=1e-4, threshold_mode='rel', cooldown=0, + min_lr=0, max_lr=1e-4, eps=1e-8, verbose=False): + + if (reset_patience == None): + reset_patience = patience + if(reset_plateau_var == None): + reset_plateau_var = plateau_var + + self.dynamic_lr = Dynamic_LR(optimizer, mode=mode, factor=factor, patience = patience, + plateau_var=plateau_var, threshold=threshold, threshold_mode =threshold_mode, + cooldown=cooldown, min_lr=min_lr, max_lr=max_lr, eps=eps, verbose=verbose) + + self.full_reset = Full_Reset(mode=mode, patience=reset_patience, plateau_var=reset_plateau_var, + threshold=threshold, threshold_mode=threshold_mode, cooldown=cooldown, + eps=eps) + + def step(self, model, metrics, epoch=None): + self.dynamic_lr.step(model=model, metrics=metrics, epoch=epoch) + self.full_reset.step(model=model, metrics=metrics, epoch=epoch) \ No newline at end of file diff --git a/root_gnn_dgl/root_gnn_base/dataset.py b/root_gnn_dgl/root_gnn_base/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..dd34b01e3e1267e52ad909ddfefa076df5faf112 --- /dev/null +++ b/root_gnn_dgl/root_gnn_base/dataset.py @@ -0,0 +1,686 @@ +from dgl.data import DGLDataset +import dgl +import ROOT +import torch +import os +import glob +import time +import numpy as np +from root_gnn_base import utils + +def node_features_from_tree(ch, node_branch_names, node_branch_types, node_feature_scales): + lengths = [] + for branch, node_type in zip(node_branch_names[0], node_branch_types): + if node_type == 'single': + lengths.append(1) + elif node_type == 'vector': + lengths.append(len(getattr(ch, branch))) + else: + print('Unknown node branch type: {}'.format(node_type)) + features = [] + for node_feat in node_branch_names: + if node_feat == 'CALC_E': + features.append(features[0]*torch.cosh(features[1])) + continue + elif node_feat == 'NODE_TYPE': + feat = [] + for i, length in enumerate(lengths): + feat.extend([i,]*length) + features.append(torch.tensor(feat)) + continue + feat = [] + itype = 0 + for length, branch, node_type in zip(lengths, node_feat, node_branch_types): + if isinstance(branch, (int, float, complex)): + feat.extend([branch,]*length) + elif branch == 'CALC_E': + this_type_starts_at = sum(lengths[:itype]) + this_type_ends_at = sum(lengths[:itype+1]) + feat.extend(features[0][this_type_starts_at:this_type_ends_at]*torch.cosh(features[1][this_type_starts_at:this_type_ends_at])) + elif node_type == 'single': + feat.append(getattr(ch, branch)) + elif node_type == 'vector': + feat.extend(getattr(ch, branch)) + itype += 1 + features.append(torch.tensor(feat)) + return torch.stack(features, dim=1) * node_feature_scales, lengths + +def full_connected_graph(n_nodes, self_loops=True): + senders = [] + receivers = [] + senders = np.arange(n_nodes*n_nodes) // n_nodes + receivers = np.arange(n_nodes*n_nodes) % n_nodes + if not self_loops and n_nodes > 1: + mask = senders != receivers + senders = senders[mask] + receivers = receivers[mask] + return dgl.graph((senders, receivers)) + +def check_selection(ch, selection): + var, cut, op = selection + if op == '>': + return getattr(ch, var) > cut + elif op == '<': + return getattr(ch, var) < cut + elif op == '==': + return getattr(ch, var) == cut + +def check_selections(ch, selections): + for selection in selections: + if not check_selection(ch, selection): + return False + return True + +#Base dataset class for making graphs from ROOT ntuples. +class RootDataset(DGLDataset): + def __init__(self, name=None, raw_dir=None, save_dir=None, label=1, file_names = '*.root', node_branch_names=None, node_branch_types=None, node_feature_scales=None, + selections=[], save=True, tree_name = 'nominal_Loose', fold_var = 'eventNumber', weight_var = None, chunks = 1, process_chunks = None, global_features = [], tracking_info = [], **kwargs): + print(f'Unused args while creating RootDataset: {kwargs}') + self.label = label + self.counts = [] + self.selections = selections + self.save_to_disk = save + self.file_names = file_names + self.node_branch_names = node_branch_names + self.node_branch_types = node_branch_types + self.node_feature_scales = torch.tensor([float(sf) for sf in node_feature_scales]) + self.tree_name = tree_name + self.fold_var = fold_var + self.tracking_info = tracking_info + self.tracking_info.insert(0, fold_var) + if weight_var == None: + weight_var = 1 + self.tracking_info.insert(1, weight_var) + self.global_features = global_features + self.chunks = chunks + self.process_chunks = process_chunks + if self.process_chunks is None: + self.process_chunks = [i for i in range(self.chunks)] + self.times = [0, 0] + super().__init__(name=name, raw_dir=raw_dir, save_dir=save_dir) + + def get_list_of_branches(self): + branches = [] + for feat in self.node_branch_names: + if isinstance(feat, list): + for branch in feat: + if branch == 'CALC_E': + continue + if isinstance(branch, str): + branches.append(branch) + for feat in self.global_features: + if isinstance(feat, str): + branches.append(feat) + for feat in self.tracking_info: + if isinstance(feat, str): + branches.append(feat) + for selection in self.selections: + branches.append(selection[0]) + return branches + + def make_graph(self, ch): + t1 = time.time() + features, _ = node_features_from_tree(ch, self.node_branch_names, self.node_branch_types, self.node_feature_scales) + features = features[features[:,0] != 0] + t2 = time.time() + g = full_connected_graph(features.shape[0], self_loops=False) + g.ndata['features'] = features + t3 = time.time() + self.times[0] += t2 - t1 + self.times[1] += t3 - t2 + return g + + def process(self): + times = [0, 0, 0] + oldtime = time.time() + if isinstance(self.file_names, str): + self.files = glob.glob(os.path.join(self.raw_dir, self.file_names)) + else: + self.files = [] + for file_name in self.file_names: + self.files.extend(glob.glob(os.path.join(self.raw_dir, file_name))) + self.chain = ROOT.TChain(self.tree_name) + + if len(self.files) == 0: + print('No files found in {}'.format(os.path.join(self.raw_dir, self.file_names))) + for file in self.files: + utils.set_timeout(60*2) + self.chain.Add(file) + utils.unset_timeout() + branches = self.get_list_of_branches() + self.chain.SetBranchStatus('*', 0) + for branch in branches: + self.chain.SetBranchStatus(branch, 1) + newtime = time.time() + times[0] += newtime - oldtime + chunks = np.array_split(np.arange(self.chain.GetEntries()), self.chunks) + chunks = [chunk for i, chunk in enumerate(chunks) if i in self.process_chunks] + + self.graph_chunks = [] + self.label_chunks = [] + self.tracking_chunks = [] + self.global_chunks = [] + chunk_id = -1 + for chunk in chunks: + chunk_id += 1 + graphs = [] + labels = [] + tracking = [] + globals = [] + for ientry in chunk: + if (ientry % 10000 == 0): + print('Processing event {}/{}'.format(ientry, self.chain.GetEntries()), flush=True) + self.chain.GetEntry(ientry) + passed = True + for selection in self.selections: + if not check_selection(self.chain, selection): + passed = False + continue + oldtime = newtime + newtime = time.time() + times[1] += newtime - oldtime + if passed: + graphs.append(self.make_graph(self.chain)) + labels.append( self.label ) + tracking.append(torch.zeros(len(self.tracking_info), dtype=torch.double)) + globals.append(torch.zeros(len(self.global_features))) + for i_ti, tr_branch in enumerate(self.tracking_info): + if isinstance(tr_branch, str): + tracking[-1][i_ti] = getattr(self.chain, tr_branch) + else: + tracking[-1][i_ti] = tr_branch + for i_gl, gl_branch in enumerate(self.global_features): + globals[-1][i_gl] = getattr(self.chain, gl_branch) + oldtime = newtime + newtime = time.time() + times[2] += newtime - oldtime + + labels = torch.tensor(labels) + tracking = torch.stack(tracking) + globals = torch.stack(globals) + + if (self.chunks > 1): + self.save_chunk(chunk_id, graphs, labels, tracking, globals) + else: + self.labels = labels + self.tracking = tracking + self.global_features = globals + self.graphs = graphs + self.save() + return + self.graphs = self.graph_chunks[0] + for chunk in self.graph_chunks[1:]: + self.graphs += chunk + self.labels = torch.cat(self.label_chunks) + self.tracking = torch.cat(self.tracking_chunks) + self.global_features = torch.cat(self.global_chunks) + print('Time spent: Creating TChain: {}s, Getting Entries and Selection: {}s, Graph Creation: {}s'.format(*times)) + print('Time spent in node_features_from_tree: {}s, full_connected_graph: {}s'.format(*self.times)) + + def save(self): + """save the graph list and the labels""" + if not self.save_to_disk: + return + graph_path = os.path.join(self.save_dir, self.name + '.bin') + if self.chunks == 1: + # print(len(self.graphs)) + # print(len(self.labels)) + # print(len(self.tracking)) + # print(len(self.globals)) + print(f'Saving dataset to {os.path.join(self.save_dir, self.name + ".bin")}') + dgl.save_graphs(str(graph_path), self.graphs, {'labels': torch.tensor(self.labels), 'tracking': torch.tensor(self.tracking), 'global': torch.tensor(self.global_features)}) + else: + print(len(self.graph_chunks)) + for i in range(len(self.process_chunks)): + print(f'Saving dataset to {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[i]}.bin")}') + dgl.save_graphs(str(graph_path).replace('.bin', f'_{self.process_chunks[i]}.bin'), self.graph_chunks[i], {'labels': self.label_chunks[i], 'tracking': self.tracking_chunks[i], 'global': self.global_chunks[i]}) + + def save_chunk(self, chunk_id, graphs, labels, tracking, globals): + if not self.save_to_disk: + return + graph_path = os.path.join(self.save_dir, self.name + '.bin') + print(f'Saving dataset to {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[chunk_id]}.bin")}') + dgl.save_graphs(str(graph_path).replace('.bin', f'_{self.process_chunks[chunk_id]}.bin'), graphs, {'labels': labels, 'tracking': tracking, 'global': globals}) + + def has_cache(self): + print(f'Checking for cache of {self.name}') + if not self.save_to_disk: + print('Skipping load.') + return False + if self.chunks == 1: + graph_path = os.path.join(self.save_dir, self.name + '.bin') + return os.path.exists(graph_path) + else: + for i in range(len(self.process_chunks)): + graph_path = os.path.join(self.save_dir, self.name + f'_{self.process_chunks[i]}.bin') + if not os.path.exists(graph_path): + print(f'File {graph_path} does not exist, processing.') + return False + return True + + def load(self): + if self.chunks == 1: + print(f'Loading dataset from {os.path.join(self.save_dir, self.name + ".bin")}') + graphs, label_dict = dgl.load_graphs(os.path.join(self.save_dir, self.name + '.bin')) + self.graphs = graphs + self.labels = label_dict['labels'] + self.tracking = label_dict['tracking'] + self.global_features = label_dict['global'] + else: + self.graphs = [] + self.labels = [] + self.tracking = [] + self.global_features = [] + for i in range(self.chunks): + try: + print(f'Loading dataset from {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[i]}.bin")}') + graphs, label = dgl.load_graphs(os.path.join(self.save_dir, self.name + f'_{self.process_chunks[i]}.bin')) + self.graphs.extend(graphs) + self.labels.append(label['labels']) + self.tracking.append(label['tracking']) + self.global_features.append(label['global']) + except Exception as e: + print(e) + self.labels = torch.cat(self.labels) + self.tracking = torch.cat(self.tracking) + self.global_features = torch.cat(self.global_features) + + def __getitem__(self, idx): + return self.graphs[idx], self.labels[idx], self.tracking[idx], self.global_features[idx] + + def __len__(self): + return len(self.graphs) + +#Dataset with edge features added (deta, dphi, dR) +class EdgeDataset(RootDataset): + def make_graph(self, ch): + g = super().make_graph(ch) + u, v = g.edges() + deta = g.ndata['features'][u, 1] - g.ndata['features'][v, 1] + dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2] + dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi) + dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi) + dR = torch.sqrt(deta**2 + dphi**2) + g.edata['features'] = torch.stack([deta, dphi, dR], dim=1) + return g + +class tHbbEdgeDataset(RootDataset): + def __init__(self, exclude_branches=None, **kwargs): + self.exclude_branches = exclude_branches + super().__init__(**kwargs) + + def get_list_of_branches(self): + br = super().get_list_of_branches() + for sector in self.exclude_branches: + if sector == None: + continue + for excl in sector: + if type(excl) == str: + br.append(excl) + return br + + def make_graph(self, ch): + features, lengths = node_features_from_tree(ch, self.node_branch_names, self.node_branch_types, self.node_feature_scales) + + include_mask = torch.ones(features.shape[0], dtype=torch.bool) + node_idx = 0 + for sector, length in zip(self.exclude_branches, lengths): + if sector == None: + node_idx += length + continue + for excl in sector: + if type(excl) == int: + include_mask[excl + node_idx] = False + elif type(excl) == str: + include_mask[getattr(self.chain, excl) + node_idx] = False + g = full_connected_graph(features[include_mask].shape[0], self_loops=False) + g.ndata['features'] = features[include_mask] + + u, v = g.edges() + deta = g.ndata['features'][u, 1] - g.ndata['features'][v, 1] + dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2] + dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi) + dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi) + dR = torch.sqrt(deta**2 + dphi**2) + g.edata['features'] = torch.stack([deta, dphi, dR], dim=1) + return g + +class LazyDataset(EdgeDataset): + def __init__(self, buffer_size = 2, **kwargs): + self.buffer = [None,] * buffer_size + self.buffer_ptr = 0 + self.get_item_calls = 0 + self.buffer_indices = [-1,] * buffer_size + super().__init__(**kwargs) + + def __getitem__(self, idx): + self.get_item_calls += 1 + chunk_idx = -1 + sum = 0 + ev_idx = -999 + for i, count in enumerate(self.counts): + sum += count + if idx < sum: + chunk_idx = i + ev_idx = idx - sum + count + break + buf_idx = self.buffer_get(chunk_idx) + if ev_idx >= len(self.buffer[buf_idx][0]): + print(f'Getting event {ev_idx} from chunk {chunk_idx} from buffer {buf_idx}. Calls: {self.get_item_calls}') + print(len(self.buffer)) + print(self.counts) + print(len(self.buffer[buf_idx][0])) + return self.buffer[buf_idx][0][ev_idx], self.buffer[buf_idx][1]['labels'][ev_idx], self.buffer[buf_idx][1]['tracking'][ev_idx], self.buffer[buf_idx][1]['global'][ev_idx] + + def buffer_get(self, buffer_idx): + if buffer_idx in self.buffer_indices: + for i in range(len(self.buffer)): + if self.buffer_indices[i] == buffer_idx: + return i + else: + print(f'Loading dataset from {os.path.join(self.save_dir, self.name + f"_{buffer_idx}.bin")}', flush=True) + self.buffer_ptr = (self.buffer_ptr + 1) % len(self.buffer) + self.buffer[self.buffer_ptr] = dgl.load_graphs(os.path.join(self.save_dir, self.name + f'_{buffer_idx}.bin')) + self.buffer_indices[self.buffer_ptr] = buffer_idx + return self.buffer_ptr + + def load(self): + self.counts = [] + self.tracking = [] + try: + for i in range(self.chunks): + print(f'Loading dataset from {os.path.join(self.save_dir, self.name + f"_{self.process_chunks[i]}.bin")}') + l = dgl.data.graph_serialize.load_labels_v2(os.path.join(self.save_dir, self.name + f'_{self.process_chunks[i]}.bin')) + self.counts.append(len(l['tracking'])) + self.tracking.append(l['tracking']) + self.tracking = torch.cat(self.tracking) + except Exception as e: + print(e) + + def __len__(self): + return sum(self.counts) + +class MultiLabelDataset(EdgeDataset): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def get_list_of_branches(self): + br = super().get_list_of_branches() + for l in self.label: + if isinstance(l, str): + br.append(l) + if isinstance(l, dict): + br.append(l['branch']) + return br + + def get_label(self, ch): + label = [] + for l in self.label: + if isinstance(l, str): + label.append((getattr(ch, l))) + if isinstance(l, dict): + label.append(getattr(ch, l['branch'])*float(l['scale'])) + if isinstance(l, float) or isinstance(l, int): + label.append(l) + + return torch.tensor(label) + + def process(self): + times = [0, 0, 0] + oldtime = time.time() + if isinstance(self.file_names, str): + self.files = glob.glob(os.path.join(self.raw_dir, self.file_names)) + else: + self.files = [] + for file_name in self.file_names: + self.files.extend(glob.glob(os.path.join(self.raw_dir, file_name))) + self.chain = ROOT.TChain(self.tree_name) + if len(self.files) == 0: + print('No files found in {}'.format(os.path.join(self.raw_dir, self.file_names))) + for file in self.files: + utils.set_timeout(60*2) + self.chain.Add(file) + utils.unset_timeout() + branches = self.get_list_of_branches() + self.chain.SetBranchStatus('*', 0) + for branch in branches: + self.chain.SetBranchStatus(branch, 1) + newtime = time.time() + times[0] += newtime - oldtime + chunks = np.array_split(np.arange(self.chain.GetEntries()), self.chunks) + chunks = [chunk for i, chunk in enumerate(chunks) if i in self.process_chunks] + self.graph_chunks = [] + self.label_chunks = [] + self.tracking_chunks = [] + self.global_chunks = [] + chunk_id = -1 + for chunk in chunks: + chunk_id += 1 + graphs = [] + labels = [] + tracking = [] + globals = [] + for ientry in chunk: + if (ientry % 10000 == 0): + print('Processing event {}/{}'.format(ientry, self.chain.GetEntries()), flush=True) + self.chain.GetEntry(ientry) + passed = True + for selection in self.selections: + if not check_selection(self.chain, selection): + passed = False + continue + oldtime = newtime + newtime = time.time() + times[1] += newtime - oldtime + if passed: + graphs.append(self.make_graph(self.chain)) + labels.append(self.get_label(self.chain)) + tracking.append(torch.zeros(len(self.tracking_info), dtype=torch.double)) + globals.append(torch.zeros(len(self.global_features))) + for i_ti, tr_branch in enumerate(self.tracking_info): + if isinstance(tr_branch, str): + tracking[-1][i_ti] = getattr(self.chain, tr_branch) + else: + tracking[-1][i_ti] = tr_branch + for i_gl, gl_branch in enumerate(self.global_features): + globals[-1][i_gl] = getattr(self.chain, gl_branch) + oldtime = newtime + newtime = time.time() + times[2] += newtime - oldtime + + labels = torch.stack(labels) + self.save_chunk(chunk_id, graphs, labels, torch.stack(tracking), torch.stack(globals)) + # self.graph_chunks.append(graphs) + # self.label_chunks.append(labels) + # self.tracking_chunks.append(torch.stack(tracking)) + # self.global_chunks.append(torch.stack(globals)) + # self.counts.append(len(graphs)) + return + self.graphs = self.graph_chunks[0] + for chunk in self.graph_chunks[1:]: + self.graphs += chunk + + self.labels = torch.cat(self.label_chunks) + self.tracking = torch.cat(self.tracking_chunks) + self.global_features = torch.cat(self.global_chunks) + print('Time spent: Creating TChain: {}s, Getting Entries and Selection: {}s, Graph Creation: {}s'.format(*times)) + print('Time spent in node_features_from_tree: {}s, full_connected_graph: {}s'.format(*self.times)) + +class LazyMultiLabelDataset(MultiLabelDataset, LazyDataset): + def __init__(self, buffer_size = 2, **kwargs): + LazyDataset.__init__(self, buffer_size=buffer_size, **kwargs) + +class MultiLabeltHbbDataset(MultiLabelDataset, tHbbEdgeDataset): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + def get_list_of_branches(self): + br = super().get_list_of_branches() + for sector in self.exclude_branches: + if sector == None: + continue + for excl in sector: + if type(excl) == str: + br.append(excl) + return br + + +class AugmentedDataset(RootDataset): + + def __init__(self, seed = 2, feature_index = None, node_mapping = None, **kwargs): + self.seed = seed + np.random.seed(seed) + if(feature_index == None): + self.feature_index = {"pt": 0, "eta": 1, "phi": 2, "energy": 3, "btag": 4, "charge": 5, "node_type": 6} + if (node_mapping == None): + self.node_mapping = {"jet": 0, "ele": 1, "mu": 2, "ph": 3, "MET": 4} + super().__init__(**kwargs) + + def detector_noise(self, node_features): + noise = np.zeros_like(node_features) + + node_types = node_features[:, self.feature_index["node_type"]] + pts = node_features[:, self.feature_index["pt"]] + etas = node_features[:, self.feature_index["eta"]] + energies = node_features[:, self.feature_index["energy"]] + + # Noise calculation for jets + jet_mask = (node_types == self.node_mapping["jet"]) + jet_pts = pts[jet_mask] + jet_etas = etas[jet_mask] + + if (jet_mask.sum() > 0): + jet_resolutions = np.where( + jet_pts <= 0.1, 0.0, + np.where( + np.abs(jet_etas) <= 0.5, np.sqrt(0.06**2 + jet_pts**2 * 1.3e-3**2), + np.where( + np.abs(jet_etas) <= 1.5, np.sqrt(0.10**2 + jet_pts**2 * 1.7e-3**2), + np.where( + np.abs(jet_etas) <= 2.5, np.sqrt(0.25**2 + jet_pts**2 * 3.1e-3**2), + 0.0 + ) + ) + ) + ) + noise[jet_mask, self.feature_index["pt"]] = np.random.normal(loc=0.0, scale=jet_resolutions) + + # Noise calculation for electrons + ele_mask = (node_types == self.node_mapping["ele"]) + ele_pts = pts[ele_mask] + ele_etas = etas[ele_mask] + + if (ele_mask.sum() > 0): + ele_resolutions = np.where( + np.abs(ele_etas) <= 0.5, np.sqrt(0.03**2 + ele_pts**2 * 1.3e-3**2), + np.where( + np.abs(ele_etas) <= 1.5, np.sqrt(0.05**2 + ele_pts**2 * 1.7e-3**2), + np.where( + np.abs(ele_etas) <= 2.5, np.sqrt(0.15**2 + ele_pts**2 * 3.1e-3**2), + 0.0 + ) + ) + ) + noise[ele_mask, self.feature_index["pt"]] = np.random.normal(loc=0.0, scale=ele_resolutions) + + # Noise calculation for muons + mu_mask = (node_types == self.node_mapping["mu"]) + mu_pts = pts[mu_mask] + mu_etas = etas[mu_mask] + + if (mu_mask.sum() > 0): + mu_resolutions = np.where( + np.abs(mu_etas) <= 0.5, np.sqrt(0.01**2 + mu_pts**2 * 1.0e-4**2), + np.where( + np.abs(mu_etas) <= 1.5, np.sqrt(0.015**2 + mu_pts**2 * 1.5e-4**2), + np.where( + np.abs(mu_etas) <= 2.5, np.sqrt(0.025**2 + mu_pts**2 * 3.5e-4**2), + 0.0 + ) + ) + ) + noise[mu_mask, self.feature_index["pt"]] = np.random.normal(loc=0.0, scale=mu_resolutions) + + # Noise calculation for photons + ph_mask = (node_types == self.node_mapping["ph"]) + ph_etas = etas[ph_mask] + ph_energies = energies[ph_mask] + + if (ph_mask.sum() > 0): + ph_resolutions = np.where( + np.abs(ph_etas) <= 3.2, np.sqrt(ph_energies**2 * 0.0017**2 + ph_energies * 0.101**2), + np.where( + np.abs(ph_etas) <= 4.9, np.sqrt(ph_energies**2 * 0.0350**2 + ph_energies * 0.285**2), + 0.0 + ) + ) + noise[ph_mask, self.feature_index["energy"]] = np.random.normal(loc=0.0, scale=ph_resolutions) + return noise + + def make_graph(self, ch): + g = super().make_graph(ch) + + g.ndata['augmented_features'] = g.ndata['features'] + + num_nodes = len(g.ndata['features'][:, 0]) + + # Rotations: phi -> phi + delta_phi + phi_index = self.feature_index["phi"] + # Generate a single delta_phi for all nodes + delta_phi = np.random.uniform(low=-np.pi, high=np.pi) + + # Apply the same delta_phi to all nodes + g.ndata['augmented_features'][:, phi_index] = (g.ndata['augmented_features'][:, phi_index] + delta_phi + np.pi) % (2 * np.pi) - np.pi + + # Reflections: eta -> -1 * eta, phi -> -1 * phi + eta_index = self.feature_index["eta"] + + eta_reflection = np.random.choice([-1, 1]) + phi_reflection = np.random.choice([-1, 1]) + + g.ndata['augmented_features'][:, eta_index] = g.ndata['augmented_features'][:, eta_index] * eta_reflection + g.ndata['augmented_features'][:, phi_index] = g.ndata['augmented_features'][:, phi_index] * phi_reflection + + + # Detector Noise: pt -> pt + normal(pt, noise(pt)) + noise = self.detector_noise(g.ndata['augmented_features']) + g.ndata['augmented_features'] = g.ndata['augmented_features'] + noise + + pt_index = self.feature_index["pt"] + if (g.ndata['augmented_features'][-1][self.feature_index["node_type"]] == self.node_mapping["MET"]): + # Initialize sums of px and py + sum_px = 0 + sum_py = 0 + + # Loop over all nodes except the last one (MET node) + for i in range(len(g.ndata['augmented_features']) - 1): + pt = g.ndata['augmented_features'][i][pt_index] + phi = g.ndata['augmented_features'][i][phi_index] + + # Compute px and py + px = pt * np.cos(phi) + py = pt * np.sin(phi) + + # Sum px and py + sum_px += px + sum_py += py + + # Calculate MET + g.ndata['augmented_features'][-1][pt_index] = np.sqrt(sum_px**2 + sum_py**2) + + u, v = g.edges() + deta = g.ndata['features'][u, 1] - g.ndata['features'][v, 1] + dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2] + dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi) + dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi) + dR = torch.sqrt(deta**2 + dphi**2) + g.edata['features'] = torch.stack([deta, dphi, dR], dim=1) + + deta = g.ndata['augmented_features'][u, 1] - g.ndata['augmented_features'][v, 1] + dphi = g.ndata['augmented_features'][u, 2] - g.ndata['augmented_features'][v, 2] + dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi) + dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi) + dR = torch.sqrt(deta**2 + dphi**2) + g.edata['augmented_features'] = torch.stack([deta, dphi, dR], dim=1) + + return g \ No newline at end of file diff --git a/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py b/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..05c893bdfc423b6e95ef5c87a382b36c68e68a4e --- /dev/null +++ b/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py @@ -0,0 +1,44 @@ +from root_gnn_base import dataset +import dgl +import torch +import numpy as np + +def radius_graph(features, radii, self_loops=False): + senders = [] + receivers = [] + n_nodes = features.shape[0] + senders = np.arange(n_nodes*n_nodes) // n_nodes + receivers = np.arange(n_nodes*n_nodes) % n_nodes + if not self_loops and n_nodes > 1: + mask = senders != receivers + senders = senders[mask] + receivers = receivers[mask] + for k, r in radii.items(): + d = features[senders, k] - features[receivers, k] + mask = np.abs(d) < r + senders = senders[mask] + receivers = receivers[mask] + return dgl.graph((senders, receivers)) + +class PhotonIDDataset(dataset.LazyMultiLabelDataset): + def __init__(self, eta_radius, phi_radius, **kwargs): + self.eta_radius = eta_radius + self.phi_radius = phi_radius + super().__init__(**kwargs) + def make_graph(self, ch): + features, _ = dataset.node_features_from_tree(ch, self.node_branch_names, self.node_branch_types, self.node_feature_scales) + features = features[features[:,0] != 0] + #Delta Eta, Delta Phi, Adjacent Layer + g = radius_graph(features, {1: self.eta_radius, 2: self.phi_radius, 6: 1.1}, self_loops=True) #Self loops ensure last cell is included even if disconnected + g.ndata['features'] = features + u, v = g.edges() + deta = features[u, 1] - features[v, 1] + dphi = g.ndata['features'][u, 2] - g.ndata['features'][v, 2] + dphi = torch.where(dphi > np.pi, dphi - 2*np.pi, dphi) + dphi = torch.where(dphi < -np.pi, dphi + 2*np.pi, dphi) + dR = torch.sqrt(deta**2 + dphi**2) + dx = features[u, 3] - features[v, 3] + dy = features[u, 4] - features[v, 4] + dz = features[u, 5] - features[v, 5] + g.edata['features'] = torch.stack([deta, dphi, dR, dx, dy, dz], dim=1) + return g \ No newline at end of file diff --git a/root_gnn_dgl/root_gnn_base/similarity.py b/root_gnn_dgl/root_gnn_base/similarity.py new file mode 100644 index 0000000000000000000000000000000000000000..9a6aaadc11ad929671e346acb794d7268fcb2468 --- /dev/null +++ b/root_gnn_dgl/root_gnn_base/similarity.py @@ -0,0 +1,158 @@ +import numpy as np +import scipy +from sklearn.decomposition import PCA +from sklearn.metrics.pairwise import cosine_similarity +from sklearn.metrics.pairwise import euclidean_distances +from sklearn.preprocessing import StandardScaler + +from scipy.stats import wasserstein_distance + +def cka(rep_a, rep_b, size=None): + """ + Computes the Centered Kernel Alignment (CKA) between two large representation matrices rep_a and rep_b. + If size is provided, it performs CKA on a randomly selected subset of the data. + + Parameters: + rep_a : np.ndarray + First representation matrix of size (n_samples, n_features_a). + rep_b : np.ndarray + Second representation matrix of size (n_samples, n_features_b). + size : int, optional + Number of samples to use for the CKA calculation. If None, use the full dataset. + + Returns: + float + CKA similarity between rep_a and rep_b. + """ + + def gram_linear(x): + """Compute the Gram (kernel) matrix using a linear kernel.""" + return x @ x.T + + def center_gram(gram): + """Center the Gram matrix.""" + n = gram.shape[0] + identity = np.eye(n) + ones = np.ones((n, n)) / n + return gram - ones @ gram - gram @ ones + ones @ gram @ ones + + # If sample_size is specified, randomly sample a subset of the data + if size is not None and size < rep_a.shape[0]: + indices = np.random.choice(rep_a.shape[0], size, replace=False) + rep_a = rep_a[indices] + rep_b = rep_b[indices] + + # Compute the Gram matrices + gram_a = gram_linear(rep_a) + gram_b = gram_linear(rep_b) + + # Center the Gram matrices + centered_gram_a = center_gram(gram_a) + centered_gram_b = center_gram(gram_b) + + # Compute the CKA similarity + numerator = np.sum(centered_gram_a * centered_gram_b) + denominator = np.sqrt(np.sum(centered_gram_a**2) * np.sum(centered_gram_b**2)) + + return numerator / denominator if denominator != 0 else 0 + +def cca(X, Y, size = None, num_components=10): + """ + Perform Canonical Correlation Analysis (CCA) between two datasets. + + Parameters: + X : np.ndarray + First dataset, shape (n_samples, n_features_X). + Y : np.ndarray + Second dataset, shape (n_samples, n_features_Y). + num_components : int + Number of CCA components to return. + + Returns: + w_X : np.ndarray + Canonical weights for the first dataset, shape (n_features_X, num_components). + w_Y : np.ndarray + Canonical weights for the second dataset, shape (n_features_Y, num_components). + corrs : np.ndarray + Array of canonical correlations for each component. + """ + + # If sample size is specified, randomly sample a subset of the data + if size is not None and size < X.shape[0]: + indices = np.random.choice(X.shape[0], size, replace=False) + X = X[indices] + Y = Y[indices] + + # Standardize both datasets (mean = 0, variance = 1) + scaler_X = StandardScaler() + scaler_Y = StandardScaler() + + X = scaler_X.fit_transform(X) + Y = scaler_Y.fit_transform(Y) + + # Covariance matrices + C_XX = np.cov(X, rowvar=False) # Covariance of X + C_YY = np.cov(Y, rowvar=False) # Covariance of Y + C_XY = np.cov(X, Y, rowvar=False)[:X.shape[1], X.shape[1]:] # Cross-covariance of X and Y + + # Regularization term to avoid singular matrices + reg = 1e-6 + inv_C_XX = np.linalg.inv(C_XX + reg * np.eye(C_XX.shape[0])) + inv_C_YY = np.linalg.inv(C_YY + reg * np.eye(C_YY.shape[0])) + + # Solve the generalized eigenvalue problem for CCA + # (inv_C_XX @ C_XY @ inv_C_YY @ C_XY.T) and vice versa for Y + A = inv_C_XX @ C_XY @ inv_C_YY @ C_XY.T + B = inv_C_YY @ C_XY.T @ inv_C_XX @ C_XY + + # Perform eigenvalue decomposition + eigvals_X, eigvecs_X = np.linalg.eigh(A) + eigvals_Y, eigvecs_Y = np.linalg.eigh(B) + + # Sort the eigenvalues and eigenvectors in descending order + idx_X = np.argsort(eigvals_X)[::-1] + idx_Y = np.argsort(eigvals_Y)[::-1] + + eigvecs_X = eigvecs_X[:, idx_X] + eigvecs_Y = eigvecs_Y[:, idx_Y] + + # Canonical weights (the first `num_components` components) + w_X = eigvecs_X[:, :num_components] + w_Y = eigvecs_Y[:, :num_components] + + # Canonical correlations (square root of the eigenvalues, constrained to [0,1]) + corrs = np.sqrt(np.clip(eigvals_X[:num_components], 0, 1)) + + return np.mean(corrs) + return w_X, w_Y, corrs + +def pca(X, Y, size=1000, n_components=3, bins=30): + + pca_X = PCA(n_components=n_components) + X_pca = pca_X.fit_transform(X) + + pca_Y = PCA(n_components=n_components) + Y_pca = pca_Y.fit_transform(Y) + + # Step 2: Determine common bin edges based on the range of PCA components + min_value = min(X_pca.min(), Y_pca.min()) + max_value = max(X_pca.max(), Y_pca.max()) + bin_edges = np.linspace(min_value, max_value, bins + 1) + + # Step 3: Calculate histograms for each PCA component using the same bins + histograms_X = [np.histogram(X_pca[:, i], bins=bin_edges, density=True)[0] for i in range(n_components)] + histograms_Y = [np.histogram(Y_pca[:, i], bins=bin_edges, density=True)[0] for i in range(n_components)] + + # Step 4: Calculate Wasserstein distance between corresponding histograms + total_distance = 0 + for i in range(n_components): + total_distance += wasserstein_distance(histograms_X[i], histograms_Y[i]) + + # Step 5: Normalize the total distance for a similarity score + # Calculate the maximum possible distance (theoretical max could be based on histogram size) + # This could be replaced with a more complex calculation if necessary. + max_distance = 1.0 # Replace this with a suitable maximum based on your dataset properties. + + similarity_score = 1 - (total_distance / max_distance) + + return max(0, min(1, similarity_score)) # Ensure the score stays in [0, 1] diff --git a/root_gnn_dgl/root_gnn_base/uproot_dataset.py b/root_gnn_dgl/root_gnn_base/uproot_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..04f06f895ef15e713f85e46ab42c76f757aedd5e --- /dev/null +++ b/root_gnn_dgl/root_gnn_base/uproot_dataset.py @@ -0,0 +1,54 @@ +from root_gnn_base import dataset +import torch +import uproot +import glob +import os +import awkward as ak +import numpy as np +import time + +def node_features_from_ak(ch, node_branch_names, node_branch_types, node_feature_scales): + node_types = [] + n_types = len(node_branch_names[0]) + for i in range(n_types): + features = [] + branch_type = node_branch_types[i] + for j in range(len(node_branch_names)): + if node_branch_names[j] == 'CALC_E': + features.append(features[0] * np.cosh(features[1])) + elif node_branch_names[j] == 'NODE_TYPE': + features.append(ak.full_like(features[0], i)) + elif isinstance(node_branch_names[j][i], str): + features.append(ch[node_branch_names[j][i]]) + elif isinstance(node_branch_names[j][i], (int, float)): + features.append(ak.full_like(features[0], node_branch_names[j][i])) + if branch_type == 'single': + features = [f[:,np.newaxis] for f in features] + node_types.append(ak.Array(features)) + node_features = ak.concatenate(node_types, axis=2) * node_feature_scales #axis order at this point is (feature, event, node) + return node_features + +class UprootDataset(dataset.RootDataset): + def process(self): + starttime = time.time() + self.files = glob.glob(os.path.join(self.raw_dir, self.file_names)) + branches = self.get_list_of_branches() + self.chain = uproot.concatenate([f + ':' + self.tree_name for f in self.files], branches, num_workers=4) + node_features = node_features_from_ak(self.chain, self.node_branch_names, self.node_branch_types, self.node_feature_scales) + loadtime = time.time() + n_nodes = ak.num(node_features[0], axis=1) #number of nodes for each event + ftime = time.time() + self.graphs = [dataset.full_connected_graph(n, False) for n in n_nodes] + itime = time.time() + for i in range(len(self.graphs)): + if i % 10000 == 0: + print(f'Processing event {i}/{len(self.graphs)}') + self.graphs[i].ndata['features'] = torch.transpose(torch.tensor(node_features[:,i,:]),0,1).to(torch.float) + self.label = torch.stack([torch.full((len(self.graphs),),torch.tensor(self.label)), torch.tensor(ak.values_astype(self.chain[self.fold_var], np.int64))], dim=1) + gtime = time.time() + print() + print(f'load time: {loadtime - starttime} s') + print(f'feature time: {ftime - loadtime} s') + print(f'graph time: {itime - ftime} s') + print(f'graph data time: {gtime - itime} s') + diff --git a/root_gnn_dgl/root_gnn_base/utils.py b/root_gnn_dgl/root_gnn_base/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8f4ff676c074a05c90499318c5fe65f3b2311eb8 --- /dev/null +++ b/root_gnn_dgl/root_gnn_base/utils.py @@ -0,0 +1,307 @@ +import importlib +import yaml +import os +import torch +import numpy as np +import matplotlib.pyplot as plt +import dgl +import signal + +def buildFromConfig(conf, run_time_args = {}): + if 'module' in conf: + module = importlib.import_module(conf['module']) + cls = getattr(module, conf['class']) + return cls(**conf['args'], **run_time_args) + else: + print('No module specified in config. Returning None.') + +def cycler(iterable): + while True: + #print('Cycler is cycling...') + for i in iterable: + yield i + +def include_config(conf): + if 'include' in conf: + for i in conf['include']: + with open(i) as f: + conf.update(yaml.load(f, Loader=yaml.FullLoader)) + del conf['include'] + +def load_config(config_file): + with open(config_file) as f: + conf = yaml.load(f, Loader=yaml.FullLoader) + include_config(conf) + return conf + +#Timeout function from https://stackoverflow.com/questions/492519/timeout-on-a-function-call +class TimeoutException(Exception): + pass + +def timeout_handler(signum, frame): + raise TimeoutException() + +def set_timeout(timeout): + signal.signal(signal.SIGALRM, timeout_handler) + signal.alarm(timeout) + +def unset_timeout(): + signal.alarm(0) + signal.signal(signal.SIGALRM, signal.SIG_DFL) + +def make_padding_graph(batch, pad_nodes, pad_edges): + senders = [] + receivers = [] + senders = torch.arange(0,pad_edges) // pad_nodes + receivers = torch.arange(1,pad_edges+1) % pad_nodes + if pad_nodes < 0 or pad_edges < 0 or pad_edges > pad_nodes * pad_nodes / 2: + print('Batch is larger than padding size or e > n^2/2. Repeating edges as necessary.') + print(f'Batch nodes: {batch.num_nodes()}, Batch edges: {batch.num_edges()}, Padding nodes: {pad_nodes}, Padding edges: {pad_edges}') + senders = senders % pad_nodes + padg = dgl.graph((senders[:pad_edges], receivers[:pad_edges]), num_nodes = pad_nodes) + for k in batch.ndata.keys(): + padg.ndata[k] = torch.zeros( (pad_nodes, batch.ndata[k].shape[1]) ) + for k in batch.edata.keys(): + padg.edata[k] = torch.zeros( (pad_edges, batch.edata[k].shape[1]) ) + return dgl.batch([batch, padg.to(batch.device)]) + +def pad_size(graphs, edges, nodes, edge_per_graph=3, node_per_graph=14): + pad_nodes = ((nodes // (node_per_graph * graphs))+1) * graphs * node_per_graph + pad_edges = ((edges // (edge_per_graph * graphs))+1) * graphs * edge_per_graph + return pad_nodes, pad_edges + +def pad_batch_to_step_per_graph(batch, edge_per_graph=3, node_per_graph=14): + n_graphs = batch.batch_num_nodes().shape[0] + pad_nodes = (batch.num_nodes() + node_per_graph * n_graphs) % int(n_graphs * node_per_graph) + pad_edges = (batch.num_edges() + edge_per_graph * n_graphs) % int(n_graphs * edge_per_graph) + return make_padding_graph(batch, pad_nodes, pad_edges) + +def pad_batch(batch, edges = 104000, nodes = 16000): + if edges == 0 and nodes == 0: + return batch + pad_nodes = 0 + pad_edges = 0 + pad_nodes = nodes - batch.num_nodes() + pad_edges = edges - batch.num_edges() + return make_padding_graph(batch, pad_nodes, pad_edges) + +def pad_batch_num_nodes(batch, max_num_nodes, hid_size = 64): + print(f"Padding each graph to have {max_num_nodes} nodes") + + unbatched = dgl.unbatch(batch) + for g in unbatched: + num_nodes_to_add = max_num_nodes - g.number_of_nodes() + if num_nodes_to_add > 0: + g.add_nodes(num_nodes_to_add) # Add isolated nodes + + batch = dgl.batch(unbatched) + + padding_mask = torch.zeros((batch.ndata['features'].shape[0]), dtype=torch.bool) + global_update_weights = torch.ones((batch.ndata['features'].shape[0], hid_size)) + + for i in range(len(batch.ndata['features'])): + if (torch.count_nonzero(batch.ndata['features'][i]) == 0): + padding_mask[i] = True + global_update_weights[i] = 0 + + batch.ndata['w'] = global_update_weights + batch.ndata['padding_mask'] = padding_mask + + return batch + + +def fold_selection(fold_config, sample): + n_folds = fold_config['n_folds'] + folds_opt = fold_config[sample] + folds = [] + if type(folds_opt) == int: + return lambda x : x.tracking[:,0] % n_folds == folds_opt + elif type(folds_opt) == list: + print("fold type is list") + print(f"fold_config = {fold_config}") + print(f"folds_opt = {folds_opt}") + return lambda x : sum([x.tracking[:,0] % n_folds == f for f in folds_opt]) == 1 + else: + raise ValueError("Invalid fold selection option with type {}".format(type(folds_opt))) + +def fold_selection_name(fold_config, sample): + n_folds = fold_config['n_folds'] + folds_opt = fold_config[sample] + if type(folds_opt) == int: + return f'n_{n_folds}_f_{folds_opt}' + elif type(folds_opt) == list: + return f'n_{n_folds}_f_{"_".join([str(f) for f in folds_opt])}' + else: + raise ValueError("Invalid fold selection option with type {}".format(type(folds_opt))) + +#Return the index and checkpoint of the last epoch. +def get_last_epoch(config, max_ep = -1, device = None): + last_epoch = -1 + checkpoint = None + if max_ep < 0: + max_ep = config['Training']['epochs'] + for ep in range(max_ep): + if os.path.exists(os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt')): + last_epoch = ep + else: + print(f'Epoch {ep} not found. Stopping at epoch {last_epoch}') + print('File not found: ', os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt')) + break + if last_epoch >= 0: + checkpoint = torch.load(os.path.join(config['Training_Directory'], f'model_epoch_{last_epoch}.pt'), map_location=device) + return last_epoch, checkpoint + +#Return the index and checkpoint of the last epoch. +def get_specific_epoch(config, target_epoch, device = None, from_ryan = False): + last_epoch = -1 + checkpoint = None + for ep in range(target_epoch + 1): + if (from_ryan): + if os.path.exists(os.path.join('/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/' + config['Training_Directory'], f'model_epoch_{ep}.pt')): + last_epoch = ep + else: + print(f'Epoch {ep} not found. Stopping at epoch {last_epoch}') + print('File not found: ', os.path.join('/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/' + config['Training_Directory'], f'model_epoch_{ep}.pt')) + break + else: + if os.path.exists(os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt')): + last_epoch = ep + else: + print(f'Epoch {ep} not found. Stopping at epoch {last_epoch}') + print('File not found: ', os.path.join(config['Training_Directory'], f'model_epoch_{ep}.pt')) + break + if last_epoch >= 0: + if (from_ryan): + checkpoint = torch.load('/global/cfs/cdirs/atlas/berobert/root_gnn_dgl/' + os.path.join(config['Training_Directory'], f'model_epoch_{last_epoch}.pt'), map_location=device) + else: + checkpoint = torch.load(os.path.join(config['Training_Directory'], f'model_epoch_{last_epoch}.pt'), map_location=device) + return last_epoch, checkpoint + +#Convert training logs into dict for plotting. +def read_log(config): + lines = [] + with open(config['Training_Directory'] + '/training.log', 'r') as f: + lines = f.readlines() + lines = [ l for l in lines if 'Epoch' in l ] + nlines = len(lines) + labels = [] + for field in lines[0].split('|'): + labels.append(field.split()[0]) + log = {label : np.zeros(nlines) for label in labels} + for i, line in enumerate(lines): + for field in line.split('|'): + spl = field.split() + log[spl[0]][i] = float(spl[1]) + return log + +#Plot training logs. +def plot_log(log, output_file): + fig, ax = plt.subplots(2, 2, figsize=(10,10)) + #Time + + ax[0][0].plot(log['Epoch'], np.cumsum(log['Time']), label='Time') + ax[0][0].set_xlabel('Epoch') + ax[0][0].set_ylabel('Time (s)') + ax[0][0].legend() + + """ + ax[0][0].plot(log['Epoch'], log['LR'], label='Learning Rate') + ax[0][0].set_xlabel('Epoch') + ax[0][0].set_ylabel('Learning Rate') + ax[0][0].set_yscale('log') + ax[0][0].legend() + """ + + #Loss + ax[0][1].plot(log['Epoch'], log['Loss'], label='Train Loss') + ax[0][1].plot(log['Epoch'], log['Test_Loss'], label='Test Loss') + ax[0][1].set_xlabel('Epoch') + ax[0][1].set_ylabel('Loss') + ax[0][1].legend() + + #Accuracy + ax[1][0].plot(log['Epoch'], log['Accuracy'], label='Test Accuracy') + ax[1][0].set_xlabel('Epoch') + ax[1][0].set_ylabel('Accuracy') + ax[1][0].set_ylim((0.44, 0.56)) + ax[1][0].legend() + + #AUC + ax[1][1].plot(log['Epoch'], log['Test_AUC'], label='Test AUC') + ax[1][1].set_xlabel('Epoch') + ax[1][1].set_ylabel('AUC') + ax[1][1].legend() + + fig.savefig(output_file) + +class EarlyStop(): + def __init__(self, patience=15, threshold=1e-8, mode='min'): + self.patience = patience + self.threshold = threshold + self.mode = mode + self.count = 0 + self.current_best = np.inf if mode == 'min' else -np.inf + self.should_stop = False + + def update(self, value): + if self.mode == 'min': # Minimizing loss + if value < self.current_best - self.threshold: + self.current_best = value + self.count = 0 + else: + self.count += 1 + elif self.mode == 'max': # Maximizing metric + if value > self.current_best + self.threshold: + self.current_best = value + self.count = 0 + else: + self.count += 1 + + # Check if patience is exceeded + if self.count >= self.patience: + self.should_stop = True + + def reset(self): + self.count = 0 + self.current_best = np.inf if self.mode == 'min' else -np.inf + self.should_stop = False + + def to_str(self): + status = ( + f"EarlyStop Status:\n" + f" Mode: {'Minimize' if self.mode == 'min' else 'Maximize'}\n" + f" Patience: {self.patience}\n" + f" Threshold: {self.threshold:.3e}\n" + f" Current Best: {self.current_best:.6f}\n" + f" Consecutive Epochs Without Improvement: {self.count}\n" + f" Stopping Triggered: {'Yes' if self.should_stop else 'No'}" + ) + return status + + def to_dict(self): + + return { + 'patience': self.patience, + 'threshold': self.threshold, + 'mode': self.mode, + 'count': self.count, + 'current_best': self.current_best, + 'should_stop': self.should_stop, + } + + @classmethod + def load_from_dict(cls, state_dict): + instance = cls( + patience=state_dict['patience'], + threshold=state_dict['threshold'], + mode=state_dict['mode'] + ) + instance.count = state_dict['count'] + instance.current_best = state_dict['current_best'] + instance.should_stop = state_dict['should_stop'] + return instance + + +def graph_augmentation(graph): + print("Augmenting Graph") + return \ No newline at end of file diff --git a/root_gnn_dgl/run_demo.sh b/root_gnn_dgl/run_demo.sh new file mode 100644 index 0000000000000000000000000000000000000000..c17c36add96523aaef1682922de611e9504542ec --- /dev/null +++ b/root_gnn_dgl/run_demo.sh @@ -0,0 +1,59 @@ +#!/bin/bash + +# Pretraining +datasets=("ttH" "tHjb" "ggF" "VBF" "WH" "ZH" "ttyy" "tttt" "SingleT_schan" "ttbar" "ttW" "ttt") +chunks=3 + +for data in "${datasets[@]}"; do + python scripts/prep_data.py --config configs/stats_100K/pretraining_multiclass.yaml --dataset "$data" --shuffle_mode --chunk 0 + for ((i=0; i 1: + chunked_dest = f"{base_dest}_chunk{args.chunkno}" + else: + chunked_dest = base_dest + chunked_dest += '.root' if args.write else '.npz' + args.destination = chunked_dest + + # --- FILE EXISTENCE CHECK --- + if os.path.exists(args.destination): + print(f'File {args.destination} already exists.') + if args.clobber: + print('Clobbering.') + else: + print('Exiting.') + return + else: + print(f'Writing to {args.destination}') + + import time + start = time.time() + import ROOT + import torch + from array import array + import numpy as np + from root_gnn_base import batched_dataset as dataset + from root_gnn_base import utils + end = time.time() + print('Imports finished in {:.2f} seconds'.format(end - start)) + + start = time.time() + dset_config = config['Datasets'][list(config['Datasets'].keys())[0]] + if dset_config['class'] == 'LazyDataset': + dset_config['class'] = 'EdgeDataset' + elif dset_config['class'] == 'LazyMultiLabelDataset': + dset_config['class'] = 'MultiLabelDataset' + elif dset_config['class'] == 'PhotonIDDataset': + dset_config['class'] = 'UnlazyPhotonIDDataset' + elif dset_config['class'] == 'kNNDataset': + dset_config['class'] = 'UnlazyKNNDataset' + dset_config['args']['raw_dir'] = os.path.split(args.target)[0] + dset_config['args']['file_names'] = os.path.split(args.target)[1] + dset_config['args']['save'] = False + dset_config['args']['chunks'] = args.chunks + dset_config['args']['process_chunks'] = [args.chunkno,] + dset_config['args']['selections'] = [] + + dset_config['args']['save_dir'] = os.path.dirname(args.destination) + + if args.tree != '': + dset_config['args']['tree_name'] = args.tree + + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + dstart = time.time() + dset = utils.buildFromConfig(dset_config) + dend = time.time() + print('Dataset finished in {:.2f} seconds'.format(dend - dstart)) + + print(dset) + + batch_size = config['Training']['batch_size'] + lstart = time.time() + loader = CustomPreBatchedDataset( + dset, + batch_size, + chunkno=args.chunkno, + chunks=args.chunks + ) + loader.process() + lend = time.time() + print('Loader finished in {:.2f} seconds'.format(lend - lstart)) + sample_graph, _, _, global_sample = loader[0] + + print('dset length =', len(dset)) + print('loader length =', len(loader)) + + all_scores = {} + all_labels = {} + all_tracking = {} + with torch.no_grad(): + for config_file, branch in zip(args.config, args.branch_name): + config = load_config(config_file) + model = utils.buildFromConfig(config['Model'], {'sample_graph' : sample_graph, 'sample_global': global_sample}).to(device) + if args.ckpt < 0: + ep, checkpoint = utils.get_best_epoch(config, var=args.var, mode='max', device=device) + else: + ep, checkpoint = utils.get_specific_epoch(config, args.ckpt, device=device) + # Remove distributed/compiled prefixes if present + mds_copy = {} + for key in checkpoint['model_state_dict'].keys(): + newkey = key.replace('module.', '') + newkey = newkey.replace('_orig_mod.', '') + mds_copy[newkey] = checkpoint['model_state_dict'][key] + model.load_state_dict(mds_copy) + model.eval() + + end = time.time() + print('Model and dataset finished in {:.2f} seconds'.format(end - start)) + print('Starting inference') + start = time.time() + + finish_fn = torch.nn.Sigmoid() + if 'Loss' in config: + finish_fn = utils.buildFromConfig(config['Loss']['finish']) + + scores = [] + labels = [] + tracking_info = [] + ibatch = 0 + + for batch, label, track, globals in loader.dataloader: + batch = batch.to(device) + pred = model(batch, globals.to(device)) + ibatch += 1 + if (finish_fn.__class__.__name__ == "ContrastiveClusterFinish"): + scores.append(pred.detach().cpu().numpy()) + else: + scores.append(finish_fn(pred).detach().cpu().numpy()) + labels.append(label.detach().cpu().numpy()) + tracking_info.append(track.detach().cpu().numpy()) + + score_size = scores[0].shape[1] if len(scores[0].shape) > 1 else 1 + scores = np.concatenate(scores) + labels = np.concatenate(labels) + tracking_info = np.concatenate(tracking_info) + end = time.time() + + print('Inference finished in {:.2f} seconds'.format(end - start)) + all_scores[branch] = scores + all_labels[branch] = labels + all_tracking[branch] = tracking_info + + if args.write: + from ROOT import std + # Open the original ROOT file + infile = ROOT.TFile.Open(args.target) + tree = infile.Get(dset_config['args']['tree_name']) + + # Create the destination directory if it doesn't exist + os.makedirs(os.path.split(args.destination)[0], exist_ok=True) + + # Create a new ROOT file to write the modified tree + outfile = ROOT.TFile.Open(args.destination, 'RECREATE') + + # Clone the original tree structure + outtree = tree.CloneTree(0) + + # Create branches for all scores + branch_vectors = {} + for branch, scores in all_scores.items(): + if isinstance(scores[0], (list, tuple, np.ndarray)) and len(scores[0]) > 1: + # Create a new branch for vectors + branch_vectors[branch] = std.vector('float')() + outtree.Branch(branch, branch_vectors[branch]) + else: + # Create a new branch for single floats + branch_vectors[branch] = array('f', [0]) + outtree.Branch(branch, branch_vectors[branch], f'{branch}/F') + + # Fill the tree + for i in range(tree.GetEntries()): + tree.GetEntry(i) + + for branch, scores in all_scores.items(): + branch_data = branch_vectors[branch] + if isinstance(branch_data, array): # Check if it's a single float array + branch_data[0] = float(scores[i]) + else: # Assume it's a std::vector + branch_data.clear() + for value in scores[i]: + branch_data.push_back(float(value)) + + outtree.Fill() + + # Write the modified tree to the new file + print(f'Writing to file {args.destination}') + print(f'Input entries: {tree.GetEntries()}, Output entries: {outtree.GetEntries()}') + print(f'Wrote scores to {args.branch_name}') + outtree.Write() + outfile.Close() + infile.Close() + else: + os.makedirs(os.path.split(args.destination)[0], exist_ok=True) + np.savez(args.destination, scores=all_scores, labels=all_labels, tracking_info=all_tracking) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/root_gnn_dgl/scripts/prep_data.py b/root_gnn_dgl/scripts/prep_data.py new file mode 100644 index 0000000000000000000000000000000000000000..87306c834cc1e84e115a85b52e609c2deb956ec7 --- /dev/null +++ b/root_gnn_dgl/scripts/prep_data.py @@ -0,0 +1,44 @@ +import sys +import os +file_path = os.getcwd() +sys.path.append(file_path) + +import root_gnn_base.utils as utils +import argparse +from root_gnn_base.batched_dataset import PreBatchedDataset +from root_gnn_base.batched_dataset import LazyPreBatchedDataset + +def main(): + parser = argparse.ArgumentParser() + add_arg = parser.add_argument + add_arg('--config', type=str, required=True) + add_arg('--dataset', type=str, required=True) + add_arg('--chunk', type=int, default=0) + add_arg('--shuffle_mode', action='store_true', help='Shuffle the dataset before training.') + add_arg('--drop_last', action='store_false', help='Set drop_last to False if the flag is provided. Defaults to True.') + args = parser.parse_args() + + config = utils.load_config(args.config) + dset_config = config['Datasets'][args.dataset] + batch_size = config['Training']['batch_size'] + if not args.shuffle_mode: + dset = utils.buildFromConfig(dset_config, {'process_chunks': [args.chunk,]}) + else: + dset = utils.buildFromConfig(dset_config) + if 'batch_size' in dset_config: + batch_size = dset_config['batch_size'] + + shuffle_chunks = dset_config.get('shuffle_chunks', 10) + padding_mode = dset_config.get('padding_mode', 'STEPS') + fold_conf = dset_config["folding"] + print(f"shuffle_chunks = {shuffle_chunks}, args.chunk = {args.chunk}, padding_mode = {padding_mode}") + if dset_config["class"] == "LazyMultiLabelDataset": + LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last) + LazyPreBatchedDataset(start_dataset = dset, batch_size = batch_size, mask_fn = utils.fold_selection(fold_conf, "test"), suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last) + + else: + PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "train"), suffix = utils.fold_selection_name(fold_conf, "train"), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last) + PreBatchedDataset(dset, batch_size, utils.fold_selection(fold_conf, "test"), suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, chunkno = args.chunk, padding_mode = padding_mode, drop_last=args.drop_last) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/root_gnn_dgl/scripts/training_script.py b/root_gnn_dgl/scripts/training_script.py new file mode 100644 index 0000000000000000000000000000000000000000..ae733e860ba0459c470890b6be006bfe7829fe46 --- /dev/null +++ b/root_gnn_dgl/scripts/training_script.py @@ -0,0 +1,813 @@ +import argparse +import time +import datetime +import yaml +import os + +start_time = time.time() + +import dgl +import torch +import torch.nn as nn + +import sys +file_path = os.getcwd() +sys.path.append(file_path) +import root_gnn_base.batched_dataset as datasets +from root_gnn_base import utils +import root_gnn_base.custom_scheduler as lr_utils +from models import GCN + +import numpy as np +from sklearn.metrics import roc_auc_score +import resource +import gc + +import torch.distributed as dist +import torch.multiprocessing as mp +from torch.utils.data.distributed import DistributedSampler +from torch.nn.parallel import DistributedDataParallel as DDP + +print("import time: {:.4f} s".format(time.time() - start_time)) + +def mem(): + print(f'Current memory usage: {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024 / 1024} GB') + +def gpu_mem(): + print() + print('GPU Memory Usage:') + sum = 0 + # for obj in gc.get_objects(): + # try: + # if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)): + # print(obj.numel() if len(obj.size()) > 0 else 0, type(obj), obj.size()) + # sum += obj.numel() if len(obj.size()) > 0 else 0 + # except: + # pass + print(f'Current GPU memory usage: {torch.cuda.memory_allocated() / 1024 / 1024 / 1024} GB') + print(f'Current GPU cache usage: {torch.cuda.memory_cached() / 1024 / 1024 / 1024} GB') + print(f'Current GPU max memory usage: {torch.cuda.max_memory_allocated() / 1024 / 1024 / 1024} GB') + print(f'Current GPU max cache usage: {torch.cuda.max_memory_cached() / 1024 / 1024 / 1024} GB') + print(f'Numel in current tensors: {sum}') + mem() + + +## epoch stores the epoch number I want to evaluate the model at +def evaluate(val_loaders, model, config, device, epoch = -1): + print("Evaluating") + + if (epoch != -1) : + print(f"Evalulating at epoch {epoch}") + last_ep, checkpoint = utils.get_specific_epoch(config, epoch, from_ryan=False) + print(f"Evaluating at epoch = {last_ep}") + else: + starting_epoch = 0 + last_ep, checkpoint = utils.get_last_epoch(config) + + if checkpoint != None: + ep = last_ep + state_dict = checkpoint['model_state_dict'] + new_state_dict = {} + for k, v in state_dict.items(): + new_key = k.replace('module.', '') + new_state_dict[new_key] = v + model.load_state_dict(new_state_dict) + starting_epoch = checkpoint['epoch'] + 1 + print(f"Loaded epoch {checkpoint['epoch']} from checkpoint") + + if 'Loss' not in config: + loss_fcn = nn.BCEWithLogitsLoss(reduction='none') + else: + loss_fcn = utils.buildFromConfig(config['Loss'], {'reduction': 'none'}) + if len(val_loaders) == 0: + return "No validation data" + start = time.time() + scores = [] + labels = [] + weights = [] + before_decoder = [] + after_decoder = [] + tracking = [] + + batch_size = config["Training"]["batch_size"] + + batch_limit = int(np.ceil(1e5 / batch_size)) + + model.eval() + with torch.no_grad(): + for loader in val_loaders: + batch_count = 0 + for batch, label, track, global_feats in loader: + #Don't use compiled model for testing since we can't control the batch size. + #We could before, but it assumes each dataset has the same number of batches... + before_global_decoder, after_global_decoder, after_classify = model.representation(batch.to(device), global_feats.to(device)) + + scores.append(after_classify.to("cpu")) + before_decoder.append(before_global_decoder.to("cpu")) + after_decoder.append(after_global_decoder.to("cpu")) + labels.append(label.to("cpu")) + weights.append(track[:,1].to("cpu")) + tracking.append(track.to("cpu")) + + batch_count += 1 + if batch_count >= batch_limit: + break + + if scores == []: #If validation set is empty. + return + logits = torch.concatenate(scores) + scores = torch.sigmoid(logits) + labels = torch.concatenate(labels) + weights = torch.concatenate(weights) + before_decoder = torch.concatenate(before_decoder) + after_decoder = torch.concatenate(after_decoder) + tracking = torch.concatenate(tracking) + + logits = logits.to("cpu").numpy() + scores = scores.to("cpu").numpy() + labels = labels.to("cpu").numpy() + before_decoder = before_decoder.to("cpu").numpy() + after_decoder = after_decoder.to("cpu").numpy() + tracking = tracking.to("cpu").numpy() + + # Save the NumPy arrays to a .npz file + outfile = f"{config['Training_Directory']}/evaluation_{epoch}.npz" + + np.savez(outfile, logits=logits, scores=scores, labels=labels, before_decoder=before_decoder, after_decoder=after_decoder, tracking=tracking) + + print(f"saved scores to {outfile}") + return + + +def train(train_loaders, test_loaders, model, device, config, args, rank): + nocompile = args.nocompile + restart = args.restart + # define train/val samples, loss function and optimizer + if 'Loss' not in config: + loss_fcn = nn.BCEWithLogitsLoss(reduction='none') + finish_fn = torch.nn.Sigmoid() + else: + loss_fcn = utils.buildFromConfig(config['Loss'], {'reduction':'none'}) + finish_fn = utils.buildFromConfig(config['Loss']['finish']) + + optimizer = torch.optim.Adam(model.parameters(), lr=config['Training']['learning_rate']) + if 'gamma' in config['Training']: + gamma = config['Training']['gamma'] + else: + gamma = 1 + + if 'dynamic_lr' in config['Training']: + factor = config['Training']['dynamic_lr']['factor'] + patience = config['Training']['dynamic_lr']['patience'] + else: + factor = 1 + patience = 1 + + early_termination = utils.EarlyStop() + if 'early_termination' in config['Training']: + early_termination.patience = config['Training']['early_termination']['patience'] + early_termination.threshold = config['Training']['early_termination']['threshold'] + early_termination.mode = config['Training']['early_termination']['mode'] + + scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma = gamma) + #scheduler_reset = custom_scheduler.Dynamic_LR(optimizer, 'max', factor = factor, patience = patience) + custom_scheduler = None + if ('custom_scheduler' in config['Training']): + run_time_args = {} + scheduler_class = config['Training']['custom_scheduler']['class'] + if (scheduler_class == 'Dynamic_LR' or + scheduler_class == 'Dynamic_LR_AND_Partial_Reset' or + scheduler_class == 'Dynamic_LR_AND_Full_Reset'): + + run_time_args={'optimizer': optimizer} + + custom_scheduler = utils.buildFromConfig(config['Training']['custom_scheduler'], run_time_args=run_time_args) + + starting_epoch = 0 + if not restart: + last_ep, checkpoint = utils.get_last_epoch(config) + if checkpoint != None: + ep = starting_epoch - 1 + if nocompile: + new_state_dict = {} + for k, v in checkpoint['model_state_dict'].items(): + new_key = k.replace('module.', '') + new_state_dict[new_key] = v + checkpoint['model_state_dict'] = new_state_dict + if (args.multinode or args.multigpu): + new_state_dict = {} + for k, v in checkpoint['model_state_dict'].items(): + new_key = 'module.' + k + new_state_dict[new_key] = v + checkpoint['model_state_dict'] = new_state_dict + model.load_state_dict(checkpoint['model_state_dict']) + else: + model._orig_mod.load_state_dict(checkpoint['model_state_dict']) + optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + starting_epoch = checkpoint['epoch'] + 1 + if 'early_stop' in checkpoint: + early_termination = utils.EarlyStop.load_from_dict(checkpoint['early_stop']) + print(early_termination.to_str()) + print("EarlyStop state restored successfully.") + if early_termination.should_stop: + print(f"Early Termination at Epoch {epoch}") + return + else: + print("'early_stop' not found in checkpoint. Initializing a new EarlyStop instance.") + early_termination = utils.EarlyStop() + print(f"Loaded epoch {checkpoint['epoch']} from checkpoint") + log = open(config['Training_Directory'] + '/training.log', 'a', buffering=1) + else: + log = open(config['Training_Directory'] + '/training.log', 'w', buffering=1) + + train_cyclers = [] + for loader in train_loaders: + train_cyclers.append(utils.cycler((loader))) + + if args.savecache: + max_batch = [None,] * len(train_loaders) + for dset_i, loader in enumerate(train_loaders): + mbs = 0 + for batch_i, batch in enumerate(loader): + if batch[0].num_nodes() > mbs: + mbs = batch[0].num_nodes() + max_batch[dset_i] = batch[0] + print(f'Max batch size for dataset {dset_i}: {mbs}') + big_batch = dgl.batch(max_batch).to(device) + with torch.no_grad(): + model(big_batch) + + cumulative_times = [0,0,0,0,0] + log.write(f'Training {config["Training_Name"]} {datetime.datetime.now()} \n') + print(f"Starting training for {config['Training']['epochs']} epochs") + + if hasattr(train_loaders[0].dataset, 'padding_mode'): + is_padded = train_loaders[0].dataset.padding_mode != 'NONE' + if (train_loaders[0].dataset.padding_mode == 'NODE'): + is_padded = False + else: + is_padded = False + + lr_utils.print_LR(optimizer) + + # torch.save({ + # 'epoch': 0, + # 'model_state_dict': model.state_dict(), + # 'optimizer_state_dict': optimizer.state_dict(), + # }, os.path.join(config['Training_Directory'], f"model_epoch_{0}.pt")) + # exit() + + + # training loop + # gpu_mem() + for epoch in range(starting_epoch, config['Training']['epochs']): + start = time.time() + run = start + if (args.multigpu or args.multinode): + dist.barrier() + if (epoch == 2): + # torch.cuda.cudart().cudaProfilerStart() + pass + + # training + model.train() + ibatch = 0 + total_loss = 0 + for batched_graph, labels, _, global_feats in train_loaders[0]: + # # need to fix padded case + # if is_padded: + # tglobals.append(torch.zeros(1, len(global_feats[0]))) + + batch_start = time.time() + logits = torch.tensor([]) + tlabels = torch.tensor([]) + weights = torch.tensor([]) + batch_lengths = [] + for cycler in train_cyclers: + graph, label, track, global_feats = next(cycler) + graph = graph.to(device) + label = label.to(device) + track = track.to(device) + global_feats = global_feats.to(device) + if is_padded: #Padding the globals to match padded graphs. + global_feats = torch.concatenate((global_feats, torch.zeros(1, len(global_feats[0])).to(device))) + load = time.time() + if (len(logits) == 0): + logits = model(graph, global_feats) + tlabels = label + weights = track[:,1] + else: + logits = torch.concatenate((logits, model(graph, global_feats)), dim=0) + tlabels = torch.concatenate((tlabels, label), dim=0) + weights = torch.concatenate((weights, track[:,1]), dim=0) + batch_lengths.append(logits.shape[0] - 1) + + if is_padded: + keepmask = torch.full_like(logits[:,0], True, dtype=torch.bool) + keepmask[batch_lengths] = False + logits = logits[keepmask] + tlabels = tlabels.to(torch.float) + if logits.shape[1] == 1 and loss_fcn.__class__.__name__ == 'BCEWithLogitsLoss': + logits = logits[:,0] + tlabels = tlabels.to(torch.float) + if loss_fcn.__class__.__name__ == 'CrossEntropyLoss': + tlabels = tlabels.to(torch.long) + # loss = loss_fcn(logits, tlabels.to(device)) # changed logits from logits[:,0] and left labels as int for multiclass. Does this break binary? Yes. + # loss = torch.sum(weights * loss) / torch.sum(weights) + + + if args.abs: + weights = torch.abs(weights) + + loss = loss_fcn(logits, tlabels.to(device)) + # Normalize loss within each label + unique_labels = torch.unique(tlabels) # Get unique labels + normalized_loss = 0.0 + + for label in unique_labels: + # Mask for samples belonging to the current label + label_mask = (tlabels == label) + + # Extract weights and losses for the current label + label_weights = weights[label_mask] + label_losses = loss[label_mask] + + + # Compute normalized loss for the current label + label_loss = torch.sum(label_weights * label_losses) / torch.sum(label_weights) + + # Add to the total normalized loss + normalized_loss += label_loss + loss = normalized_loss / len(unique_labels) + + + optimizer.zero_grad() + loss.backward() + optimizer.step() + total_loss += loss.detach().cpu().item() + ibatch += 1 + cumulative_times[0] += batch_start - run + cumulative_times[1] += load - batch_start + run = time.time() + cumulative_times[2] += run - load + if ibatch % 1000 == 0: + print(f'Batch {ibatch} out of {len(train_loaders[0])}', end='\r') + # gpu_mem() + + if (args.multigpu): + print(f'Rank {rank} Epoch Done.') + elif (args.multinode): + print(f'Rank {args.global_rank} Epoch Done.') + else: + print("Epoch Done.") + # validation + + scores = [] + labels = [] + weights = [] + model.eval() + with torch.no_grad(): + for loader in test_loaders: + for batch, label, track, global_feats in loader: + #Don't use compiled model for testing since we can't control the batch size. + #We could before, but it assumes each dataset has the same number of batches... + if is_padded: + global_feats = torch.cat([global_feats, torch.zeros(1, len(global_feats[0]))]) + if nocompile: + batch_scores = model(batch.to(device), global_feats.to(device)) + else: + batch_scores = model._orig_mod(batch.to(device), global_feats.to(device)) + if is_padded: + scores.append(batch_scores[:-1,:]) + else: + scores.append(batch_scores) + labels.append(label) + weights.append(track[:,1]) + eval_end = time.time() + cumulative_times[3] += eval_end - run + + if scores == []: #If validation set is empty. + continue + logits = torch.concatenate(scores).to(device) + labels = torch.concatenate(labels).to(device) + weights = torch.concatenate(weights).to(device) + + if (args.multigpu or args.multinode): + gathered_logits = [torch.zeros_like(logits) for _ in range(dist.get_world_size())] + gathered_labels = [torch.zeros_like(labels) for _ in range(dist.get_world_size())] + gathered_weights = [torch.zeros_like(weights) for _ in range(dist.get_world_size())] + + if (args.multigpu or args.multinode): + dist.barrier() + if (args.multigpu and rank != 0) or (args.multinode and args.global_rank != 0): + dist.gather(logits, dst=0) + dist.gather(labels, dst=0) + dist.gather(weights, dst=0) + continue + else: + dist.gather(logits, gather_list=gathered_logits) + dist.gather(labels, gather_list=gathered_labels) + dist.gather(weights, gather_list=gathered_weights) + + logits = torch.concatenate(gathered_logits) + labels = torch.concatenate(gathered_labels) + weights = torch.concatenate(gathered_weights) + + wgt_mask = weights > 0 + + if args.abs: + weights = torch.abs(weights) + + print(f"Num batches trained = {ibatch}") + + #Note: This section is a bit ugly. Very conditional. Should maybe config defined behavior? + if (loss_fcn.__class__.__name__ == "ContrastiveClusterLoss"): + scores = logits + preds = scores + accuracy = 0 + test_auc = 0 + acc = 0 + contrastive_cluster_loss = finish_fn(logits) + + elif (loss_fcn.__class__.__name__ == "MultiLabelLoss"): + scores = finish_fn(logits) + preds = torch.round(scores) + multilabel_accuracy = [] + threshold = 0.1 # 10% threshold + + for i in range(len(labels[0])): + # accurate_count = torch.sum(torch.abs(preds[:, i].to("cpu") - labels[:, i].to("cpu")) / labels[:, i].to("cpu") <= threshold) + # multilabel_accruacy.append(accurate_count / len(labels)) + multilabel_accuracy.append(torch.sum(preds[:, i].to("cpu") == labels[:, i].to("cpu")) / len(labels)) + test_auc = 0 + acc = np.mean(multilabel_accuracy) + + elif logits.shape[1] == 1 and loss_fcn.__class__.__name__ == 'BCEWithLogitsLoss': #Proxy for binary classification. + test_auc = 0 + acc = 0 + logits = logits[:,0] + scores = finish_fn(logits) + labels =labels.to(torch.float) + preds = scores > 0.5 + test_auc = roc_auc_score(labels[wgt_mask].to("cpu") == 1, scores[wgt_mask].to("cpu"), sample_weight=weights[wgt_mask].to("cpu")) + acc = torch.sum(preds.to("cpu") == labels.to("cpu")) / len(labels) + + elif logits.shape[1] == 1 and loss_fcn.__class__.__name__ == 'MSELoss': + logits = logits[:,0] + scores = finish_fn(logits) + labels = labels.to(torch.float) + acc = 0 + test_auc = 0 + + else: + preds = torch.argmax(logits, dim=1) + scores = finish_fn(logits) + if labels.dim() == 1: #Multi-class + acc = torch.sum(preds.to("cpu") == labels.to("cpu")) / len(labels) #TODO: Make each class weighted equally? + + labels = labels.to("cpu") + weights = weights.to("cpu") + logits = logits.to("cpu") + wgt_mask = wgt_mask.to("cpu") + + labels_onehot = np.zeros((len(labels), len(scores[0]))) + labels_onehot[np.arange(len(labels)), labels] = 1 + + try: + #test_auc = roc_auc_score(labels[wgt_mask].to("cpu") == 1, scores[wgt_mask].to("cpu"), multi_class='ovr', sample_weight=weights[wgt_mask].to("cpu")) + test_auc = roc_auc_score(labels_onehot[wgt_mask], scores[wgt_mask].to("cpu"), multi_class='ovr', sample_weight=weights[wgt_mask].to("cpu")) + except ValueError: + test_auc = np.nan + else: #Multi-loss + acc = torch.sum(preds.to("cpu") == labels[:,0].to("cpu")) / len(labels) + try: + test_auc = roc_auc_score(labels[:,0][wgt_mask].to("cpu") == 1, scores[wgt_mask].to("cpu"), multi_class='ovr', sample_weight=weights[wgt_mask].to("cpu")) + except ValueError: + test_auc = np.nan + + + # print(f"logits = {logits[:10]}") + # print(f"preds = {preds[:2]}") + # print(f"labels = {labels[:10]}") + + # print(f"len(Unique logits) = {len(torch.unique(logits))}") + # print(f"Average of labels = {torch.mean(labels)}") + # print(f"unique logits = {torch.unique(logits)[0]:.4f}, {torch.unique(logits)[-1]:.4f}") + + + if (loss_fcn.__class__.__name__ == "MultiLabelLoss"): + multilabel_log_str = "MultiLabel_Accuracy " + for accuracy in multilabel_accuracy: + multilabel_log_str += f" | {accuracy:.4f}" + log.write(multilabel_log_str + '\n') + print(multilabel_log_str, flush=True) + elif (loss_fcn.__class__.__name__ == "ContrastiveClusterLoss"): + contrastive_cluster_log_str = "ContrastiveClusterLoss " + contrastive_cluster_log_str += f"Contrastive Loss: {contrastive_cluster_loss[0]:.4f}, Clustering Loss: {contrastive_cluster_loss[1]:.4f}, Variance Loss: {contrastive_cluster_loss[2]:.4f}" + log.write(contrastive_cluster_log_str + '\n') + print(contrastive_cluster_log_str, flush=True) + + # test_loss = loss_fcn(logits, labels.to(device)) + # test_loss = loss_fcn(logits, labels) + # test_loss = torch.sum(weights * test_loss) / torch.sum(weights) + + test_loss = loss_fcn(logits, labels) + # Normalize loss within each label + unique_labels = torch.unique(labels) # Get unique labels + normalized_loss = 0.0 + + for label in unique_labels: + # Mask for samples belonging to the current label + label_mask = (labels == label) + + # Extract weights and losses for the current label + label_weights = weights[label_mask] + label_losses = test_loss[label_mask] + # Compute normalized loss for the current label + label_loss = torch.sum(label_weights * label_losses) / torch.sum(label_weights) + + # Add to the total normalized loss + normalized_loss += label_loss + test_loss = normalized_loss / len(unique_labels) + + + end = time.time() + log_str = "Epoch {:05d} | LR {:.4e} | Loss {:.4f} | Accuracy {:.4f} | Test_Loss {:.4f} | Test_AUC {:.4f} | Time {:.4f} s".format( + epoch, optimizer.param_groups[0]['lr'], total_loss/ibatch, acc, test_loss, test_auc, end - start + ) + log.write(log_str + '\n') + print(log_str, flush=True) + + state_dict = model.state_dict() + if not nocompile: + state_dict = model._orig_mod.state_dict() + + new_state_dict = {} + for k, v in state_dict.items(): + new_key = k.replace('module.', '') + new_state_dict[new_key] = v + state_dict = new_state_dict + + # print('Testing done') + # gpu_mem() + + if epoch == 2: + # torch.cuda.cudart().cudaProfilerStop() + pass + + torch.save({ + 'epoch': epoch, + 'model_state_dict': state_dict, + 'optimizer_state_dict': optimizer.state_dict(), + 'early_stop': early_termination.to_dict() + }, os.path.join(config['Training_Directory'], f"model_epoch_{epoch}.pt")) + np.savez(os.path.join(config['Training_Directory'], f'model_epoch_{epoch}.npz'), scores=scores.to("cpu"), labels=labels.to("cpu")) + save_end = time.time() + cumulative_times[4] += save_end - eval_end + + early_termination.update(test_loss) + if early_termination.should_stop: + log_str = f"Early Termination at Epoch {epoch}" + log.write(log_str + "\n") + print(log_str) + log_str = early_termination.to_str() + log.write(log_str + "\n") + print(log_str) + break + + if (custom_scheduler): + custom_scheduler.step(model, {'test_auc':test_auc}) + scheduler.step() + + print(f"Load: {cumulative_times[0]:.4f} s") + print(f"Batch: {cumulative_times[1]:.4f} s") + print(f"Train: {cumulative_times[2]:.4f} s") + print(f"Eval: {cumulative_times[3]:.4f} s") + print(f"Save: {cumulative_times[4]:.4f} s") + log.close() + +def find_free_port(): + import socket + from contextlib import closing + + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s: + s.bind(('', 0)) + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + return str(s.getsockname()[1]) + +def init_process_group(world_size, rank, port): + os.environ['MASTER_ADDR'] = 'localhost' + # os.environ['MASTER_PORT'] = find_free_port() + os.environ['MASTER_PORT'] = port + + dist.init_process_group( + backend="nccl", # change to 'nccl' for multiple GPUs (other was gloo) + init_method='env://', + world_size=world_size, + rank=rank, + timeout=datetime.timedelta(seconds=300), + ) + +def main(rank=0, args=None, world_size=1, port=24500, seed=12345): + + #Prevent simultaneous file access + #sleep_time = 120 * rank + #time.sleep(sleep_time) + + #Load config file + config = utils.load_config(args.config) + + if (args.directory): + print(f"New training directory: { config['Training_Directory'] + args.directory}") + config['Training_Directory'] = config['Training_Directory'] + args.directory + + if not os.path.exists(config['Training_Directory']): + os.makedirs(config['Training_Directory'], exist_ok=True) + with open(config['Training_Directory'] + '/config.yaml', 'w') as f: + yaml.dump(config, f) + batch_size = config["Training"]["batch_size"] + + if(args.plot): + rl = utils.read_log(config) + utils.plot_log(rl, config['Training_Directory'] + '/training.png') + print('Log at ' + config['Training_Directory'] + '/training.log') + print('Plotted at ' + config['Training_Directory'] + '/training.png') + exit() + + if (args.multigpu): + print(f"Setting up multigpu") + start_time = time.time() + init_process_group(world_size, rank, port) + print("multigpu setup time: {:.4f} s".format(time.time() - start_time)) + device = torch.device(f'cuda:{rank}') + torch.cuda.device(device) + elif (args.multinode): + device = torch.device(f'cuda:{rank}') + torch.cuda.device(device) + print(f"global rank = {args.global_rank}, local rank = {rank}, device = {device}") + else: + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + if (args.cpu): + print(f"Using CPU") + device = "cpu" + + train_loaders = [] + test_loaders = [] + val_loaders = [] + load_start = time.time() + + torch.backends.cuda.matmul.allow_tf32 = True + + ldr_type = datasets.LazyPreBatchedDataset if args.lazy else datasets.PreBatchedDataset + + #Load datasets + if (pargs.statistics): + pargs.statistics = int(pargs.statistics) + print(f"Training Dataset Size: {pargs.statistics}") + num_batches = int(np.ceil(pargs.statistics / batch_size)) + np.random.seed(pargs.seed) + + for dset_conf in config["Datasets"]: + dset = utils.buildFromConfig(config["Datasets"][dset_conf]) + if 'batch_size' in config["Datasets"][dset_conf]: + batch_size = config["Datasets"][dset_conf]['batch_size'] + fold_conf = config["Datasets"][dset_conf]["folding"] + shuffle_chunks = config["Datasets"][dset_conf].get("shuffle_chunks", 10) + padding_mode = config["Datasets"][dset_conf].get("padding_mode", "STEPS") + mask_fn = utils.fold_selection(fold_conf, "train") + if args.preshuffle: + # ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'train'), chunks = shuffle_chunks, padding_mode = padding_mode, use_ddp = args.multigpu, rank=rank, world_size=world_size) + ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'train'), chunks = shuffle_chunks, padding_mode = padding_mode) + gsamp, _, _, global_samp = ldr[0] + sampler = None + + if (pargs.statistics): + sampler = np.random.choice(range(len(ldr)), size=num_batches) + + if (args.multigpu): + sampler = DistributedSampler(ldr, num_replicas=world_size, rank=rank, shuffle=False, drop_last=True) + # num_batches = len(ldr) + # sampler = list(sampler) + # if (sampler[0] >= num_batches % world_size): + # sampler.pop() + if (args.multinode): + sampler = DistributedSampler(ldr, num_replicas=world_size, rank=pargs.global_rank, shuffle=False, drop_last=True) + train_loaders.append(torch.utils.data.DataLoader(ldr, batch_size = None, num_workers = 0, sampler = sampler)) + sampler = None + ldr = ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=mask_fn, suffix = utils.fold_selection_name(fold_conf, 'test'), chunks = shuffle_chunks, padding_mode = padding_mode) + if (args.multigpu): + sampler = DistributedSampler(ldr, num_replicas=world_size, rank=rank, shuffle=False, drop_last=True) + # num_batches = len(ldr) + # sampler = list(sampler) + # if (rank >= num_batches % world_size): + # sampler.pop() + if (args.multinode): + sampler = DistributedSampler(ldr, num_replicas=world_size, rank=pargs.global_rank, shuffle=False, drop_last=True) + + test_loaders.append(torch.utils.data.DataLoader(ldr, batch_size = None, num_workers = 0, sampler=sampler)) + + if "validation" in fold_conf: + val_loaders.append(torch.utils.data.DataLoader((ldr_type(start_dataset=dset, batch_size=batch_size, mask_fn=utils.fold_selection(fold_conf, "validation"), suffix = utils.fold_selection_name(fold_conf, 'validation'), chunks = shuffle_chunks, padding_mode = padding_mode, rank=rank, world_size=1)), batch_size = None, num_workers = 0, sampler = sampler)) + else: + print("No validation set for dataset ", dset_conf) + else: + train_loaders.append(datasets.GetBatchedLoader(dset, batch_size, utils.fold_selection(fold_conf, "train"))) + gsamp, _, _, global_samp = dset[0] + test_loaders.append(datasets.GetBatchedLoader(dset, batch_size, utils.fold_selection(fold_conf, "test"))) + if "validation" in fold_conf: + val_loaders.append(datasets.GetBatchedLoader(dset, batch_size, utils.fold_selection(fold_conf, "validation"))) + else: + print("No validation set for dataset ", dset_conf) + + load_end = time.time() + print("Load time: {:.4f} s".format(load_end - load_start)) + + model = utils.buildFromConfig(config["Model"], {'sample_graph': gsamp, 'sample_global': global_samp, 'seed': seed}).to(device) + if not args.nocompile: + model = torch.compile(model) + if args.multigpu: + print(f"Trying to create DDP model") + start_time = time.time() + model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device]) + print("model creation time: {:.4f} s".format(time.time() - start_time)) + if (args.multinode): + print(f"Trying to create DDP model") + start_time = time.time() + model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[device]) + print("model creation time: {:.4f} s".format(time.time() - start_time)) + + # total_params = 0 + # for param_dict in model.parameters(): + # for param in param_dict['params']: + # if param.requires_grad: + # total_params += param.numel() + # print(f"Number of trainable parameters = {total_params}") + + if(type(model) == GCN.Clustering): + print("clustering") + + if args.evaluate != None: + evaluate(test_loaders, model, config, device, args.evaluate) + exit() + + # model training + print("Training...") + gpu_mem() + train(train_loaders, test_loaders, model, device, config, args, rank) + + # test the model + # print("Testing...") + # evaluate(val_loaders, model, config, device) + + # if args.multigpu or args.multinode: + # dist.destroy_process_group() + + # if rank == 0: + # rl = utils.read_log(config) + # utils.plot_log(rl, config['Training_Directory'] + '/training.png') + # print('Log at ' + config['Training_Directory'] + '/training.log') + # print('Plotted at ' + config['Training_Directory'] + '/training.png') + +if __name__ == "__main__": + #Handle CLI arguments + parser = argparse.ArgumentParser() + add_arg = parser.add_argument + add_arg("--config", type=str, help="Config file.", required=True) + add_arg("--restart", action="store_true", help="Restart training from scratch.") + add_arg("--preshuffle", action="store_true", help="Shuffle data before training.") + add_arg("--lazy", action="store_true", help="Lazy loading of data.") + add_arg("--nocompile", action="store_true", help="Disable JIT compilation.") + add_arg("--evaluate", type = int, help="Skip training and go to evaluation.") + add_arg("--plot", action="store_true", help="Plot training logs.") + add_arg("--multigpu", action="store_true", help="Use multiple GPUs.") + add_arg("--multinode", action="store_true", help="Use multiple nodes.") + add_arg("--savecache", action="store_true", help="") + add_arg("--cpu", action="store_true", help="Uses the cpu only") + add_arg("--statistics", type=float, help="Size of training data") + add_arg("--directory", type=str, help="Append to Training Directory") + add_arg("--seed", type=int, default=2, help="Sets random seed") + add_arg("--abs", action="store_true", help="Use abs value of per-event weight") + + pargs = parser.parse_args() + + if pargs.multigpu: + port = find_free_port() + torch.backends.cudnn.enabled = False + mp.spawn(main, args=(pargs, 4, port), nprocs=4, join=True) + if pargs.multinode: + global_rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + print(f"global_rank = {global_rank}, local_rank = {local_rank}, world_size = {world_size}") + + dist.init_process_group(backend="nccl") + torch.backends.cudnn.enabled = False + + pargs.global_rank = global_rank + + main(rank = local_rank, args=pargs, world_size=world_size) + else: + main(0, pargs) + + diff --git a/root_gnn_dgl/setup/environment.yml b/root_gnn_dgl/setup/environment.yml new file mode 100644 index 0000000000000000000000000000000000000000..d55c25d7b4d0f7b0b2f625c5dc47b4ed1bd47f12 --- /dev/null +++ b/root_gnn_dgl/setup/environment.yml @@ -0,0 +1,391 @@ +name: dgl +channels: + - pytorch + - dglteam/label/cu118 + - nvidia + - conda-forge + - defaults +dependencies: + - _libgcc_mutex=0.1 + - _openmp_mutex=4.5 + - _sysroot_linux-64_curr_repodata_hack=3 + - afterimage=1.21 + - anyio=3.7.1 + - appdirs=1.4.4 + - argon2-cffi=21.3.0 + - argon2-cffi-bindings=21.2.0 + - arrow=1.2.3 + - asttokens=2.2.1 + - async-lru=2.0.4 + - atk-1.0=2.38.0 + - attrs=23.1.0 + - awkward-pandas=2023.8.0 + - aws-c-auth=0.7.0 + - aws-c-cal=0.6.0 + - aws-c-common=0.8.23 + - aws-c-compression=0.2.17 + - aws-c-event-stream=0.3.1 + - aws-c-http=0.7.11 + - aws-c-io=0.13.28 + - aws-c-mqtt=0.8.14 + - aws-c-s3=0.3.13 + - aws-c-sdkutils=0.1.11 + - aws-checksums=0.1.16 + - aws-crt-cpp=0.20.3 + - aws-sdk-cpp=1.10.57 + - babel=2.12.1 + - backcall=0.2.0 + - backports=1.0 + - backports.functools_lru_cache=1.6.5 + - beautifulsoup4=4.12.2 + - binutils=2.38 + - binutils_impl_linux-64=2.38 + - binutils_linux-64=2.38.0 + - blas=1.0 + - bleach=6.0.0 + - brotlipy=0.7.0 + - bzip2=1.0.8 + - c-ares=1.19.1 + - c-compiler=1.5.2 + - ca-certificates=2025.4.26 + - cached-property=1.5.2 + - cached_property=1.5.2 + - cairo=1.16.0 + - certifi=2024.8.30 + - cffi=1.15.1 + - cfitsio=4.2.0 + - charset-normalizer=2.0.4 + - comm=0.1.4 + - compilers=1.5.2 + - cryptography=41.0.2 + - cuda-cudart=11.8.89 + - cuda-cupti=11.8.87 + - cuda-libraries=11.8.0 + - cuda-nvrtc=11.8.89 + - cuda-nvtx=11.8.86 + - cuda-runtime=11.8.0 + - cxx-compiler=1.5.2 + - davix=0.8.4 + - debugpy=1.6.8 + - decorator=5.1.1 + - defusedxml=0.7.1 + - dgl=1.1.1.cu118 + - entrypoints=0.4 + - exceptiongroup=1.1.3 + - executing=1.2.0 + - expat=2.5.0 + - ffmpeg=4.3 + - fftw=3.3.10 + - filelock=3.9.0 + - flit-core=3.9.0 + - font-ttf-dejavu-sans-mono=2.37 + - font-ttf-inconsolata=3.000 + - font-ttf-source-code-pro=2.038 + - font-ttf-ubuntu=0.83 + - fontconfig=2.14.2 + - fonts-conda-ecosystem=1 + - fonts-conda-forge=1 + - fortran-compiler=1.5.2 + - fqdn=1.5.1 + - freetype=2.12.1 + - fribidi=1.0.10 + - ftgl=2.4.0 + - gcc=11.2.0 + - gcc_impl_linux-64=11.2.0 + - gcc_linux-64=11.2.0 + - gdk-pixbuf=2.42.8 + - gettext=0.21.1 + - gflags=2.2.2 + - gfortran=11.2.0 + - gfortran_impl_linux-64=11.2.0 + - gfortran_linux-64=11.2.0 + - giflib=5.2.1 + - gl2ps=1.4.2 + - glew=2.1.0 + - glog=0.6.0 + - gmp=6.2.1 + - gmpy2=2.1.2 + - gnutls=3.6.15 + - graphite2=1.3.13 + - graphviz=6.0.2 + - gsl=2.7 + - gsoap=2.8.123 + - gtk2=2.24.33 + - gts=0.7.6 + - gxx=11.2.0 + - gxx_impl_linux-64=11.2.0 + - gxx_linux-64=11.2.0 + - harfbuzz=7.3.0 + - icu=72.1 + - idna=3.4 + - importlib-metadata=6.8.0 + - importlib-resources=6.0.1 + - importlib_metadata=6.8.0 + - importlib_resources=6.0.1 + - intel-openmp=2023.1.0 + - ipykernel=6.25.1 + - ipyparallel=8.6.1 + - ipython=8.12.2 + - isoduration=20.11.0 + - jedi=0.19.0 + - jinja2=3.1.2 + - jpeg=9e + - json5=0.9.14 + - jsonpointer=2.0 + - jsonschema=4.19.0 + - jsonschema-specifications=2023.7.1 + - jsonschema-with-format-nongpl=4.19.0 + - jupyter-lsp=2.2.0 + - jupyter_client=8.3.0 + - jupyter_core=5.3.0 + - jupyter_events=0.7.0 + - jupyter_server=2.7.0 + - jupyter_server_terminals=0.4.4 + - jupyterlab=4.0.5 + - jupyterlab_pygments=0.2.2 + - jupyterlab_server=2.24.0 + - kernel-headers_linux-64=3.10.0 + - keyutils=1.6.1 + - krb5=1.20.1 + - lame=3.100 + - lcms2=2.12 + - ld_impl_linux-64=2.38 + - lerc=3.0 + - libabseil=20230125.3 + - libarrow=12.0.1 + - libblas=3.9.0 + - libbrotlicommon=1.0.9 + - libbrotlidec=1.0.9 + - libbrotlienc=1.0.9 + - libcblas=3.9.0 + - libcrc32c=1.1.2 + - libcublas=11.11.3.6 + - libcufft=10.9.0.58 + - libcufile=1.7.1.12 + - libcurand=10.3.3.129 + - libcurl=8.1.2 + - libcusolver=11.4.1.48 + - libcusparse=11.7.5.86 + - libcxx=15.0.7 + - libcxxabi=15.0.7 + - libdeflate=1.12 + - libedit=3.1.20191231 + - libev=4.33 + - libevent=2.1.12 + - libexpat=2.5.0 + - libffi=3.4.4 + - libgcc-devel_linux-64=11.2.0 + - libgcc-ng=13.1.0 + - libgd=2.3.3 + - libgfortran-ng=11.2.0 + - libgfortran5=11.2.0 + - libglib=2.76.4 + - libglu=9.0.0 + - libgomp=13.1.0 + - libgoogle-cloud=2.12.0 + - libgrpc=1.56.2 + - libiconv=1.17 + - libidn2=2.3.4 + - libllvm13=13.0.1 + - libllvm14=14.0.6 + - libnghttp2=1.52.0 + - libnpp=11.8.0.86 + - libnsl=2.0.0 + - libnuma=2.0.18 + - libnvjpeg=11.9.0.86 + - libpng=1.6.39 + - libprotobuf=4.23.3 + - librsvg=2.54.4 + - libsodium=1.0.18 + - libsqlite=3.42.0 + - libssh2=1.11.0 + - libstdcxx-devel_linux-64=11.2.0 + - libstdcxx-ng=13.1.0 + - libtasn1=4.19.0 + - libthrift=0.18.1 + - libtiff=4.4.0 + - libtool=2.4.7 + - libunistring=0.9.10 + - libutf8proc=2.8.0 + - libuuid=2.38.1 + - libwebp=1.2.4 + - libwebp-base=1.2.4 + - libxcb=1.15 + - libxml2=2.10.4 + - libzlib=1.2.13 + - llvmlite=0.40.1 + - lz4-c=1.9.4 + - markupsafe=2.1.1 + - matplotlib-inline=0.1.6 + - metakernel=0.29.5 + - mistune=3.0.0 + - mkl=2023.1.0 + - mkl-service=2.4.0 + - mkl_fft=1.3.6 + - mkl_random=1.2.2 + - mpc=1.1.0 + - mpfr=4.0.2 + - mpmath=1.3.0 + - nbclient=0.8.0 + - nbconvert-core=7.7.3 + - nbformat=5.9.2 + - ncurses=6.4 + - nest-asyncio=1.5.6 + - nettle=3.7.3 + - networkx=3.1 + - nlohmann_json=3.11.2 + - notebook=7.0.2 + - notebook-shim=0.2.3 + - numba=0.57.1 + - numpy=1.24.3 + - numpy-base=1.24.3 + - openh264=2.1.1 + - openssl=3.3.1 + - orc=1.9.0 + - overrides=7.4.0 + - packaging=23.0 + - pandas=2.0.3 + - pandocfilters=1.5.0 + - pango=1.50.14 + - parso=0.8.3 + - pcre=8.45 + - pcre2=10.40 + - pexpect=4.8.0 + - pickleshare=0.7.5 + - pillow=9.4.0 + - pip=23.2.1 + - pixman=0.40.0 + - pkgutil-resolve-name=1.3.10 + - platformdirs=2.6.0 + - pooch=1.4.0 + - portalocker=2.7.0 + - prometheus_client=0.17.1 + - prompt-toolkit=3.0.39 + - prompt_toolkit=3.0.39 + - psutil=5.9.0 + - pthread-stubs=0.4 + - ptyprocess=0.7.0 + - pure_eval=0.2.2 + - pyarrow=12.0.1 + - pycparser=2.21 + - pygments=2.16.1 + - pyopenssl=23.2.0 + - pysocks=1.7.1 + - pythia8=8.309 + - python=3.8.17 + - python-dateutil=2.8.2 + - python-fastjsonschema=2.18.0 + - python-json-logger=2.0.7 + - python-tzdata=2024.2 + - python_abi=3.8 + - pytorch=2.0.1 + - pytorch-cuda=11.8 + - pytorch-mutex=1.0 + - pytz=2023.3 + - pyyaml=6.0 + - pyzmq=25.1.1 + - rdma-core=28.9 + - re2=2023.03.02 + - readline=8.2 + - referencing=0.30.2 + - requests=2.31.0 + - rfc3339-validator=0.1.4 + - rfc3986-validator=0.1.1 + - root=6.28.0 + - root_base=6.28.0 + - rpds-py=0.9.2 + - s2n=1.3.46 + - scipy=1.10.1 + - scitokens-cpp=0.7.3 + - send2trash=1.8.2 + - setuptools=68.0.0 + - six=1.16.0 + - snappy=1.1.10 + - sniffio=1.3.0 + - soupsieve=2.3.2.post1 + - sqlite=3.41.2 + - stack_data=0.6.2 + - sympy=1.11.1 + - sysroot_linux-64=2.17 + - tbb=2021.8.0 + - terminado=0.17.1 + - tinycss2=1.2.1 + - tk=8.6.12 + - tomli=2.0.1 + - torchaudio=2.0.2 + - torchtriton=2.0.0 + - torchvision=0.15.2 + - tornado=6.3.2 + - tqdm=4.65.0 + - traitlets=5.9.0 + - typing_extensions=4.12.2 + - typing_utils=0.1.0 + - ucx=1.14.1 + - uri-template=1.3.0 + - urllib3=1.26.16 + - vdt=0.4.3 + - vector-classes=1.4.3 + - wcwidth=0.2.6 + - webcolors=1.13 + - webencodings=0.5.1 + - websocket-client=1.6.1 + - wheel=0.38.4 + - xorg-fixesproto=5.0 + - xorg-kbproto=1.0.7 + - xorg-libice=1.1.1 + - xorg-libsm=1.2.4 + - xorg-libx11=1.8.6 + - xorg-libxau=1.0.11 + - xorg-libxcursor=1.2.0 + - xorg-libxdmcp=1.1.3 + - xorg-libxext=1.3.4 + - xorg-libxfixes=5.0.3 + - xorg-libxft=2.3.8 + - xorg-libxpm=3.5.16 + - xorg-libxrender=0.9.11 + - xorg-libxt=1.3.0 + - xorg-renderproto=0.11.1 + - xorg-xextproto=7.3.0 + - xorg-xproto=7.0.31 + - xrootd=5.5.4 + - xxhash=0.8.1 + - xz=5.2.6 + - yaml=0.2.5 + - zeromq=4.3.4 + - zipp=3.16.2 + - zlib=1.2.13 + - zstd=1.5.2 + - pip: + - awkward==2.6.4 + - awkward-cpp==33 + - contourpy==1.1.0 + - cramjam==2.8.3 + - cycler==0.11.0 + - fonttools==4.42.0 + - fsspec==2024.3.1 + - h5py==3.9.0 + - pip-install==1.3.5 + - joblib==1.3.2 + - kiwisolver==1.4.4 + - matplotlib==3.7.2 + - nvidia-cublas-cu12==12.1.3.1 + - nvidia-cuda-cupti-cu12==12.1.105 + - nvidia-cuda-nvrtc-cu12==12.1.105 + - nvidia-cuda-runtime-cu12==12.1.105 + - nvidia-cudnn-cu12==8.9.2.26 + - nvidia-cufft-cu12==11.0.2.54 + - nvidia-curand-cu12==10.3.2.106 + - nvidia-cusolver-cu12==11.4.5.107 + - nvidia-cusparse-cu12==12.1.0.106 + - nvidia-nccl-cu12==2.20.5 + - nvidia-nvjitlink-cu12==12.4.127 + - nvidia-nvtx-cu12==12.1.105 + - pyparsing==3.0.9 + - scikit-learn==1.3.0 + - threadpoolctl==3.2.0 + - torch==2.3.0 + - triton==2.3.0 + - typing-extensions==4.11.0 + - tzdata==2024.1 + - uproot==5.3.7 +prefix: /global/homes/j/joshuaho/.conda/envs/dgl diff --git a/root_gnn_dgl/setup/test_setup.py b/root_gnn_dgl/setup/test_setup.py new file mode 100644 index 0000000000000000000000000000000000000000..f41377b0c50f35b55050ca88fa3f129a89b53d05 --- /dev/null +++ b/root_gnn_dgl/setup/test_setup.py @@ -0,0 +1,48 @@ +import os +import importlib.util +import sys + +def test_imports(directories): + """ + Test importing all Python files in the specified directories. + + Parameters: + - directories: List of directory paths to test. + """ + print("Testing Conda environment...") + + for directory in directories: + print(f"\nChecking directory: {directory}") + + # Check if the directory exists + if not os.path.isdir(directory): + print(f"Directory not found: {directory}") + continue + + # Iterate through all files in the directory + for filename in os.listdir(directory): + # Only consider Python files + if filename.endswith(".py"): + filepath = os.path.join(directory, filename) + module_name = os.path.splitext(filename)[0] # Remove .py extension + + try: + # Dynamically import the module + spec = importlib.util.spec_from_file_location(module_name, filepath) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + print(f"Successfully imported: {filepath}") + except Exception as e: + # Print the file and the error message if import fails + print(f"Failed to import: {filepath}") + print(f"Error: {e}") + +if __name__ == "__main__": + # Automatically append the current directory to sys.path + current_directory = os.getcwd() + sys.path.append(current_directory) + print(f"Current directory added to sys.path: {current_directory}") + + # List of directories to check + directories = ["scripts", "root_gnn_base", "models"] + test_imports(directories) \ No newline at end of file diff --git a/scripts/compare_encoders.py b/scripts/compare_encoders.py deleted file mode 100644 index bbd186cfc3289b61a14fb2f30a7ed3693fea7136..0000000000000000000000000000000000000000 --- a/scripts/compare_encoders.py +++ /dev/null @@ -1,138 +0,0 @@ -"""Linear CKA between object encoders, as in Ho et al. (2026): of their event -embeddings, and layer by layer against the first encoder. - -All encoders have the architecture of the config (``model.backbone`` of a -pretraining config, or the transformer or edge network of a classifier's): - - uv run python scripts/compare_encoders.py configs/gamgam_jepa_finetune.yaml \\ - --encoders random=random jepa=outputs/delphes_jepa/last.pt \\ - contrastive=outputs/delphes_contrastive/last.pt \\ - masked=outputs/delphes_masked/last.pt --events 5000 - -``random`` is a randomly initialized encoder (the reference). The events are -``--events`` random events of ``--split``, drawn with ``data.seed``. Writes -``--output`` (JSON; undefined values, of an encoder that gives every event -the same embedding, are null) and prints Markdown tables. -""" - -from __future__ import annotations - -import argparse -import dataclasses -import json -import math -from pathlib import Path - -from gnn4colliders.config import load_config -from gnn4colliders.data import load_events -from gnn4colliders.distributed import resolve_device -from gnn4colliders.probe import layer_embeddings, linear_cka, load_encoder - - -def _table(rows: list[str], columns: list[str], values: list[list[float]]) -> str: - lines = [ - "| | " + " | ".join(columns) + " |", - "|---" * (len(columns) + 1) + "|", - ] - for row, numbers in zip(rows, values): - lines.append(f"| {row} | " + " | ".join(f"{v:.3f}" for v in numbers) + " |") - return "\n".join(lines) - - -def _encoders(parser: argparse.ArgumentParser, entries: list[str]) -> dict[str, str]: - """``label=path`` entries as {label: path}, labels distinct.""" - encoders: dict[str, str] = {} - for entry in entries: - label, separator, path = entry.partition("=") - if not separator or not label or not path: - parser.error( - f"--encoders takes label=path (or label=random), not {entry!r}" - ) - if label in encoders: - parser.error(f"--encoders label {label!r} is given twice") - encoders[label] = path - return encoders - - -def _finite(values: list[float]) -> list[float | None]: - return [None if math.isnan(value) else value for value in values] - - -def main() -> None: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("config", type=Path) - parser.add_argument("--encoders", nargs="+", required=True, help="label=path") - parser.add_argument("--events", type=int, default=5000) - parser.add_argument("--split", default="validation") - parser.add_argument("--override", nargs="*", default=[], help="key=value") - parser.add_argument("--output", type=Path, default=Path("outputs/cka.json")) - args = parser.parse_args() - encoders = _encoders(parser, args.encoders) - config = load_config(args.config, args.override) - events = load_events(config.data.cache) - folds = { - "train": config.data.train_folds, - "validation": config.data.validation_folds, - "test": config.data.test_folds, - }[args.split] - indices = events.split_indices(folds) - indices = events.random_subset(indices, args.events, config.data.seed) - device = resolve_device(config.train.device) - layers = {} - for label, path in encoders.items(): - model = dataclasses.replace( - config.model, pretrained=None if path == "random" else path - ) - encoder_config = dataclasses.replace(config, model=model) - encoder = load_encoder(encoder_config, events).to(device) - layers[label] = layer_embeddings( - encoder, - events, - indices, - batch_size=config.data.batch_size, - device=device, - ) - labels = list(layers) - final = [[linear_cka(layers[a][-1], layers[b][-1]) for b in labels] for a in labels] - reference = labels[0] - layerwise = { - label: [linear_cka(x, y) for x, y in zip(layers[reference], layers[label])] - for label in labels[1:] - } - depth = len(layers[reference]) - args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text( - json.dumps( - { - "events": len(indices), - "split": args.split, - "encoders": encoders, - "embeddings": { - label: _finite(row) for label, row in zip(labels, final) - }, - "layerwise_against": reference, - "layerwise": {label: _finite(row) for label, row in layerwise.items()}, - }, - indent=2, - allow_nan=False, - ) - + "\n" - ) - print(f"Linear CKA of the event embeddings ({len(indices)} {args.split} events):\n") - print(_table(labels, labels, final)) - # An encoder's CKA with itself is undefined only if it does not vary. - undefined = [ - label for index, label in enumerate(labels) if math.isnan(final[index][index]) - ] - if undefined: - print(f"\nnan: {', '.join(undefined)} give every event the same embedding") - if layerwise: - print(f"\nLayer by layer against {reference} (last: the event embedding):\n") - columns = [f"layer {index}" for index in range(depth)] - print(_table(list(layerwise), columns, list(layerwise.values()))) - - -if __name__ == "__main__": - main() diff --git a/scripts/label_efficiency.py b/scripts/label_efficiency.py deleted file mode 100644 index bf63c7fc8e1634b54cb7d4a905ca82350fa710e7..0000000000000000000000000000000000000000 --- a/scripts/label_efficiency.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Validation AUC against the number of labelled training events. - -For every config and every ``--events`` count N, runs the model with -``data.train_events=N`` in ``//n``: supervised models are -trained, then evaluated; in-context models (TabPFN, TabICL, PFN) are only evaluated, -with the N events as context. Runs whose metrics exist are skipped, so an -interrupted sweep restarts where it stopped. Prints a table of AUC (%) and -evaluation seconds. The event cache must exist (``gnn4colliders prepare``). - - uv run python scripts/label_efficiency.py configs/gamgam_tabpfn.yaml \\ - configs/gamgam.yaml --events 100 1000 10000 50000 - -Meta-trained in-context models (the PFN; TabPFN or TabICL after -meta-training) need ``--override checkpoint=``. Overrides apply to -every config; ``name=...`` keeps variants apart (run one config at a time -then, as names must differ): - - uv run python scripts/label_efficiency.py configs/gamgam_tabpfn.yaml \\ - --override name=gamgam_tabpfn_multiclass \\ - model.tabpfn_checkpoint=tabpfn-v3.5-20260909_multiclass.safetensors -""" - -from __future__ import annotations - -import argparse -import json -from pathlib import Path - -from gnn4colliders.cli import main as gnn4colliders -from gnn4colliders.config import load_config -from gnn4colliders.models import IN_CONTEXT_MODELS - - -def _run(config: Path, events: int, out: Path, overrides: list[str]) -> dict: - name = load_config(config, overrides).name - run_dir = out / name / f"n{events}" - metrics = run_dir / "validation_metrics.json" - if not metrics.exists(): - args = [*overrides, f"data.train_events={events}", f"run_dir={run_dir}"] - resolved = load_config(config, args) - if resolved.model.name == "pfn" and resolved.checkpoint is None: - # The run directory is new: the meta-trained weights are elsewhere. - raise SystemExit("the PFN needs --override checkpoint=.pt") - if resolved.model.name not in IN_CONTEXT_MODELS: - if gnn4colliders(["train", str(config), *args]) != 0: - raise SystemExit(f"training {name} with {events} events failed") - if gnn4colliders(["evaluate", str(config), *args, "split=validation"]) != 0: - raise SystemExit(f"evaluating {name} with {events} events failed") - return {"name": name, "events": events, **json.loads(metrics.read_text())} - - -def main() -> None: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("configs", nargs="+", type=Path) - parser.add_argument("--events", nargs="+", type=int, required=True) - parser.add_argument("--override", nargs="*", default=[], help="key=value") - parser.add_argument("--out", type=Path, default=Path("outputs/label_efficiency")) - args = parser.parse_args() - names = [load_config(config, args.override).name for config in args.configs] - if len(set(names)) < len(names): - raise SystemExit(f"configs share a run name ({names}); override name per run") - runs = [ - _run(config, events, args.out, args.override) - for config in args.configs - for events in args.events - ] - print("| model | " + " | ".join(f"N = {n}" for n in args.events) + " |") - print("|---" * (len(args.events) + 1) + "|") - for name in dict.fromkeys(run["name"] for run in runs): - cells = [] - for events in args.events: - run = next(r for r in runs if r["name"] == name and r["events"] == events) - auc = 100 * run["metrics"]["auc"] - seconds = run.get("seconds") - total = f", {sum(seconds.values()):.0f} s" if seconds else "" - cells.append(f"{auc:.2f}{total}") - print(f"| {name} | " + " | ".join(cells) + " |") - - -if __name__ == "__main__": - main() diff --git a/scripts/label_efficiency.sh b/scripts/label_efficiency.sh deleted file mode 100755 index 6a91b1b4706e7a8e1ce9654dfbd92ed0bf13bdeb..0000000000000000000000000000000000000000 --- a/scripts/label_efficiency.sh +++ /dev/null @@ -1,29 +0,0 @@ -#!/bin/bash -# Label-efficiency sweep (scripts/label_efficiency.py) on one GPU: -# sbatch -A --export=ALL,CONFIGS="configs/gamgam_tabpfn.yaml configs/gamgam.yaml" \ -# scripts/label_efficiency.sh [key=value ...] -# EVENTS (default "100 1000 10000 50000") are the labelled training events; -# key=value overrides apply to every config. Restartable: finished runs are -# skipped. TabPFN-3.5 needs its weights in TABPFN_MODEL_CACHE_DIR -# (docs/usage.md). Prepare the event cache first. -#SBATCH -N 1 -#SBATCH -C gpu -#SBATCH -q shared -#SBATCH -G 1 -#SBATCH -c 32 -#SBATCH -t 06:00:00 -#SBATCH -o logs/label_efficiency_%j.out - -set -euo pipefail -cd "${SLURM_SUBMIT_DIR:-$(pwd)}" -read -r -a configs <<< "${CONFIGS:-configs/gamgam_tabpfn.yaml}" -read -r -a events <<< "${EVENTS:-100 1000 10000 50000}" -export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" -export UV_CACHE_DIR="${UV_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/uv}" -export HF_HOME="${HF_HOME:-${SCRATCH:-${HOME}/.cache}/huggingface}" -export TABPFN_MODEL_CACHE_DIR="${TABPFN_MODEL_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/tabpfn}" -mkdir -p logs -overrides=() -if (($#)); then overrides=(--override "$@"); fi -uv run --all-extras python scripts/label_efficiency.py "${configs[@]}" \ - --events "${events[@]}" --out "${OUT:-outputs/label_efficiency}" ${overrides[@]+"${overrides[@]}"} diff --git a/scripts/label_efficiency_study.sh b/scripts/label_efficiency_study.sh deleted file mode 100755 index 375e2da230edf1a2ebca5a656e22a19ba9d38e6f..0000000000000000000000000000000000000000 --- a/scripts/label_efficiency_study.sh +++ /dev/null @@ -1,120 +0,0 @@ -#!/bin/bash -# Every run of the label-efficiency study in docs/design.md ("Label -# efficiency on GamGam"), with its exact overrides, one after another on one -# GPU (the study ran them as parallel jobs, 3.5 hours on 4 GPUs): -# bash scripts/label_efficiency_study.sh -# OUT (default outputs/study) holds runs/, label_efficiency/ and probes/; -# DELPHES is the Delphes event cache. Restartable: finished runs are skipped. -# TabPFN-3.5 needs its weights in TABPFN_MODEL_CACHE_DIR (docs/usage.md). -# Prepare the Delphes and GamGam event caches first. - -set -euo pipefail -cd "$(dirname "$0")/.." -OUT="${OUT:-outputs/study}" -D="${DELPHES:-${SCRATCH:-outputs}/gnn4colliders/delphes/events.pt}" -export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" -export UV_CACHE_DIR="${UV_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/uv}" -export HF_HOME="${HF_HOME:-${SCRATCH:-${HOME}/.cache}/huggingface}" -export TABPFN_MODEL_CACHE_DIR="${TABPFN_MODEL_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/tabpfn}" -G4C=(uv run --all-extras gnn4colliders) -EVAL=data.eval_events=20000 -PGNN=Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt -# The encoder sizes of configs/delphes_{jepa,masked}.yaml. -TRANSFORMER=(model.backbone=transformer model.hid_size=128 model.n_layers=4 - model.n_heads=8 model.ffn_size=512 model.pair_bias=true) - -sweep() { # [key=value ...]: N = 100, 1000, 10000 - uv run --all-extras python scripts/label_efficiency.py "$1" \ - --events 100 1000 10000 --out "${OUT}/label_efficiency" --override "${EVAL}" "${@:2}" -} -steps() { # [key=value ...]: a fixed number of steps per N - local n batch epochs - while read -r n batch epochs; do - uv run --all-extras python scripts/label_efficiency.py "$1" --events "${n}" \ - --out "${OUT}/label_efficiency" --override name="$2" "${EVAL}" \ - data.batch_size="${batch}" train.max_epochs="${epochs}" train.early_stopping=null "${@:3}" - done <<< $'100 100 1000\n1000 128 400\n10000 256 100' -} -train() { # [key=value ...] - [[ -e "${OUT}/runs/$2/last.pt" ]] || - "${G4C[@]}" train "$1" data.cache="${D}" "${EVAL}" run_dir="${OUT}/runs/$2" "${@:3}" -} -probes() { # : k-NN and linear probes of its last epoch - for n in 100 1000 10000; do - local out="${OUT}/probes/$1/n${n}.json" - [[ -e "${out}" ]] || "${G4C[@]}" probe configs/gamgam_jepa_finetune.yaml split=validation \ - model.pretrained="${OUT}/runs/$1/last.pt" data.train_events="${n}" "${EVAL}" output="${out}" - done -} - -# Supervised, from scratch and from Pretrained_GNN. -steps configs/gamgam.yaml gamgam_steps -steps configs/gamgam_finetune.yaml gamgam_finetune_steps -steps configs/gamgam_transformer.yaml gamgam_transformer_steps - -# Label-free pretraining on Delphes, probes, and fine-tuning. -for objective in jepa masked contrastive; do - train "configs/delphes_${objective}.yaml" "delphes_${objective}" \ - data.train_events=4000000 train.max_epochs=8 - probes "delphes_${objective}" -done -train configs/delphes_contrastive.yaml delphes_contrastive_noscale \ - data.train_events=4000000 train.max_epochs=4 model.scale_jitter=0.0 -probes delphes_contrastive_noscale -for n in 100 1000 10000; do # the random encoder - out="${OUT}/probes/random/n${n}.json" - [[ -e "${out}" ]] || "${G4C[@]}" probe configs/gamgam_jepa_finetune.yaml split=validation \ - model.pretrained=null data.train_events="${n}" "${EVAL}" output="${out}" -done -for objective in jepa masked; do - steps configs/gamgam_jepa_finetune.yaml "gamgam_${objective}_finetune_steps" \ - model.pretrained="${OUT}/runs/delphes_${objective}/last.pt" -done - -# In-context models as released, and TabPFN with other inputs. -sweep configs/gamgam_tabpfn.yaml -sweep configs/gamgam_tabicl.yaml -sweep configs/gamgam_tabdpt.yaml -for config in configs/gamgam_tabpfn.yaml configs/gamgam_tabicl.yaml; do - uv run --all-extras python scripts/label_efficiency.py "${config}" --events 50000 \ - --out "${OUT}/label_efficiency" --override "${EVAL}" -done -sweep configs/gamgam_tabpfn.yaml name=gamgam_tabpfn_physics model.relative_phi=true model.pair_objects=4 -sweep configs/gamgam_tabpfn.yaml name=gamgam_tabpfn_pgnn model.tabular_inputs=both \ - model.backbone=edge_network model.hid_size=64 model.n_layers=4 model.pretrained="${PGNN}" -for objective in jepa masked; do - sweep configs/gamgam_tabpfn.yaml name="gamgam_tabpfn_${objective}" model.tabular_inputs=both \ - "${TRANSFORMER[@]}" model.pretrained="${OUT}/runs/delphes_${objective}/last.pt" -done - -# Meta-training on Delphes tasks. One step at learning rate 0 keeps the LoRA -# updates at zero: the released network by its direct forward pass. -train configs/delphes_tabpfn.yaml delphes_tabpfn_zero data.eval_events=2000 \ - train.max_epochs=1 train.learning_rate=0.0 task.steps_per_epoch=1 task.validation_steps=1 -sweep configs/gamgam_tabpfn.yaml name=gamgam_tabpfn_direct checkpoint="${OUT}/runs/delphes_tabpfn_zero/best.pt" -train configs/delphes_tabpfn.yaml delphes_tabpfn_all12 train.max_epochs=6 -sweep configs/gamgam_tabpfn.yaml name=gamgam_tabpfn_meta_all12 checkpoint="${OUT}/runs/delphes_tabpfn_all12/best.pt" -train configs/delphes_pfn.yaml delphes_pfn_all12 train.max_epochs=25 -sweep configs/gamgam_pfn.yaml name=gamgam_pfn_all12 checkpoint="${OUT}/runs/delphes_pfn_all12/best.pt" - -# The supervised runs' kept epochs and the leading in-context runs on all -# validation events. -full() { # [ key=value ...]; by default the run's own config - local n dir - for n in 100 1000 10000; do - dir="${OUT}/label_efficiency/$1/n${n}" - [[ -e "${dir}/validation_full_metrics.json" ]] || - "${G4C[@]}" evaluate "${2:-${dir}/config.yaml}" split=validation \ - data.train_events="${n}" data.eval_events=null run_dir="${dir}" \ - output="${dir}/validation_full_metrics.json" "${@:3}" - done -} -for run in gamgam_steps gamgam_finetune_steps gamgam_transformer_steps \ - gamgam_jepa_finetune_steps gamgam_masked_finetune_steps; do - full "${run}" -done -full gamgam_tabpfn configs/gamgam_tabpfn.yaml -full gamgam_tabpfn_pgnn configs/gamgam_tabpfn.yaml model.tabular_inputs=both \ - model.backbone=edge_network model.hid_size=64 model.n_layers=4 model.pretrained="${PGNN}" -full gamgam_tabpfn_masked configs/gamgam_tabpfn.yaml model.tabular_inputs=both \ - "${TRANSFORMER[@]}" model.pretrained="${OUT}/runs/delphes_masked/last.pt" diff --git a/scripts/meta_train.sh b/scripts/meta_train.sh deleted file mode 100755 index 55105bf644c73531e58bbbf1f62b71f7127807b0..0000000000000000000000000000000000000000 --- a/scripts/meta_train.sh +++ /dev/null @@ -1,36 +0,0 @@ -#!/bin/bash -# Meta-train an in-context model on one GPU, then classify GamGam in context -# with its best checkpoint for each number of labelled events: -# sbatch -A --export=ALL,CONFIG=configs/delphes_tabicl.yaml,EVAL_CONFIG=configs/gamgam_tabicl.yaml \ -# scripts/meta_train.sh [key=value ...] -# CONFIG is configs/delphes_{pfn,tabpfn,tabicl}.yaml and EVAL_CONFIG the -# matching configs/gamgam_*.yaml (none: meta-training only). key=value -# overrides apply to both (e.g. task.prior_classes=[1,6,7,8,9,10,11], which -# leaves the GamGam processes out; model.lora_rank=0). Prepare the Delphes -# and GamGam event caches first. -#SBATCH -N 1 -#SBATCH -C gpu -#SBATCH -q shared -#SBATCH -G 1 -#SBATCH -c 32 -#SBATCH -t 12:00:00 -#SBATCH -o logs/meta_train_%j.out - -set -euo pipefail -cd "${SLURM_SUBMIT_DIR:-$(pwd)}" -CONFIG="${CONFIG:-configs/delphes_pfn.yaml}" -name="$(basename "${CONFIG}" .yaml)" -RUN="${RUN:-outputs/${name}}" -export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" -export UV_CACHE_DIR="${UV_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/uv}" -export HF_HOME="${HF_HOME:-${SCRATCH:-${HOME}/.cache}/huggingface}" -export TABPFN_MODEL_CACHE_DIR="${TABPFN_MODEL_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/tabpfn}" -mkdir -p logs -uv run --all-extras gnn4colliders train "${CONFIG}" run_dir="${RUN}" "$@" -if [[ -n "${EVAL_CONFIG:-}" ]]; then - read -r -a events <<< "${EVENTS:-100 1000 10000 50000}" - evaluated="$(basename "${EVAL_CONFIG}" .yaml)_from_${name}" - uv run --all-extras python scripts/label_efficiency.py "${EVAL_CONFIG}" \ - --events "${events[@]}" --out "${OUT:-outputs/label_efficiency}" \ - --override name="${evaluated}" checkpoint="${RUN}/best.pt" "$@" -fi diff --git a/scripts/plot_label_efficiency.py b/scripts/plot_label_efficiency.py deleted file mode 100644 index 7235641d170266a20cd5c0a8af4e113c775eedca..0000000000000000000000000000000000000000 --- a/scripts/plot_label_efficiency.py +++ /dev/null @@ -1,184 +0,0 @@ -"""Plot validation AUC against the number of labelled training events, and -write the same numbers as Markdown and LaTeX tables. - -Reads the sweeps of ``scripts/label_efficiency.py`` (``//n/ -_metrics.json``, one series per ````) and the probes of -``scripts/pretrain.sh`` (``/probe/_n.json``, two series per -run: k-NN and linear): - - uv run --extra plots python scripts/plot_label_efficiency.py \\ - --sweeps outputs/label_efficiency --probes outputs/delphes_jepa \\ - --output plots/label_efficiency - -writes ``plots/label_efficiency.{pdf,png,md,tex}``. ``--metric`` chooses -another metric (``accuracy``, ``f1``, ...). -""" - -from __future__ import annotations - -import argparse -import json -import re -from pathlib import Path - -#: Categorical colours in fixed order (validated for colour-vision -#: deficiency on a light background) and marker shapes, so that series are -#: not told apart by colour alone. -COLOURS = ( - "#2a78d6", - "#eb6834", - "#1baf7a", - "#eda100", - "#e87ba4", - "#008300", - "#6250d6", - "#e34948", -) -MARKERS = ("o", "s", "^", "D", "v", "P", "X", "*") -TEXT, MUTED, GRID = "#2b2b29", "#6b6a64", "#e4e3de" - - -def _sweeps(root: Path, split: str, metric: str) -> dict[str, dict[int, float]]: - series: dict[str, dict[int, float]] = {} - for path in sorted(root.glob(f"*/n*/{split}_metrics.json")): - events = int(path.parent.name.removeprefix("n")) - value = json.loads(path.read_text())["metrics"][metric] - series.setdefault(path.parent.parent.name, {})[events] = value - return series - - -def _probes(run: Path, split: str, metric: str) -> dict[str, dict[int, float]]: - series: dict[str, dict[int, float]] = {} - for path in sorted((run / "probe").glob(f"{split}_n*.json")): - match = re.fullmatch(rf"{split}_n(\d+)\.json", path.name) - if match is None: - continue - summary = json.loads(path.read_text()) - for probe, label in (("knn", "k-NN"), ("linear", "linear")): - name = f"{run.name} ({label} probe)" - series.setdefault(name, {})[int(match[1])] = summary[probe][metric] - return series - - -def _tables(series: dict[str, dict[int, float]], metric: str) -> tuple[str, str]: - counts = sorted({count for values in series.values() for count in values}) - - def cell(values: dict[int, float], count: int) -> str: - return f"{100 * values[count]:.2f}" if count in values else "--" - - header = ["model", *(f"N = {count}" for count in counts)] - markdown = [ - "| " + " | ".join(header) + " |", - "|---" * len(header) + "|", - *( - "| " + " | ".join([name, *(cell(v, c) for c in counts)]) + " |" - for name, v in series.items() - ), - ] - latex = [ - r"\begin{tabular}{l" + "r" * len(counts) + "}", - r"\hline", - " & ".join(["Model", *(f"$N={count}$" for count in counts)]) + r" \\", - r"\hline", - *( - " & ".join([name.replace("_", r"\_"), *(cell(v, c) for c in counts)]) - + r" \\" - for name, v in series.items() - ), - r"\hline", - r"\end{tabular}", - f"% {metric} (%) against labelled training events N", - ] - return "\n".join(markdown) + "\n", "\n".join(latex) + "\n" - - -def _plot(series: dict[str, dict[int, float]], metric: str, output: Path) -> None: - import matplotlib - - matplotlib.use("Agg") - import matplotlib.pyplot as plt - from matplotlib.ticker import FuncFormatter, LogLocator - - if len(series) > len(COLOURS): - raise SystemExit(f"at most {len(COLOURS)} series; plot fewer per figure") - figure, axes = plt.subplots(figsize=(6.4, 4.2), layout="constrained") - figure.set_facecolor("#fcfcfb") - axes.set_facecolor("#fcfcfb") - for index, (name, values) in enumerate(series.items()): - counts = sorted(values) - percent = [100 * values[count] for count in counts] - axes.plot( - counts, - percent, - color=COLOURS[index], - marker=MARKERS[index], - markersize=7, - linewidth=2, - markeredgecolor="#fcfcfb", - markeredgewidth=1.5, - label=name, - ) - # A direct label at the last point when there are few series. - if len(series) <= 4: - axes.annotate( - name, - (counts[-1], percent[-1]), - xytext=(6, 0), - textcoords="offset points", - va="center", - fontsize=8, - color=TEXT, - ) - axes.set_xscale("log") - axes.xaxis.set_major_locator(LogLocator(base=10)) - axes.xaxis.set_major_formatter(FuncFormatter(lambda x, _: f"{x:,.0f}")) - if len(series) <= 4: # room on the right for the direct labels - low, high = axes.get_xlim() - axes.set_xlim(low, high * 4) - axes.set_xlabel("labelled training events", color=TEXT) - label = "AUC" if metric == "auc" else metric - axes.set_ylabel(f"validation {label} (%)", color=TEXT) - axes.grid(True, which="major", color=GRID, linewidth=0.8) - axes.tick_params(colors=MUTED) - for spine in ("top", "right"): - axes.spines[spine].set_visible(False) - for spine in ("left", "bottom"): - axes.spines[spine].set_color(GRID) - if len(series) > 1: - axes.legend(frameon=False, fontsize=8, labelcolor=TEXT) - for suffix in (".pdf", ".png"): - figure.savefig(output.with_suffix(suffix), dpi=200) - plt.close(figure) - - -def main() -> None: - parser = argparse.ArgumentParser( - description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter - ) - parser.add_argument("--sweeps", nargs="*", type=Path, default=[]) - parser.add_argument("--probes", nargs="*", type=Path, default=[]) - parser.add_argument("--split", default="validation") - parser.add_argument("--metric", default="auc") - parser.add_argument("--output", type=Path, default=Path("plots/label_efficiency")) - args = parser.parse_args() - series: dict[str, dict[int, float]] = {} - found = [_sweeps(root, args.split, args.metric) for root in args.sweeps] - found += [_probes(run, args.split, args.metric) for run in args.probes] - for part in found: - for name, values in part.items(): - merged = series.setdefault(name, {}) - if merged.keys() & values.keys(): - raise SystemExit(f"two results for {name} at the same N; plot apart") - merged.update(values) - if not series: - raise SystemExit("no results found") - args.output.parent.mkdir(parents=True, exist_ok=True) - markdown, latex = _tables(series, args.metric) - args.output.with_suffix(".md").write_text(markdown) - args.output.with_suffix(".tex").write_text(latex) - _plot(series, args.metric, args.output) - print(markdown, end="") - - -if __name__ == "__main__": - main() diff --git a/scripts/pretrain.sh b/scripts/pretrain.sh deleted file mode 100755 index 21bc3726ad4548dfefa4da22ae4d26b984716662..0000000000000000000000000000000000000000 --- a/scripts/pretrain.sh +++ /dev/null @@ -1,38 +0,0 @@ -#!/bin/bash -# Label-free pretraining of an object encoder on one GPU, then k-NN and -# linear probes of its last epoch on GamGam for each number of labelled -# events: -# sbatch -A --export=ALL,CONFIG=configs/delphes_contrastive.yaml scripts/pretrain.sh [key=value ...] -# CONFIG is configs/delphes_{jepa,contrastive,masked}.yaml (or -# delphes_jepa_edge.yaml with PROBE_CONFIG=configs/gamgam_finetune.yaml); -# PROBE_CONFIG (default configs/gamgam_jepa_finetune.yaml) must have the -# encoder's sizes. The probes evaluate EVAL_EVENTS (default 20000) random -# validation events. key=value overrides apply to the pretraining. For -# several GPUs, pretrain with scripts/train_ddp.sh instead. Prepare the -# Delphes and GamGam event caches first. -#SBATCH -N 1 -#SBATCH -C gpu -#SBATCH -q shared -#SBATCH -G 1 -#SBATCH -c 32 -#SBATCH -t 12:00:00 -#SBATCH -o logs/pretrain_%j.out - -set -euo pipefail -cd "${SLURM_SUBMIT_DIR:-$(pwd)}" -CONFIG="${CONFIG:-configs/delphes_jepa.yaml}" -PROBE_CONFIG="${PROBE_CONFIG:-configs/gamgam_jepa_finetune.yaml}" -name="$(basename "${CONFIG}" .yaml)" -RUN="${RUN:-outputs/${name}}" -read -r -a events <<< "${EVENTS:-100 1000 10000}" -export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" -export UV_CACHE_DIR="${UV_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/uv}" -export HF_HOME="${HF_HOME:-${SCRATCH:-${HOME}/.cache}/huggingface}" -mkdir -p logs -uv run gnn4colliders train "${CONFIG}" run_dir="${RUN}" "$@" -for count in "${events[@]}"; do - uv run gnn4colliders probe "${PROBE_CONFIG}" split=validation \ - model.pretrained="${RUN}/last.pt" data.train_events="${count}" \ - data.eval_events="${EVAL_EVENTS:-20000}" \ - output="${RUN}/probe/validation_n${count}.json" -done diff --git a/scripts/shrink_root_sample.py b/scripts/shrink_root_sample.py deleted file mode 100644 index 92527f29d092f36afb8bb9a5d3a2e182800ac268..0000000000000000000000000000000000000000 --- a/scripts/shrink_root_sample.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Create a small ROOT fixture from a Delphes or ATLAS Open Data sample.""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import uproot - -DELPHES_BRANCHES = ( - "jet_pt", - "jet_eta", - "jet_phi", - "jet_btag", - "ph_pt", - "ph_eta", - "ph_phi", - "ele_pt", - "ele_eta", - "ele_phi", - "ele_charge", - "mu_pt", - "mu_eta", - "mu_phi", - "mu_charge", - "MET_met", - "MET_phi", - "weight", - "Number", -) - -#: Branches read by ``configs/gamgam.yaml``. -ATLAS_OPENDATA_BRANCHES = ( - "jet_pt", - "jet_eta", - "jet_phi", - "jet_MV2c10", - "photon_pt", - "photon_eta", - "photon_phi", - "lep_pt", - "lep_eta", - "lep_phi", - "lep_charge", - "lep_type", - "met_et", - "met_phi", - "mcWeight", - "eventNumber", -) - -PRESETS = { - "delphes": ("output", DELPHES_BRANCHES), - "atlas_opendata": ("mini", ATLAS_OPENDATA_BRANCHES), -} -BRANCHES = DELPHES_BRANCHES - - -def shrink_sample( - source: Path, target: Path, entries: int, *, preset: str = "delphes" -) -> None: - """Copy the active branches and first ``entries`` events to ``target``.""" - if entries < 1: - raise ValueError("entries must be positive") - tree, branches = PRESETS[preset] - target.parent.mkdir(parents=True, exist_ok=True) - with uproot.open(source) as source_file: - arrays = source_file[tree].arrays( - branches, entry_start=0, entry_stop=entries, library="ak" - ) - with uproot.recreate(target) as target_file: - target_file[tree] = arrays - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("source", type=Path) - parser.add_argument("target", type=Path) - parser.add_argument("--entries", type=int, default=64) - parser.add_argument("--preset", choices=sorted(PRESETS), default="delphes") - args = parser.parse_args() - shrink_sample(args.source, args.target, args.entries, preset=args.preset) - print(f"Wrote {args.entries} events to {args.target}") - - -if __name__ == "__main__": - main() diff --git a/scripts/summarize_ensemble.py b/scripts/summarize_ensemble.py deleted file mode 100644 index b5be3a9577d2e1b3af31c727069d2c0886d5251f..0000000000000000000000000000000000000000 --- a/scripts/summarize_ensemble.py +++ /dev/null @@ -1,64 +0,0 @@ -"""Summarize an ensemble of training runs as mean +- standard deviation. - -Each run directory must contain ``history.jsonl`` and ``test_metrics.json`` -(written by ``scripts/train.sh``). The standard deviation -is over runs, as in Table 3 of Ho et al. (2026). -""" - -from __future__ import annotations - -import argparse -import json -import statistics -from pathlib import Path - -METRICS = ("accuracy", "auc", "f1", "precision", "recall") - - -def _run(path: Path) -> dict[str, float]: - history = [ - json.loads(line) for line in (path / "history.jsonl").read_text().splitlines() - ] - test = json.loads((path / "test_metrics.json").read_text()) - best = max(history, key=lambda row: row["validation"]["auc"]) - seconds = [row["seconds"] for row in history] - return { - **{name: 100 * test["metrics"][name] for name in METRICS}, - "epochs": len(history), - "best_epoch": best["epoch"], - "seconds_per_epoch": statistics.median(seconds), - "training_hours": sum(seconds) / 3600, - } - - -def _format(values: list[float]) -> str: - mean = statistics.fmean(values) - spread = statistics.stdev(values) if len(values) > 1 else 0.0 - return f"{mean:.2f} +- {spread:.2f}" - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("runs", nargs="+", type=Path) - args = parser.parse_args() - runs = {path.name: _run(path) for path in args.runs} - keys = list(next(iter(runs.values()))) - width = max(len(name) for name in runs) - print(" ".join([f"{'run':{width}}", *(f"{key:>17}" for key in keys)])) - for name, values in runs.items(): - print(" ".join([f"{name:{width}}", *(f"{values[key]:17.3f}" for key in keys)])) - print( - " ".join( - [ - f"{'mean':{width}}", - *( - f"{_format([run[key] for run in runs.values()]):>17}" - for key in keys - ), - ] - ) - ) - - -if __name__ == "__main__": - main() diff --git a/scripts/train.sh b/scripts/train.sh deleted file mode 100755 index 17489f078a0f35fd12292d01c1874e78d582c0b6..0000000000000000000000000000000000000000 --- a/scripts/train.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/bash -# Train CONFIG on one GPU, then evaluate its best epoch on the test split. -# sbatch -A --export=ALL,SEED=1 scripts/train.sh [key=value ...] -# The paper's ensembles vary only the seed (SEED=1..5). Prepare the event -# cache first: uv run gnn4colliders prepare configs/gamgam.yaml -#SBATCH -N 1 -#SBATCH -C gpu -#SBATCH -q shared -#SBATCH -G 1 -#SBATCH -c 32 -#SBATCH -t 06:00:00 -#SBATCH -o logs/train_%j.out - -set -euo pipefail -cd "${SLURM_SUBMIT_DIR:-$(pwd)}" -CONFIG="${CONFIG:-configs/gamgam.yaml}" -SEED="${SEED:-42}" -RUN="${RUN:-outputs/$(basename "${CONFIG}" .yaml)/seed${SEED}}" -# PyTorch otherwise sizes CPU thread pools to every core on the node. -export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" -# Home-directory file locks are unavailable on Perlmutter compute nodes. -export UV_CACHE_DIR="${UV_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/uv}" -# Hugging Face downloads (the Delphes samples are 52 GB) go to scratch too. -export HF_HOME="${HF_HOME:-${SCRATCH:-${HOME}/.cache}/huggingface}" -mkdir -p logs -args=("${CONFIG}" run_dir="${RUN}" train.seed="${SEED}" data.seed="${SEED}" "$@") -uv run gnn4colliders train "${args[@]}" -uv run gnn4colliders evaluate "${args[@]}" split=test diff --git a/scripts/train_ddp.sh b/scripts/train_ddp.sh deleted file mode 100755 index 30e9a3466f67dee1e824c1cf5093c71748cc89aa..0000000000000000000000000000000000000000 --- a/scripts/train_ddp.sh +++ /dev/null @@ -1,25 +0,0 @@ -#!/bin/bash -# Data-parallel training on all GPUs of one or more nodes: -# sbatch -A -N 2 scripts/train_ddp.sh [key=value ...] -# data.batch_size is per GPU. CUDA graphs need a single process. -#SBATCH -N 1 -#SBATCH -C gpu -#SBATCH --ntasks-per-node=1 -#SBATCH --gpus-per-node=4 -#SBATCH -c 64 -#SBATCH -t 02:00:00 -#SBATCH -o logs/train_ddp_%j.out - -set -euo pipefail -cd "${SLURM_SUBMIT_DIR:-$(pwd)}" -CONFIG="${CONFIG:-configs/gamgam.yaml}" -export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}" -export UV_CACHE_DIR="${UV_CACHE_DIR:-${SCRATCH:-${HOME}/.cache}/uv}" -# Hugging Face downloads (the Delphes samples are 52 GB) go to scratch too. -export HF_HOME="${HF_HOME:-${SCRATCH:-${HOME}/.cache}/huggingface}" -MASTER_ADDR="$(scontrol show hostnames "${SLURM_JOB_NODELIST}" | head -n 1)" -mkdir -p logs -srun uv run torchrun \ - --nnodes="${SLURM_NNODES}" --nproc_per_node="${GPUS_PER_NODE:-4}" \ - --rdzv_backend=c10d --rdzv_endpoint="${MASTER_ADDR}:${MASTER_PORT:-29500}" \ - -m gnn4colliders.cli train "${CONFIG}" "$@" diff --git a/src/gnn4colliders/__init__.py b/src/gnn4colliders/__init__.py deleted file mode 100644 index c68fe192f771ebebb70dfab6a85ded1f0783cfca..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Public package metadata for GNN4Colliders.""" - -__version__ = "0.1.0" - -__all__ = ["__version__"] diff --git a/src/gnn4colliders/cli.py b/src/gnn4colliders/cli.py deleted file mode 100644 index 2bcb9b82aaa844ab3eb46520720ee5a93363532c..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/cli.py +++ /dev/null @@ -1,497 +0,0 @@ -"""``gnn4colliders CONFIG [key=value ...]``. - -Commands: - -* ``prepare``: read the ROOT files into the event cache ``data.cache``. -* ``train``: train and write ``best.pt``, ``last.pt``, ``history.jsonl`` - and ``config.yaml`` to ``run_dir``. Launch with ``torchrun`` for DDP. -* ``evaluate``: metrics of ``checkpoint`` on ``split``, written as JSON. -* ``predict``: per-event scores of ``checkpoint`` on ``split``, as ``.npz``. -* ``export``: ONNX model of ``checkpoint``. -* ``probe``: k-NN and linear classifiers of the frozen encodings of the - encoder of ``model.pretrained`` (random without), trained on - ``data.train_events`` events; their metrics on ``split``, as JSON. - -The in-context models (``model.name: pfn``, ``tabpfn``, ``tabicl`` or -``tabdpt``) are evaluated -without training on the task: their context is ``data.train_events`` -labelled training events. The PFN, TabPFN and TabICL are meta-trained with -``task.name: pfn``; TabPFN and TabICL without a ``checkpoint`` are the -classifiers as released, and TabDPT is only used as released. -""" - -from __future__ import annotations - -import argparse -import json -import logging -import os -import sys -import time -from pathlib import Path -from typing import Any - -import torch - -from .config import Config, load_config, save_config -from .data import PackedEvents, load_events, prepare -from .distributed import ( - DistributedContext, - barrier, - finalize, - initialize, - resolve_device, - synchronize, -) -from .graphs import collate_graph_batch -from .inference import collect_outputs, write_predictions -from .loader import EventLoader -from .models import IN_CONTEXT_MODELS, SEQUENCE_MODELS, build_model, load_weights -from .sequences import collate_sequence_batch -from .tasks import build_task -from .train import ( - EarlyStopping, - EpochResult, - Trainer, - load_checkpoint, - restore, - save_checkpoint, - seed_everything, -) - -logger = logging.getLogger("gnn4colliders") - - -def make_loaders( - config: Config, - events: PackedEvents, - splits: tuple[str, ...], - context: DistributedContext | None = None, -) -> dict[str, EventLoader]: - data = config.data - collate = ( - collate_sequence_batch - if config.model.name in SEQUENCE_MODELS - else collate_graph_batch - ) - folds = { - "train": data.train_folds, - "validation": data.validation_folds, - "test": data.test_folds, - } - context = context or DistributedContext() - loaders = {} - for split in splits: - indices = events.split_indices(folds[split]) - if split == "train" and data.train_events is not None: - indices = events.balanced_subset(indices, data.train_events, data.seed) - if split != "train" and data.eval_events is not None: - indices = events.random_subset(indices, data.eval_events, data.seed) - if config.task.name == "pfn": - from .pfn import TaskLoader - - loaders[split] = TaskLoader( - events, - indices, - config.task, - data.batch_size, - train=split == "train", - seed=data.seed, - num_workers=data.loader_workers, - pin_memory=data.pin_memory, - rank=context.rank, - ) - continue - loaders[split] = EventLoader( - events, - collate, - data.batch_size, - indices=indices, - shuffle=data.shuffle and split == "train", - seed=data.seed, - num_workers=data.loader_workers, - pin_memory=data.pin_memory, - rank=context.rank, - world_size=context.world_size, - ) - return loaders - - -def make_model(config: Config, events: PackedEvents, *, load_pretrained: bool = True): - return build_model( - config.model, - num_features=events.num_features, - num_node_types=len(config.data.object_types), - num_global_features=events.num_global_features, - load_pretrained=load_pretrained, - ) - - -def make_task(config: Config): - task = config.task - return build_task( - task.name, absolute_weights=task.absolute_weights, threshold=task.threshold - ) - - -def make_trainer(config: Config, model: torch.nn.Module, context: DistributedContext): - train = config.train - optimizers = {"adam": torch.optim.Adam, "adamw": torch.optim.AdamW} - if train.optimizer not in optimizers: - raise ValueError(f"unknown optimizer {train.optimizer!r}; use adam or adamw") - optimizer = optimizers[train.optimizer]( - [parameter for parameter in model.parameters() if parameter.requires_grad], - lr=train.learning_rate, - weight_decay=train.weight_decay, - ) - scheduler = ( - None - if train.lr_decay == 1 - else torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=train.lr_decay) - ) - return Trainer( - model, - make_task(config), - optimizer, - scheduler, - context=context, - early_stopping=EarlyStopping(**train.early_stopping) - if train.early_stopping - else None, - compile_model=train.compile, - cuda_graphs=train.cuda_graphs, - precision=train.precision, - ) - - -class _Recorder: - """Per-epoch history line and checkpoints (main process only). - - ``last.pt`` is written every epoch, ``best.pt`` when the validation - ``train.best_metric`` improves, ``epoch_NNNN.pt`` every - ``train.save_every`` epochs. - """ - - def __init__(self, config: Config, trainer: Trainer, best: float | None) -> None: - self.config, self.trainer, self.best = config, trainer, best - self.history = config.path / "history.jsonl" - self.started = time.perf_counter() - if trainer.context.is_main_process and not config.train.resume: - self.history.unlink(missing_ok=True) - - def _improved(self, validation: EpochResult | None) -> bool: - if validation is None: - return False - metric = self.config.train.best_metric - value = validation.loss if metric == "loss" else validation.metrics[metric] - better = self.best is None or ( - value < self.best if metric == "loss" else value > self.best - ) - if better: - self.best = value - return better - - def __call__(self, train: EpochResult, validation: EpochResult | None) -> None: - if not self.trainer.context.is_main_process: - return - trainer, run = self.trainer, self.config.path - now = time.perf_counter() - record: dict[str, Any] = { - "epoch": trainer.epoch, - "seconds": now - self.started, - "learning_rate": [ - float(group["lr"]) for group in trainer.optimizer.param_groups - ], - "train": {"loss": train.loss, "events": train.num_samples, **train.metrics}, - } - if validation is not None: - record["validation"] = { - "loss": validation.loss, - "events": validation.num_samples, - **validation.metrics, - } - self.started = now - with self.history.open("a") as stream: - stream.write(json.dumps(record) + "\n") - improved = self._improved(validation) - payload = {"config": self.config.to_dict(), "best_value": self.best} - save_checkpoint(run / "last.pt", trainer, **payload) - if improved: - save_checkpoint(run / "best.pt", trainer, **payload) - every = self.config.train.save_every - if every and (trainer.epoch + 1) % every == 0: - save_checkpoint(run / f"epoch_{trainer.epoch:04d}.pt", trainer, **payload) - - -def run_prepare(config: Config) -> Path: - if int(os.environ.get("WORLD_SIZE", "1")) > 1: - raise ValueError("prepare runs in one process; use data.prepare_workers") - return prepare(config.data) - - -def run_train(config: Config) -> Path: - if config.model.name in IN_CONTEXT_MODELS and config.task.name != "pfn": - raise ValueError( - f"{config.model.name} is not trained on {config.task.name}; " - "run evaluate or predict (in context)" - ) - context = initialize(config.train.device) - try: - seed_everything(config.train.seed + context.rank) - torch.backends.cuda.matmul.allow_tf32 = config.train.allow_tf32 - if context.is_main_process: - save_config(config, config.path / "config.yaml") - events = load_events(config.data.cache) - loaders = make_loaders(config, events, ("train", "validation"), context) - resume = load_checkpoint(config.train.resume) if config.train.resume else None - model = make_model(config, events, load_pretrained=resume is None) - if hasattr(model, "total_steps"): # JEPA: EMA momentum schedule - model.total_steps = config.train.max_epochs * len(loaders["train"]) - trainer = make_trainer(config, model, context) - if resume is not None: - restore(trainer, resume) - recorder = _Recorder( - config, trainer, None if resume is None else resume.get("best_value") - ) - validation = ( - loaders["validation"] if len(loaders["validation"].indices) else None - ) - trainer.fit( - loaders["train"], - validation, - max_epochs=config.train.max_epochs, - on_epoch_end=recorder, - ) - barrier(context) - return config.path / "last.pt" - finally: - finalize(context) - - -def _checkpoint(config: Config) -> Path: - if config.checkpoint: - return Path(config.checkpoint) - best = config.path / "best.pt" - return best if best.exists() else config.path / "last.pt" - - -#: The tabular in-context models used as released (without a meta-trained -#: ``checkpoint``), and how the metrics name them. -RELEASED_MODELS = { - "tabpfn": lambda model: ( - f"TabPFN {model.tabpfn_version} " - f"({model.tabpfn_checkpoint or 'default checkpoint'})" - ), - "tabicl": lambda model: f"TabICL ({model.tabicl_checkpoint})", - "tabdpt": lambda model: ( - f"TabDPT (context {model.tabdpt_context_size or 'all'} per query)" - ), -} - - -def in_context_forward( - config: Config, events: PackedEvents, device: torch.device -) -> tuple[Any, str]: - """The function from a batch to logits of an in-context model, given - ``data.train_events`` labelled training events, and what it is.""" - data, model_config = config.data, config.model - if data.train_events is None: - raise ValueError("in-context models need data.train_events (the context)") - if config.task.name != "multiclass": - raise ValueError("in-context models classify with task.name: multiclass") - train = events.split_indices(data.train_folds) - context = events.balanced_subset(train, data.train_events, data.seed) - if model_config.name in RELEASED_MODELS and config.checkpoint is None: - from . import tabular - - forward = getattr(tabular, f"{model_config.name}_forward")( - model_config, - events, - context, - num_classes=model_config.out_size, - num_node_types=len(data.object_types), - batch_size=data.batch_size, - device=device, - seed=config.train.seed, - ) - return forward, RELEASED_MODELS[model_config.name](model_config) - from .pfn import in_context_forward as pfn_forward - from .tabular import meta_tabular_forward - - checkpoint = _checkpoint(config) - model = make_model(config, events, load_pretrained=False) - load_weights(model, load_checkpoint(checkpoint)) - build = ( - meta_tabular_forward - if model_config.name in ("tabpfn", "tabicl") - else pfn_forward - ) - forward = build( - model.to(device), - events, - context, - num_classes=model_config.out_size, - batch_size=data.batch_size, - device=device, - ) - return forward, str(checkpoint) - - -def run_inference(config: Config, *, predict: bool) -> Path: - """Evaluate (metrics JSON) or predict (scores ``.npz``) ``config.split``.""" - context = initialize(config.train.device) - try: - torch.backends.cuda.matmul.allow_tf32 = config.train.allow_tf32 - events = load_events(config.data.cache) - loader = make_loaders(config, events, (config.split,), context)[config.split] - if context.device.type == "cuda": - torch.cuda.reset_peak_memory_stats(context.device) - start = time.perf_counter() - if config.model.name in IN_CONTEXT_MODELS: - forward, checkpoint = in_context_forward(config, events, context.device) - else: - checkpoint = _checkpoint(config) - forward = make_model(config, events, load_pretrained=False) - load_weights(forward, load_checkpoint(checkpoint)) - forward = forward.to(context.device).eval() - synchronize(context.device) - loaded = time.perf_counter() - outputs = collect_outputs( - forward, loader, device=context.device, context=context - ) - synchronize(context.device) - # Setup: loading the checkpoint, or the in-context model's pass over - # the context; inference: this rank's events. - seconds = {"setup": loaded - start, "inference": time.perf_counter() - loaded} - peak_memory = ( - torch.cuda.max_memory_allocated(context.device) / 2**30 - if context.device.type == "cuda" - else None - ) - task = make_task(config) - suffix = "predictions.npz" if predict else "metrics.json" - target = Path(config.output or config.path / f"{config.split}_{suffix}") - if context.is_main_process: - if predict: - write_predictions(target, outputs, task, events) - else: - metrics = task.metrics(outputs.logits, outputs.labels, outputs.weight) - summary = { - "checkpoint": str(checkpoint), - "split": config.split, - "num_samples": int(outputs.labels.shape[0]), - "metrics": metrics, - "seconds": seconds, - "peak_memory_gib": peak_memory, - } - target.parent.mkdir(parents=True, exist_ok=True) - target.write_text(json.dumps(summary, indent=2) + "\n") - print(json.dumps(metrics, indent=2)) - barrier(context) - return target - finally: - finalize(context) - - -def run_export(config: Config) -> Path: - from .export import export_onnx - - if config.model.name in IN_CONTEXT_MODELS: - raise ValueError(f"{config.model.name} is not exported (it needs a context)") - events = load_events(config.data.cache) - loader = make_loaders(config, events, (config.split,))[config.split] - checkpoint = _checkpoint(config) - model = make_model(config, events, load_pretrained=False) - load_weights(model, load_checkpoint(checkpoint)) - return export_onnx( - model, - next(iter(loader)), - config.output or config.path / "model.onnx", - metadata={"checkpoint": str(checkpoint), "model": config.to_dict()["model"]}, - ) - - -def run_probe(config: Config) -> Path: - from .probe import probe - - if int(os.environ.get("WORLD_SIZE", "1")) > 1: - raise ValueError("probe runs in one process") - data = config.data - if data.train_events is None: - raise ValueError("probe needs data.train_events (the probes' training events)") - events = load_events(data.cache) - folds = { - "train": data.train_folds, - "validation": data.validation_folds, - "test": data.test_folds, - } - train = events.balanced_subset( - events.split_indices(data.train_folds), data.train_events, data.seed - ) - queries = events.split_indices(folds[config.split]) - # The probes are not scored on the events they are fitted to. - queries = queries[~torch.isin(queries, train)] - if config.split != "train" and data.eval_events is not None: - queries = events.random_subset(queries, data.eval_events, data.seed) - device = resolve_device(config.train.device) - start = time.perf_counter() - # All the dataset's classes, whatever the config's task (a pretraining - # config probes too). - num_classes = int(events.labels.max()) + 1 - logits = probe( - config, events, train, queries, num_classes=num_classes, device=device - ) - task = build_task("multiclass") - labels, weight = events.labels[queries], events.weight[queries] - summary = { - "encoder": config.model.pretrained, - "split": config.split, - "train_events": int(len(train)), - "num_samples": int(len(queries)), - **{ - name: task.metrics(values.cpu(), labels, weight) - for name, values in logits.items() - }, - "seconds": time.perf_counter() - start, - } - target = Path(config.output or config.path / f"{config.split}_probe.json") - target.parent.mkdir(parents=True, exist_ok=True) - target.write_text(json.dumps(summary, indent=2) + "\n") - print(json.dumps({name: summary[name] for name in logits}, indent=2)) - return target - - -COMMANDS = { - "prepare": run_prepare, - "train": run_train, - "evaluate": lambda config: run_inference(config, predict=False), - "predict": lambda config: run_inference(config, predict=True), - "export": run_export, - "probe": run_probe, -} - - -def main(argv: list[str] | None = None) -> int: - parser = argparse.ArgumentParser( - prog="gnn4colliders", - description=__doc__, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.add_argument("command", choices=COMMANDS) - parser.add_argument("config", help="experiment YAML file, e.g. configs/gamgam.yaml") - parser.add_argument("overrides", nargs="*", help="dotted.key=value overrides") - args = parser.parse_args(argv) - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s") - try: - config = load_config(args.config, args.overrides) - result = COMMANDS[args.command](config) - except (ValueError, FileNotFoundError, KeyError, ImportError) as error: - print(f"gnn4colliders: {error}", file=sys.stderr) - return 2 - if int(os.environ.get("RANK", "0")) == 0: - print(f"{args.command}: wrote {result}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/src/gnn4colliders/config.py b/src/gnn4colliders/config.py deleted file mode 100644 index 7b590267401dc6dc5515cf5d021c32418f459805..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/config.py +++ /dev/null @@ -1,310 +0,0 @@ -"""Experiment configuration: one YAML file per experiment. - -A file has the sections ``data``, ``model``, ``task`` and ``train`` plus a -few top-level keys (see ``configs/gamgam.yaml``). Command-line overrides are -``dotted.key=value`` with YAML values, e.g. ``train.fast=true`` or -``data.files=[{path: a.root, label: 0}]``. Unknown keys are errors. -""" - -from __future__ import annotations - -import dataclasses -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -import yaml - - -@dataclass -class DataConfig: - #: ROOT inputs: ``{path, label}`` with optional ``sha256``/``size_bytes``. - #: ``label`` is a class index or the name of a label branch. - files: list[dict[str, Any]] - tree_name: str - feature_branches: list[Any] - object_types: list[str] - scales: list[float] - #: Prepared event cache written by ``prepare``. - cache: str - #: Read ``files`` from this Hugging Face dataset at ``hf_revision``. - hf_repo: str | None = None - hf_revision: str | None = None - class_names: list[str] | None = None - derived_branches: dict[str, Any] = field(default_factory=dict) - global_features: list[str] = field(default_factory=list) - #: The fold is ``fold_var`` (modulo ``fold_modulus`` when set). - fold_var: str = "eventNumber" - fold_modulus: int | None = None - #: Per-event weight branch; unweighted when null. - weight_var: str | None = None - train_folds: list[int] = field(default_factory=lambda: list(range(8))) - validation_folds: list[int] = field(default_factory=lambda: [8]) - test_folds: list[int] = field(default_factory=lambda: [9]) - #: Train on (and give in-context models as context) only this many - #: training events, an equal number per class drawn with ``seed``. - train_events: int | None = None - #: Evaluate (validation in training, evaluate, predict, probe) only this - #: many random validation or test events, drawn with ``seed``: 20000 - #: give AUCs to about 0.003, for quick runs. - eval_events: int | None = None - #: Events per batch and process. - batch_size: int = 1024 - shuffle: bool = True - seed: int = 42 - prepare_workers: int = 0 - #: Training-time collation workers and pinned batches. - loader_workers: int = 0 - pin_memory: bool = False - - def __post_init__(self) -> None: - folds = [ - set(self.train_folds), - set(self.validation_folds), - set(self.test_folds), - ] - if folds[0] & folds[1] or folds[0] & folds[2] or folds[1] & folds[2]: - raise ValueError("train, validation and test folds must be disjoint") - - -@dataclass -class ModelConfig: - #: ``edge_network``, ``transformer``, ``deep_sets``, ``jepa`` or - #: ``contrastive`` (pretraining), or the in-context models ``pfn``, - #: ``tabpfn``, ``tabicl`` and ``tabdpt``. - name: str = "edge_network" - hid_size: int = 64 - out_size: int = 1 - n_layers: int = 4 - n_proc_steps: int = 4 - dropout: float = 0.0 - #: Fine-tune from this checkpoint (new classifier): a supervised or JEPA - #: checkpoint of the same backbone. - pretrained: str | None = None - freeze_backbone: bool = False - # Transformer and JEPA predictor. - n_heads: int = 4 - #: Feed-forward width (default ``4 * hid_size``). - ffn_size: int | None = None - #: Attention bias from object-pair features (Particle Transformer). - pair_bias: bool = False - # JEPA and contrastive pretraining, probes and tabular embeddings. - #: The encoder: ``transformer`` or ``edge_network`` (the sizes above - #: apply to it; JEPA's predictor is a transformer). - backbone: str = "transformer" - # JEPA only. - #: What the predictor predicts for hidden objects: their ``encoding`` by - #: the target encoder, or their kinematic ``features`` (masked objects). - jepa_target: str = "encoding" - predictor_layers: int = 2 - #: Fraction of each event's objects to predict (at least one, and at - #: least one left visible). - mask_fraction: float = 0.4 - #: Initial EMA momentum of the target encoder; rises linearly to 1. - ema_momentum: float = 0.996 - #: Weight of a VICReg variance term on the context encodings (and the - #: edge network's global state) that counters collapse. - variance_weight: float = 0.0 - # Contrastive only (contrastive.py): the NT-Xent temperature and the - # views' momentum-scale jitter and object dropout probability. - temperature: float = 0.1 - scale_jitter: float = 0.1 - object_dropout: float = 0.1 - # Probes of frozen encoders (probe.py). - probe_neighbours: int = 20 - probe_l2: float = 1e-4 - # PFN only (the encoder is the transformer above). - #: Classes the PFN can tell apart in one task. - max_classes: int = 10 - #: Attention blocks over each task's events. - pfn_layers: int = 6 - # TabPFN only. - tabpfn_version: str = "v3.5" - #: A checkpoint of that version (a path, or a file name in the TabPFN - #: cache) instead of its default, e.g. - #: ``tabpfn-v3.5-20260909_multiclass.safetensors``. - tabpfn_checkpoint: str | None = None - tabpfn_estimators: int = 4 - #: Meta-training of the TabPFN or TabICL network (``task.name: pfn``): - #: rank of the LoRA updates of its in-context blocks, or 0 to train all - #: its weights. - lora_rank: int = 8 - # TabICL only. - tabicl_checkpoint: str = "tabicl-classifier-v2-20260212.ckpt" - tabicl_estimators: int = 8 - # TabDPT only. - tabdpt_estimators: int = 8 - #: Context events per query, its nearest rows (retrieval); null: all. - tabdpt_context_size: int | None = None - # The table rows of TabPFN, TabICL and TabDPT (tabular.py). - #: Objects of each type (in ``data.object_types`` order) in a table row. - slots: list[int] = field(default_factory=lambda: [6, 2, 2, 3, 1]) - #: Azimuths measured from that of the event's leading object. - relative_phi: bool = False - #: Invariant mass and Delta R of every pair of this many leading objects. - pair_objects: int = 0 - #: What a tabular model sees of an event: its ``rows``, the - #: ``embeddings`` of the encoder of ``model.backbone`` and - #: ``model.pretrained`` (sizes as above), or ``both``. - tabular_inputs: str = "rows" - - -@dataclass -class TaskConfig: - #: ``multiclass`` (``model.out_size`` classes), ``binary`` (one logit) - #: ``jepa`` or ``contrastive`` (pretraining, with the same - #: ``model.name``) or ``pfn`` - #: (meta-training with ``model.name: pfn``, ``tabpfn`` or ``tabicl``). - name: str = "multiclass" - #: Use ``|weight|`` in the loss (signed NLO weights). - absolute_weights: bool = False - threshold: float = 0.5 - # PFN meta-training: the prior over tasks (see pfn.py). - #: Labels tasks draw their classes from (default: all). - prior_classes: list[int] | None = None - #: Smallest and largest number of classes of a task. - task_classes: list[int] = field(default_factory=lambda: [2, 10]) - #: Smallest and largest number of context events (log-uniform). - context_events: list[int] = field(default_factory=lambda: [16, 2048]) - query_events: int = 256 - #: Momenta and energies of a task are scaled by a factor uniform in - #: ``[1 - energy_jitter, 1 + energy_jitter]``. - energy_jitter: float = 0.2 - steps_per_epoch: int = 1000 - validation_steps: int = 50 - - -@dataclass -class TrainConfig: - #: ``auto`` (CUDA, else MPS, else CPU), ``cpu``, ``cuda`` or ``mps``. - device: str = "auto" - seed: int = 42 - max_epochs: int = 100 - optimizer: str = "adam" - learning_rate: float = 1e-4 - weight_decay: float = 0.0 - #: Exponential learning-rate decay per epoch (1 keeps it constant). - lr_decay: float = 1.0 - #: ``{monitor, patience, min_delta}``; monitor ``train_loss``, ``loss`` - #: (validation) or a validation metric. Null disables early stopping. - early_stopping: dict[str, Any] | None = None - #: Validation quantity that selects ``best.pt``: ``loss`` or a metric. - best_metric: str = "auc" - #: Also keep ``epoch_NNNN.pt`` every this many epochs (0: never). - save_every: int = 0 - #: Continue from this checkpoint. - resume: str | None = None - #: Shorthand for compile, cuda_graphs, bfloat16 and TF32 together. - #: CUDA graphs apply to edge networks; transformer steps run uncaptured. - fast: bool = False - compile: bool = False - cuda_graphs: bool = False - #: ``float32`` or ``bfloat16`` (autocast forward). - precision: str = "float32" - allow_tf32: bool = False - - def __post_init__(self) -> None: - if self.fast: - self.compile = self.cuda_graphs = self.allow_tf32 = True - self.precision = "bfloat16" - - -@dataclass -class Config: - name: str - data: DataConfig - model: ModelConfig = field(default_factory=ModelConfig) - task: TaskConfig = field(default_factory=TaskConfig) - train: TrainConfig = field(default_factory=TrainConfig) - #: Checkpoints, history and outputs; ``outputs/`` by default. - run_dir: str | None = None - #: evaluate / predict / export: checkpoint (default ``run_dir/best.pt``), - #: split, and output file (default under ``run_dir``). - checkpoint: str | None = None - split: str = "test" - output: str | None = None - - def __post_init__(self) -> None: - self.run_dir = self.run_dir or f"outputs/{self.name}" - if self.task.name == "binary" and self.model.out_size != 1: - raise ValueError("binary classification needs model.out_size: 1") - if self.task.name == "pfn" and self.model.name not in ( - "pfn", - "tabpfn", - "tabicl", - ): - raise ValueError( - "task.name: pfn meta-trains model.name: pfn, tabpfn or tabicl" - ) - for pretraining in ("jepa", "contrastive"): - if (self.model.name == pretraining) != (self.task.name == pretraining): - raise ValueError( - f"model.name: {pretraining} trains task.name: {pretraining}" - ) - if self.split not in ("train", "validation", "test"): - raise ValueError("split must be train, validation or test") - - @property - def path(self) -> Path: - return Path(self.run_dir) - - def to_dict(self) -> dict[str, Any]: - return dataclasses.asdict(self) - - -def _build(cls: type, values: Any, where: str) -> Any: - if not isinstance(values, dict): - raise ValueError(f"{where or 'config'} must be a mapping") - known = {item.name: item for item in dataclasses.fields(cls)} - unknown = sorted(set(values) - set(known)) - if unknown: - raise ValueError(f"unknown config keys in {where or 'config'}: {unknown}") - sections = { - "data": DataConfig, - "model": ModelConfig, - "task": TaskConfig, - "train": TrainConfig, - } - kwargs = {} - for name, value in values.items(): - if cls is Config and name in sections: - value = _build(sections[name], value, name) - elif known[name].type.startswith("float") and isinstance(value, str): - value = float(value) # YAML reads 1e-4 (no dot) as a string - kwargs[name] = value - try: - return cls(**kwargs) - except TypeError as error: - raise ValueError(f"{where or 'config'}: {error}") from None - - -def apply_override(values: dict[str, Any], override: str) -> None: - """Set ``dotted.key=value`` in ``values`` (the value is parsed as YAML).""" - key, separator, text = override.partition("=") - if not separator or not key: - raise ValueError(f"override {override!r} is not key=value") - *parents, name = key.split(".") - target = values - for parent in parents: - target = target.setdefault(parent, {}) - if not isinstance(target, dict): - raise ValueError(f"override {override!r}: {parent} is not a section") - target[name] = yaml.safe_load(text) - - -def load_config(path: str | Path, overrides: list[str] = ()) -> Config: - values = yaml.safe_load(Path(path).read_text()) or {} - for override in overrides: - apply_override(values, override) - return config_from_dict(values) - - -def config_from_dict(values: dict[str, Any]) -> Config: - return _build(Config, values, "") - - -def save_config(config: Config, path: str | Path) -> Path: - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(yaml.safe_dump(config.to_dict(), sort_keys=False)) - return path diff --git a/src/gnn4colliders/contrastive.py b/src/gnn4colliders/contrastive.py deleted file mode 100644 index 4dab0b9d344d2ab4c818418ad11713f687ed2cd0..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/contrastive.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Contrastive pretraining of an object encoder (SimCLR, Chen et al. 2020; -for jets JetCLR, Dillon et al. 2022, and RS3L, Harris et al. 2024). - -Every event is seen in two random views. Each view rotates the event in -phi, flips it in eta (probability 1/2), scales its momenta and energies by -a factor uniform in ``1 +- model.scale_jitter`` and drops each object -other than the first with probability ``model.object_dropout``. The encoder -(``model.backbone``, as for JEPA) embeds each view (the mean of its object -encodings, or the edge network's decoded global state), and a projection MLP -maps the embedding to a unit vector. The NT-Xent loss at temperature -``model.temperature`` pulls the two views of an event together and pushes -them from the other events of the batch (on each rank) apart; labels are not -used. The pretrained encoder (``context_encoder``, as JEPA names it) -initializes a classifier through ``model.pretrained``. -""" - -from __future__ import annotations - -import math -from typing import Any - -import torch -import torch.nn.functional as F -from torch import nn - -from .jepa import event_embedding -from .sequences import augment_tokens - - -# Random augmentations from a generator do not trace; they are cheap. -@torch.compiler.disable -def random_view( - tokens: torch.Tensor, - mask: torch.Tensor, - *, - scale_jitter: float, - object_dropout: float, - generator: torch.Generator | None = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """An augmented view ``(tokens, mask)`` of every event (see the module - documentation).""" - events, device = tokens.shape[0], tokens.device - - def uniform(*shape: int) -> torch.Tensor: - return torch.rand(shape, generator=generator, device=device) - - scale = 1 + scale_jitter * (2 * uniform(events) - 1) - rotation = math.pi * (2 * uniform(events) - 1) - flip = torch.where(uniform(events) < 0.5, -1.0, 1.0) - view = augment_tokens(tokens, mask, scale=scale, rotation=rotation, flip=flip) - dropped = uniform(*mask.shape) < object_dropout - dropped[:, 0] = False # every event keeps an object - return view, mask & ~dropped - - -class Contrastive(nn.Module): - """The encoder and the projection MLP. - - ``forward`` returns ``[events, 3]``: each event's loss, 1 (every event - counts), and the spread of the embeddings over the batch (their standard - deviation, averaged over dimensions; it falls towards zero if the - representation collapses). - """ - - def __init__( - self, - encoder: nn.Module, - *, - temperature: float = 0.1, - scale_jitter: float = 0.1, - object_dropout: float = 0.1, - ) -> None: - super().__init__() - hid_size = encoder.hid_size - self.temperature = temperature - self.scale_jitter, self.object_dropout = scale_jitter, object_dropout - self.context_encoder = encoder - self.projector = nn.Sequential( - nn.Linear(hid_size, hid_size), nn.ReLU(), nn.Linear(hid_size, hid_size) - ) - - def forward(self, batch: Any) -> torch.Tensor: - tokens, mask = batch.tokens, batch.token_mask - # Validation sees the same views every epoch. - generator = ( - None if self.training else torch.Generator(tokens.device).manual_seed(0) - ) - embedded = [ - event_embedding( - self.context_encoder, - *random_view( - tokens, - mask, - scale_jitter=self.scale_jitter, - object_dropout=self.object_dropout, - generator=generator, - ), - ) - for _ in range(2) - ] - spread = embedded[0].float().std(dim=0).mean() - z = F.normalize(self.projector(torch.cat(embedded)).float(), dim=1) - events = tokens.shape[0] - similarity = (z @ z.T) / self.temperature - similarity = similarity.masked_fill( - torch.eye(2 * events, dtype=torch.bool, device=z.device), float("-inf") - ) - # The other view of each event is its positive. - positive = torch.arange(2 * events, device=z.device).roll(events) - loss = F.cross_entropy(similarity, positive, reduction="none") - loss = loss.view(2, events).mean(dim=0) - return torch.stack((loss, torch.ones_like(loss), spread.expand_as(loss)), dim=1) diff --git a/src/gnn4colliders/cuda_graphs.py b/src/gnn4colliders/cuda_graphs.py deleted file mode 100644 index ba3599a1c94562619b4de435ebcd072e1ec7308a..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/cuda_graphs.py +++ /dev/null @@ -1,100 +0,0 @@ -"""CUDA graph replay of complete training steps (``train.cuda_graphs``). - -The edge network launches about a thousand small kernels per step, so eager -training is bound by host launch overhead rather than GPU work. Capturing -forward, backward and optimizer step once per batch shape and replaying it -removes that overhead. Graph batches are padded into capacity buckets -(:class:`~gnn4colliders.graphs.StaticGraphBatch`) so a few captures serve a -whole epoch. - -Results agree with eager training to floating-point rounding, not bitwise: -matrix multiplications see padded row counts and the optimizer runs in its -``capturable`` form with device-tensor learning rates. -""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass -from typing import Any - -import torch - -from .graphs import GraphBatch, StaticGraphBatch, graph_batch_capacity - - -@dataclass -class _Capture: - graph: torch.cuda.CUDAGraph - buffers: StaticGraphBatch - loss: torch.Tensor - logits: torch.Tensor - - -def make_capturable(optimizer: torch.optim.Optimizer) -> None: - """Prepare an Adam optimizer for capture (idempotent). - - Learning rates become device tensors, which schedulers update in place, - so captured steps follow the schedule without being captured again. - """ - for group in optimizer.param_groups: - device = group["params"][0].device - group["capturable"] = True - if not torch.is_tensor(group["lr"]): - group["lr"] = torch.tensor(float(group["lr"]), device=device) - for parameter in group["params"]: - state = optimizer.state.get(parameter, {}) - if torch.is_tensor(state.get("step")): - state["step"] = state["step"].to(device) - - -class CudaGraphTrainStep: - """Train on graph batches by replaying one captured step per bucket. - - ``forward(batch)`` must not synchronize with the host. - """ - - def __init__( - self, - forward: Callable[[GraphBatch], torch.Tensor], - task: Any, - optimizer: torch.optim.Optimizer, - *, - node_multiple: int = 512, - edge_multiple: int = 4096, - ) -> None: - if not isinstance(optimizer, (torch.optim.Adam, torch.optim.AdamW)): - raise ValueError("CUDA graph training supports the Adam optimizers") - make_capturable(optimizer) - self.forward, self.task, self.optimizer = forward, task, optimizer - self.node_multiple, self.edge_multiple = node_multiple, edge_multiple - self._captures: dict[tuple[int, int, int], _Capture] = {} - - def _capture(self, key: tuple[int, int, int], batch: GraphBatch) -> _Capture: - # A restored checkpoint may have replaced the optimizer settings. - make_capturable(self.optimizer) - buffers = StaticGraphBatch(batch, key) - buffers.load(batch) - # Gradients allocated during capture belong to this graph's pool. - self.optimizer.zero_grad(set_to_none=True) - torch.cuda.synchronize() - graph = torch.cuda.CUDAGraph() - # Loader pin-memory threads keep using CUDA while this thread captures. - with torch.cuda.graph(graph, capture_error_mode="thread_local"): - logits = self.forward(buffers.batch)[: buffers.graphs] - loss = self.task.loss(logits, buffers.batch.labels, buffers.batch.weight) - loss.backward() - self.optimizer.step() - capture = _Capture(graph, buffers, loss, logits) - self._captures[key] = capture - return capture - - def __call__(self, batch: GraphBatch) -> tuple[torch.Tensor, torch.Tensor]: - """Train on ``batch``; return copies of its loss and logits.""" - key = graph_batch_capacity( - batch, node_multiple=self.node_multiple, edge_multiple=self.edge_multiple - ) - capture = self._captures.get(key) or self._capture(key, batch) - capture.buffers.load(batch) - capture.graph.replay() - return capture.loss.detach().clone(), capture.logits.detach().clone() diff --git a/src/gnn4colliders/data.py b/src/gnn4colliders/data.py deleted file mode 100644 index 4c83334b1e17826c1abecac533297b6695294827..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/data.py +++ /dev/null @@ -1,530 +0,0 @@ -"""ROOT event files to packed per-object features. - -``prepare`` reads ROOT trees in chunks and writes :class:`PackedEvents`: one -row per reconstructed object (jet, electron, muon, photon, MET) and per-event -labels, folds and weights. Graphs, token sequences or any other model input -are built from these rows at batch time. - -Object features follow the pretraining schema of Ho et al. (2026). Each entry -of ``feature_branches`` is one output column and lists one branch (or a -constant) per object type; ``CALC_E`` is ``pt * cosh(eta)`` and ``NODE_TYPE`` -the object-type index. ``object_types`` says whether a type is a ``vector`` -(jagged) or a ``single`` (one per event) branch. -""" - -from __future__ import annotations - -import hashlib -import multiprocessing -from collections.abc import Callable, Mapping, Sequence -from concurrent.futures import ProcessPoolExecutor -from dataclasses import dataclass, fields, is_dataclass, replace -from numbers import Number -from pathlib import Path -from typing import TYPE_CHECKING, Any - -import awkward as ak -import numpy as np -import torch -import uproot - -if TYPE_CHECKING: - from .config import DataConfig - -#: Columns of ``PackedEvents.object_features`` for the standard schema. -NODE_FEATURE_NAMES = ("pt", "eta", "phi", "energy", "btag", "charge", "node_type") -CACHE_VERSION = 2 -#: Events read from ROOT per chunk; bounds peak memory per worker. -CHUNK_EVENTS = 100_000 - -_CALC_E = "CALC_E" -_NODE_TYPE = "NODE_TYPE" - - -@dataclass(frozen=True) -class PackedEvents: - """Prepared events in a flat tensor layout. - - Rows ``object_offsets[i]:object_offsets[i + 1]`` of ``object_features`` - are the objects of event ``i``. ``file_index`` and ``entry_index`` locate - each event in ``source_files``. - """ - - object_features: torch.Tensor - object_offsets: torch.Tensor - labels: torch.Tensor - global_features: torch.Tensor | None - fold: torch.Tensor - weight: torch.Tensor - file_index: torch.Tensor - entry_index: torch.Tensor - source_files: tuple[str, ...] - tree_name: str - - def __post_init__(self) -> None: - count = len(self) - if self.object_features.ndim != 2: - raise ValueError("object_features must have shape [objects, features]") - if int(self.object_offsets[-1]) != self.object_features.shape[0]: - raise ValueError("object_offsets do not cover object_features") - for name in ("labels", "fold", "weight", "file_index", "entry_index"): - if getattr(self, name).shape[0] != count: - raise ValueError(f"{name} must have one entry per event") - if self.global_features is not None and self.global_features.shape[0] != count: - raise ValueError("global_features must have one row per event") - - def __len__(self) -> int: - return int(self.object_offsets.numel()) - 1 - - @property - def num_features(self) -> int: - return int(self.object_features.shape[1]) - - @property - def num_global_features(self) -> int: - return 0 if self.global_features is None else int(self.global_features.shape[1]) - - def gather_objects( - self, indices: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """Object rows and per-event object counts of ``indices``, in order.""" - starts = self.object_offsets[indices] - counts = self.object_offsets[indices + 1] - starts - total = int(counts.sum()) - first = torch.cumsum(counts, 0) - counts - rows = torch.arange(total) + torch.repeat_interleave( - starts - first, counts, output_size=total - ) - return self.object_features[rows], counts - - def split_indices(self, folds: Sequence[int]) -> torch.Tensor: - """Indices of the events whose fold is in ``folds``, in stored order.""" - selected = torch.tensor(sorted(folds), dtype=self.fold.dtype) - return torch.nonzero(torch.isin(self.fold, selected)).flatten() - - def balanced_subset( - self, indices: torch.Tensor, count: int, seed: int - ) -> torch.Tensor: - """Up to ``count // classes`` random events of each class of - ``indices``, in stored order. Smaller counts give subsets of larger - ones with the same seed.""" - labels = self.labels[indices] - classes = labels.unique() - generator = torch.Generator().manual_seed(seed) - chosen = [] - for label in classes: - members = indices[labels == label] - order = torch.randperm(len(members), generator=generator) - chosen.append(members[order[: count // len(classes)]]) - return torch.cat(chosen).sort().values - - @staticmethod - def random_subset(indices: torch.Tensor, count: int, seed: int) -> torch.Tensor: - """Up to ``count`` random events of ``indices`` (keeping the class - proportions on average), in stored order.""" - generator = torch.Generator().manual_seed(seed) - order = torch.randperm(len(indices), generator=generator) - return indices[order[:count]].sort().values - - -def pack_events( - objects: Sequence[torch.Tensor], - labels: Sequence[Any] | torch.Tensor, - *, - weights: Sequence[float] | torch.Tensor | None = None, - folds: Sequence[int] | torch.Tensor | None = None, - global_features: torch.Tensor | None = None, -) -> PackedEvents: - """Pack in-memory per-event ``[objects, features]`` tensors (for tests).""" - count = len(objects) - counts = torch.tensor([int(item.shape[0]) for item in objects]) - return PackedEvents( - object_features=torch.cat([item.to(torch.float32) for item in objects]), - object_offsets=torch.cat((torch.zeros(1, dtype=torch.long), counts.cumsum(0))), - labels=torch.as_tensor(labels), - global_features=global_features, - fold=torch.zeros(count, dtype=torch.long) - if folds is None - else torch.as_tensor(folds, dtype=torch.long), - weight=torch.ones(count) - if weights is None - else torch.as_tensor(weights, dtype=torch.float32), - file_index=torch.zeros(count, dtype=torch.long), - entry_index=torch.arange(count), - source_files=("memory",), - tree_name="events", - ) - - -def concatenate_events(parts: Sequence[PackedEvents]) -> PackedEvents: - """Concatenate parts that index the same ``source_files``, in order.""" - first = parts[0] - if any(part.source_files != first.source_files for part in parts): - raise ValueError("cannot concatenate events from different source files") - offsets, base = [first.object_offsets], first.object_offsets[-1] - for part in parts[1:]: - offsets.append(part.object_offsets[1:] + base) - base = base + part.object_offsets[-1] - - def cat(name: str) -> torch.Tensor | None: - values = [getattr(part, name) for part in parts] - return None if values[0] is None else torch.cat(values) - - return PackedEvents( - object_features=cat("object_features"), - object_offsets=torch.cat(offsets), - labels=cat("labels"), - global_features=cat("global_features"), - fold=cat("fold"), - weight=cat("weight"), - file_index=cat("file_index"), - entry_index=cat("entry_index"), - source_files=first.source_files, - tree_name=first.tree_name, - ) - - -def save_events(events: PackedEvents, path: str | Path) -> Path: - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - payload: dict[str, Any] = {"cache_schema_version": CACHE_VERSION} - for item in fields(PackedEvents): - value = getattr(events, item.name) - payload[item.name] = value.contiguous() if torch.is_tensor(value) else value - payload["source_files"] = list(events.source_files) - torch.save(payload, path) - return path - - -def load_events(path: str | Path, *, mmap: bool = False) -> PackedEvents: - path = Path(path) - if not path.is_file(): - raise FileNotFoundError(f"event cache does not exist: {path} (run prepare)") - payload = torch.load(str(path), map_location="cpu", weights_only=True, mmap=mmap) - if payload.get("cache_schema_version") != CACHE_VERSION: - raise ValueError(f"{path} has an unsupported cache version; prepare again") - values = {item.name: payload.get(item.name) for item in fields(PackedEvents)} - values["source_files"] = tuple(values["source_files"]) - return PackedEvents(**values) - - -def map_tensors(value: Any, function: Callable[[torch.Tensor], torch.Tensor]) -> Any: - """Apply ``function`` to every tensor field of a (nested) dataclass.""" - changes = {} - for item in fields(value): - field_value = getattr(value, item.name) - if torch.is_tensor(field_value): - changes[item.name] = function(field_value) - elif is_dataclass(field_value): - changes[item.name] = map_tensors(field_value, function) - return replace(value, **changes) - - -# --- Derived branches --------------------------------------------------------- - - -@dataclass(frozen=True) -class DerivedBranch: - """A branch computed from stored branches before feature construction. - - ATLAS Open Data stores electrons and muons in one ``lep_*`` collection - and a continuous b-tagging score; derived branches name the per-type - quantities the feature schema expects:: - - ele_pt: {source: lep_pt, where: {lep_type: 11}} - jet_btag: {source: jet_MV2c10, above: 0.8244273} - - ``where`` keeps the objects whose parallel branches equal the values; - ``above`` turns values into ``1.0`` above the threshold and ``0.0`` - otherwise (selection is applied first). - """ - - name: str - source: str - where: tuple[tuple[str, Any], ...] = () - above: float | None = None - - @property - def inputs(self) -> tuple[str, ...]: - return (self.source, *(branch for branch, _ in self.where)) - - -def parse_derived_branches( - specification: Mapping[str, Mapping[str, Any]] | None, -) -> tuple[DerivedBranch, ...]: - derived = [] - for name, options in (specification or {}).items(): - unknown = set(options) - {"source", "where", "above"} - if unknown or "source" not in options: - raise ValueError( - f"derived branch {name!r} needs a source and only where/above" - ) - above = options.get("above") - derived.append( - DerivedBranch( - name=str(name), - source=str(options["source"]), - where=tuple( - (str(k), v) for k, v in (options.get("where") or {}).items() - ), - above=None if above is None else float(above), - ) - ) - names = {item.name for item in derived} - for item in derived: - if names.intersection(item.inputs): - raise ValueError( - f"derived branch {item.name!r} depends on a derived branch" - ) - return tuple(derived) - - -def apply_derived_branches( - arrays: Any, derived: Sequence[DerivedBranch] -) -> Mapping[str, Any]: - result = {name: arrays[name] for name in arrays.fields} - for item in derived: - values = arrays[item.source] - if item.where: - mask = None - for branch, expected in item.where: - matches = arrays[branch] == expected - mask = matches if mask is None else mask & matches - values = values[mask] - if item.above is not None: - values = ak.values_astype(values > item.above, np.float32) - result[item.name] = values - return result - - -# --- Object features ---------------------------------------------------------- - - -def build_object_features( - arrays: Mapping[str, Any], - num_events: int, - feature_branches: Sequence[Any], - object_types: Sequence[str], - scales: Sequence[float], -) -> tuple[torch.Tensor, torch.Tensor]: - """Object rows of ``num_events`` events and the object count per event. - - Within an event, objects are ordered by type (the order of - ``object_types``) and then as stored. - """ - first_column = feature_branches[0] - if len(first_column) != len(object_types): - raise ValueError("one branch specification is required per object type") - lengths = [] - for branch, object_type in zip(first_column, object_types): - if object_type == "single": - lengths.append(np.ones(num_events, dtype=np.int64)) - elif object_type == "vector": - lengths.append(np.asarray(ak.num(arrays[branch], axis=1), dtype=np.int64)) - else: - raise ValueError(f"unknown object type: {object_type!r}") - per_type = np.stack(lengths, axis=1) - counts = per_type.sum(axis=1) - event_start = np.cumsum(counts) - counts - total = int(counts.sum()) - - # Destination row of every object of each type, in event-major order. - positions = [] - for type_index in range(len(object_types)): - type_lengths = per_type[:, type_index] - base = event_start + per_type[:, :type_index].sum(axis=1) - local = np.arange(int(type_lengths.sum())) - np.repeat( - np.cumsum(type_lengths) - type_lengths, type_lengths - ) - positions.append(np.repeat(base, type_lengths) + local) - - columns: list[torch.Tensor] = [] - for column_index, specification in enumerate(feature_branches): - if specification == _CALC_E: - columns.append(columns[0] * torch.cosh(columns[1])) - continue - column = np.empty(total, dtype=np.float32) - if specification == _NODE_TYPE: - for type_index, rows in enumerate(positions): - column[rows] = type_index - elif len(specification) != len(object_types): - raise ValueError(f"feature column {column_index} needs one entry per type") - else: - for type_index, branch in enumerate(specification): - if isinstance(branch, Number): - column[positions[type_index]] = branch - continue - values = arrays[branch] - if object_types[type_index] == "vector": - values = ak.flatten(values, axis=1) - column[positions[type_index]] = np.asarray(ak.to_numpy(values)) - columns.append(torch.from_numpy(column)) - features = torch.stack(columns, dim=1) - scale = torch.as_tensor(scales, dtype=torch.float32) - if scale.shape != (features.shape[1],): - raise ValueError("scales must contain one value per feature column") - return features * scale, torch.from_numpy(counts) - - -# --- ROOT input --------------------------------------------------------------- - - -def _sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as stream: - for chunk in iter(lambda: stream.read(1 << 20), b""): - digest.update(chunk) - return digest.hexdigest() - - -def resolve_files(config: DataConfig) -> list[tuple[Path, Any]]: - """``(path, label)`` of every input file, downloading Hugging Face files. - - Hub downloads are pinned to ``hf_revision`` and checked against - ``size_bytes`` and ``sha256`` when given. - """ - resolved = [] - for entry in config.files: - path = Path(entry["path"]) - if config.hf_repo: - from huggingface_hub import hf_hub_download - - path = Path( - hf_hub_download( - repo_id=config.hf_repo, - filename=str(entry["path"]), - repo_type="dataset", - revision=config.hf_revision, - ) - ) - if "size_bytes" in entry and path.stat().st_size != int(entry["size_bytes"]): - raise ValueError(f"size mismatch for {path}") - if "sha256" in entry and _sha256(path) != entry["sha256"]: - raise ValueError(f"SHA-256 mismatch for {path}") - resolved.append((path, entry["label"])) - if not resolved: - raise ValueError("data.files must list at least one ROOT file") - return resolved - - -def _branch_names(config: DataConfig, label: Any) -> list[str]: - """Stored branches needed for the features, folds, weights and labels.""" - derived = { - item.name: item for item in parse_derived_branches(config.derived_branches) - } - names: list[str] = [] - - def visit(value: Any) -> None: - if isinstance(value, (list, tuple)): - for item in value: - visit(item) - elif isinstance(value, str) and value not in (_CALC_E, _NODE_TYPE): - for stored in derived[value].inputs if value in derived else (value,): - if stored not in names: - names.append(stored) - - visit([config.feature_branches, config.global_features, config.fold_var]) - visit([config.weight_var or [], label if isinstance(label, str) else []]) - return names - - -def _per_event(arrays: Any, branch: str) -> np.ndarray: - values = arrays[branch] - if isinstance(values, ak.Array) and values.ndim != 1: - raise ValueError(f"branch {branch!r} is not one value per event") - return np.asarray(ak.to_numpy(values)) - - -def _labels(label: Any, arrays: Any, count: int) -> torch.Tensor: - if isinstance(label, str): - values = _per_event(arrays, label) - return torch.from_numpy( - values.astype(np.int64 if values.dtype.kind in "iu" else values.dtype) - ) - return torch.full((count,), label, dtype=torch.as_tensor(label).dtype) - - -def read_chunk( - config: DataConfig, - files: Sequence[tuple[Path, Any]], - file_index: int, - start: int, - stop: int, -) -> PackedEvents: - """Prepare entries ``[start, stop)`` of one input file.""" - path, label = files[file_index] - with uproot.open(path) as root_file: - arrays = root_file[config.tree_name].arrays( - _branch_names(config, label), entry_start=start, entry_stop=stop - ) - arrays = apply_derived_branches( - arrays, parse_derived_branches(config.derived_branches) - ) - count = stop - start - features, counts = build_object_features( - arrays, count, config.feature_branches, config.object_types, config.scales - ) - # Exact integer arithmetic: event numbers exceed float32 precision. - fold = _per_event(arrays, config.fold_var).astype(np.int64) - if config.fold_modulus: - fold = fold % config.fold_modulus - weight = ( - _per_event(arrays, config.weight_var).astype(np.float32) - if config.weight_var - else np.ones(count, dtype=np.float32) - ) - global_features = None - if config.global_features: - global_features = torch.from_numpy( - np.stack( - [ - _per_event(arrays, name).astype(np.float32) - for name in config.global_features - ], - axis=1, - ) - ) - return PackedEvents( - object_features=features, - object_offsets=torch.cat((torch.zeros(1, dtype=torch.long), counts.cumsum(0))), - labels=_labels(label, arrays, count), - global_features=global_features, - fold=torch.from_numpy(fold), - weight=torch.from_numpy(weight), - file_index=torch.full((count,), file_index, dtype=torch.long), - entry_index=torch.arange(start, stop), - source_files=tuple(str(path) for path, _ in files), - tree_name=config.tree_name, - ) - - -def _read_chunk(arguments: tuple[Any, ...]) -> PackedEvents: - return read_chunk(*arguments) - - -def prepare(config: DataConfig, *, workers: int | None = None) -> Path: - """Write the event cache of ``config`` to ``config.cache``. - - Chunks of at most :data:`CHUNK_EVENTS` entries are prepared in - ``workers`` processes (``config.prepare_workers`` by default) and - concatenated in file and entry order. - """ - files = resolve_files(config) - chunks = [] - for file_index, (path, _) in enumerate(files): - with uproot.open(path) as root_file: - entries = int(root_file[config.tree_name].num_entries) - chunks.extend( - (config, files, file_index, start, min(start + CHUNK_EVENTS, entries)) - for start in range(0, entries, CHUNK_EVENTS) - ) - if not chunks: - raise ValueError("the input files contain no events") - workers = config.prepare_workers if workers is None else workers - if workers > 1: - context = multiprocessing.get_context("spawn") - with ProcessPoolExecutor(max_workers=workers, mp_context=context) as pool: - parts = list(pool.map(_read_chunk, chunks)) - else: - parts = [_read_chunk(chunk) for chunk in chunks] - return save_events(concatenate_events(parts), config.cache) diff --git a/src/gnn4colliders/distributed.py b/src/gnn4colliders/distributed.py deleted file mode 100644 index dbdd91fdd76366421ef1f958f045c4557d443d6a..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/distributed.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Data-parallel (DDP) execution under ``torchrun``. - -A process launched by ``torchrun`` with ``WORLD_SIZE > 1`` joins the -``env://`` process group; anything else runs as a single process. -""" - -from __future__ import annotations - -import os -from dataclasses import dataclass -from datetime import timedelta - -import torch -import torch.distributed as dist -from torch import nn - - -@dataclass(frozen=True) -class DistributedContext: - enabled: bool = False - rank: int = 0 - local_rank: int = 0 - world_size: int = 1 - device: torch.device = torch.device("cpu") - - @property - def is_main_process(self) -> bool: - return self.rank == 0 - - -def resolve_device(device: str) -> torch.device: - if device == "auto": - if torch.cuda.is_available(): - return torch.device("cuda") - return torch.device("mps" if torch.backends.mps.is_available() else "cpu") - return torch.device(device) - - -def synchronize(device: torch.device) -> None: - """Wait for the work queued on ``device``, for timing.""" - if device.type == "cuda": - torch.cuda.synchronize(device) - elif device.type == "mps": - torch.mps.synchronize() - - -def initialize(device: str = "auto") -> DistributedContext: - world_size = int(os.environ.get("WORLD_SIZE", "1")) - requested = resolve_device(device) - if world_size < 2: - return DistributedContext(device=requested) - rank = int(os.environ["RANK"]) - local_rank = int(os.environ.get("LOCAL_RANK", rank)) - if requested.type == "cuda": - torch.cuda.set_device(local_rank) - requested = torch.device("cuda", local_rank) - if not dist.is_initialized(): - dist.init_process_group( - backend="nccl" if requested.type == "cuda" else "gloo", - timeout=timedelta(minutes=30), - ) - return DistributedContext(True, rank, local_rank, world_size, requested) - - -def finalize(context: DistributedContext) -> None: - if context.enabled and dist.is_initialized(): - dist.destroy_process_group() - - -def wrap_model(model: nn.Module, context: DistributedContext) -> nn.Module: - model = model.to(context.device) - if not context.enabled: - return model - if context.device.type == "cuda": - return nn.parallel.DistributedDataParallel( - model, device_ids=[context.local_rank] - ) - return nn.parallel.DistributedDataParallel(model) - - -def gather_tensor(tensor: torch.Tensor, context: DistributedContext) -> torch.Tensor: - """Concatenate every rank's tensor along dimension 0 (sizes may differ).""" - if not context.enabled: - return tensor - local = tensor.contiguous() - size = torch.tensor([local.shape[0]], device=local.device) - sizes = [torch.zeros_like(size) for _ in range(context.world_size)] - dist.all_gather(sizes, size) - longest = max(int(item) for item in sizes) - padded = local.new_zeros((longest, *local.shape[1:])) - padded[: local.shape[0]] = local - gathered = [torch.empty_like(padded) for _ in range(context.world_size)] - dist.all_gather(gathered, padded) - return torch.cat([part[: int(n)] for part, n in zip(gathered, sizes)]) - - -def broadcast_bool(value: bool, context: DistributedContext) -> bool: - """Rank 0's ``value`` on every rank.""" - if not context.enabled: - return value - tensor = torch.tensor([int(value)], device=context.device) - dist.broadcast(tensor, src=0) - return bool(tensor.item()) - - -def barrier(context: DistributedContext) -> None: - if context.enabled: - dist.barrier() diff --git a/src/gnn4colliders/export.py b/src/gnn4colliders/export.py deleted file mode 100644 index 95bab25f8c0c0203d9b0eaf56737602a6108bb03..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/export.py +++ /dev/null @@ -1,208 +0,0 @@ -"""ONNX export, checked with ONNX Runtime. - -The exported graph takes plain tensors. For edge networks: ``node_features``, -``edge_features``, ``edge_src``, ``edge_dst``, ``node_batch`` and -``global_features`` (for models without global features, the per-graph node -count as one column). For the transformer and Deep Sets: ``tokens`` ``[events, objects, -7]``, ``token_mask`` (true for real objects) and, if the model has them, -``global_features``. Requires the ``onnx`` extra. -""" - -from __future__ import annotations - -import json -from pathlib import Path -from types import SimpleNamespace -from typing import Any - -import torch -from torch import nn - -from .graphs import GraphBatch -from .models import EdgeNetwork, FineTunedEdgeNetwork -from .sequences import SequenceBatch - -OPSET = 17 -GRAPH_INPUTS = { - "node_features": {0: "nodes"}, - "edge_features": {0: "edges"}, - "edge_src": {0: "edges"}, - "edge_dst": {0: "edges"}, - "node_batch": {0: "nodes"}, - "global_features": {0: "graphs"}, -} -SEQUENCE_INPUTS = { - "tokens": {0: "events", 1: "objects"}, - "token_mask": {0: "events", 1: "objects"}, - "global_features": {0: "events"}, -} - - -def _membership(groups: torch.Tensor, count: int, values: torch.Tensor) -> torch.Tensor: - # ``index_add`` exports to scatter operations whose handling of repeated - # indices is unreliable; a dense membership matrix is unambiguous. - labels = torch.arange(count, device=groups.device).unsqueeze(1) - return (labels == groups.unsqueeze(0)).to(values.dtype) - - -def _sum_by_group( - values: torch.Tensor, groups: torch.Tensor, count: int -) -> torch.Tensor: - return _membership(groups, count, values) @ values - - -def _mean_by_group( - values: torch.Tensor, groups: torch.Tensor, count: int -) -> torch.Tensor: - membership = _membership(groups, count, values) - return (membership @ values) / membership.sum(dim=1, keepdim=True).clamp_min(1) - - -class ExportAdapter(nn.Module): - """The forward pass of an (fine-tuned) edge network on plain tensors.""" - - def __init__(self, model: nn.Module) -> None: - super().__init__() - if isinstance(model, FineTunedEdgeNetwork): - self.backbone = model.backbone - elif isinstance(model, EdgeNetwork): - self.backbone = model - else: - raise TypeError(f"no graph export for {type(model).__name__}") - self.classifier = model.classifier - - def forward( - self, - node_features: torch.Tensor, - edge_features: torch.Tensor, - edge_src: torch.Tensor, - edge_dst: torch.Tensor, - node_batch: torch.Tensor, - global_features: torch.Tensor, - ) -> torch.Tensor: - net = self.backbone - graphs = global_features.shape[0] - edge_batch = node_batch[edge_src] - node_h = net.node_encoder(node_features) - edge_h = net.edge_encoder(edge_features) - global_h = net.global_encoder(global_features) - for _ in range(net.n_proc_steps): - edge_h = net.edge_update( - torch.cat( - (edge_h, node_h[edge_src], node_h[edge_dst], global_h[edge_batch]), - 1, - ) - ) - messages = _sum_by_group(edge_h, edge_dst, node_h.shape[0]) - node_h = net.node_update( - torch.cat((node_h, messages, global_h[node_batch]), 1) - ) - global_h = net.global_update( - torch.cat( - ( - global_h, - _mean_by_group(node_h, node_batch, graphs), - _mean_by_group(edge_h, edge_batch, graphs), - ), - 1, - ) - ) - return self.classifier(net.global_decoder(global_h)) - - -class SequenceExportAdapter(nn.Module): - """The forward pass of a token model on plain tensors.""" - - def __init__(self, model: nn.Module) -> None: - super().__init__() - self.model = model - - def forward( - self, - tokens: torch.Tensor, - token_mask: torch.Tensor, - global_features: torch.Tensor | None = None, - ) -> torch.Tensor: - batch = SimpleNamespace( - tokens=tokens, token_mask=token_mask, global_features=global_features - ) - return self.model(batch) - - -def export_inputs(batch: GraphBatch | SequenceBatch) -> tuple[torch.Tensor, ...]: - if isinstance(batch, SequenceBatch): - if batch.global_features is None: - return batch.tokens, batch.token_mask - return batch.tokens, batch.token_mask, batch.global_features - graph = batch.graph - global_features = batch.global_features - if global_features is None: - global_features = graph.node_counts[:, None].to(graph.node_features.dtype) - return ( - graph.node_features, - graph.edge_features, - graph.edge_src, - graph.edge_dst, - graph.node_batch, - global_features, - ) - - -def export_onnx( - model: nn.Module, - example: GraphBatch | SequenceBatch, - path: str | Path, - *, - metadata: dict[str, Any] | None = None, -) -> Path: - """Export ``model``, check the file, and compare ONNX Runtime logits with - PyTorch on ``example``. Writes ``metadata`` next to it as JSON.""" - try: - import onnx - import onnxruntime - except ImportError as error: - raise ImportError( - "ONNX export needs the onnx extra: uv sync --extra onnx" - ) from error - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - model = model.cpu() - if isinstance(example, SequenceBatch): - adapter, axes = SequenceExportAdapter(model), SEQUENCE_INPUTS - else: - adapter, axes = ExportAdapter(model), GRAPH_INPUTS - adapter.eval() - inputs = export_inputs(example.to("cpu")) - names = list(axes)[: len(inputs)] - torch.onnx.export( - adapter, - inputs, - str(path), - input_names=names, - output_names=["logits"], - opset_version=OPSET, - dynamic_axes={ - **{name: axes[name] for name in names}, - "logits": {0: axes[names[-1]][0]}, - }, - # The TorchScript exporter keeps the dynamic axes and needs no - # onnxscript dependency. - dynamo=False, - ) - onnx.checker.check_model(onnx.load(str(path))) - session = onnxruntime.InferenceSession( - str(path), providers=["CPUExecutionProvider"] - ) - feeds = {name: tensor.numpy() for name, tensor in zip(names, inputs)} - with torch.inference_mode(): - expected = adapter(*inputs) - actual = torch.from_numpy(session.run(["logits"], feeds)[0]) - torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5) - info = { - "opset": OPSET, - "inputs": names, - "output": "logits", - **(metadata or {}), - } - path.with_suffix(".json").write_text(json.dumps(info, indent=2) + "\n") - return path diff --git a/src/gnn4colliders/graphs.py b/src/gnn4colliders/graphs.py deleted file mode 100644 index f5df316b8f3fab10d00413b91981355d58a01336..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/graphs.py +++ /dev/null @@ -1,229 +0,0 @@ -"""Fully connected event graphs over packed events. - -A :class:`TensorGraph` stores the nodes and directed edges of many graphs as -flat tensors with per-element graph membership (``node_batch``, -``edge_batch``), the edge-index layout common to PyTorch GNNs. Graphs of -different sizes need no padding; only CUDA graph replay -(:class:`StaticGraphBatch`) pads batches into fixed-size buffers. -""" - -from __future__ import annotations - -import math -from dataclasses import dataclass - -import torch - -from .data import NODE_FEATURE_NAMES, PackedEvents, map_tensors - -ETA = NODE_FEATURE_NAMES.index("eta") -PHI = NODE_FEATURE_NAMES.index("phi") -#: Edge feature columns. -EDGE_FEATURE_NAMES = ("deta", "dphi", "dR") - - -@dataclass(frozen=True) -class TensorGraph: - node_features: torch.Tensor - edge_features: torch.Tensor - edge_src: torch.Tensor - edge_dst: torch.Tensor - node_batch: torch.Tensor - edge_batch: torch.Tensor - node_counts: torch.Tensor - edge_counts: torch.Tensor - - -@dataclass(frozen=True) -class GraphBatch: - """Graphs of a batch of events with their labels, weights and indices. - - ``event_index`` holds the events' indices into the :class:`PackedEvents` - they were collated from. - """ - - graph: TensorGraph - labels: torch.Tensor - weight: torch.Tensor - global_features: torch.Tensor | None - event_index: torch.Tensor - - def to(self, device: torch.device | str, *, non_blocking: bool = False): - return map_tensors(self, lambda t: t.to(device, non_blocking=non_blocking)) - - def pin_memory(self) -> GraphBatch: - return map_tensors(self, torch.Tensor.pin_memory) - - -def fully_connected_edges( - node_counts: torch.Tensor, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Directed all-pairs edges of consecutive graphs of ``node_counts`` nodes. - - Edges are source-major without self-loops, except that a one-node graph - keeps one self-loop (as in the original ROOT-GNN). Returns global - ``(src, dst)`` node indices and the number of edges per graph. - """ - counts = node_counts.to(torch.long) - edge_counts = torch.where(counts > 1, counts * (counts - 1), counts.clamp(max=1)) - total = int(edge_counts.sum()) - graph = torch.repeat_interleave( - torch.arange(counts.numel(), device=counts.device), - edge_counts, - output_size=total, - ) - local = ( - torch.arange(total, device=counts.device) - - (edge_counts.cumsum(0) - edge_counts)[graph] - ) - width = (counts[graph] - 1).clamp(min=1) - src, other = local // width, local % width - dst = torch.where(counts[graph] > 1, other + (other >= src).to(torch.long), src) - offset = (counts.cumsum(0) - counts)[graph] - return src + offset, dst + offset, edge_counts - - -def edge_features( - node_features: torch.Tensor, src: torch.Tensor, dst: torch.Tensor -) -> torch.Tensor: - """``[deta, dphi, dR]`` from ``src`` to ``dst``, with dphi in [-pi, pi].""" - deta = node_features[src, ETA] - node_features[dst, ETA] - dphi = node_features[src, PHI] - node_features[dst, PHI] - dphi = torch.where(dphi > math.pi, dphi - 2 * math.pi, dphi) - dphi = torch.where(dphi < -math.pi, dphi + 2 * math.pi, dphi) - return torch.stack((deta, dphi, torch.sqrt(deta.square() + dphi.square())), dim=1) - - -def build_graph(node_features: torch.Tensor, node_counts: torch.Tensor) -> TensorGraph: - """Fully connected graphs on the device of ``node_features``.""" - counts = node_counts.to(device=node_features.device, dtype=torch.long) - src, dst, edge_counts = fully_connected_edges(counts) - graphs = torch.arange(counts.numel(), device=counts.device) - return TensorGraph( - node_features=node_features, - edge_features=edge_features(node_features, src, dst), - edge_src=src, - edge_dst=dst, - node_batch=torch.repeat_interleave( - graphs, counts, output_size=node_features.shape[0] - ), - edge_batch=torch.repeat_interleave( - graphs, edge_counts, output_size=src.numel() - ), - node_counts=counts, - edge_counts=edge_counts, - ) - - -def collate_graph_batch(events: PackedEvents, indices: torch.Tensor) -> GraphBatch: - nodes, counts = events.gather_objects(indices) - return GraphBatch( - graph=build_graph(nodes, counts), - labels=events.labels[indices], - weight=events.weight[indices], - global_features=None - if events.global_features is None - else events.global_features[indices], - event_index=indices, - ) - - -# --- Fixed-size buffers for CUDA graph replay --------------------------------- - -#: Extra graphs (and the minimum number of padding nodes) per buffer. -PADDING_GRAPHS = 64 - - -def graph_batch_capacity( - batch: GraphBatch, *, node_multiple: int, edge_multiple: int -) -> tuple[int, int, int]: - """``(graphs, nodes, edges)`` capacity bucket of ``batch``. - - Capacities are rounded up to a multiple and leave at least - :data:`PADDING_GRAPHS` padding nodes and one padding edge. - """ - graph = batch.graph - nodes = int(graph.node_features.shape[0]) + PADDING_GRAPHS - edges = int(graph.edge_features.shape[0]) - return ( - int(graph.node_counts.numel()), - -(-nodes // node_multiple) * node_multiple, - (edges // edge_multiple + 1) * edge_multiple, - ) - - -class StaticGraphBatch: - """Device buffers holding any batch of one capacity bucket. - - Real nodes and edges fill a prefix; the rest forms :data:`PADDING_GRAPHS` - extra graphs whose edges connect only padding nodes. Sums, means and - LayerNorm never mix graphs, so the real graphs' outputs do not depend on - the padding; the model returns ``graphs + PADDING_GRAPHS`` rows of which - the last are discarded. Padding is spread over many nodes and graphs - because scatter-adds into a single row serialize on atomic contention. - """ - - def __init__(self, example: GraphBatch, capacity: tuple[int, int, int]) -> None: - graphs, nodes, edges = capacity - graph = example.graph - - def buffer(like: torch.Tensor, rows: int) -> torch.Tensor: - return like.new_zeros((rows, *like.shape[1:])) - - self.graphs, self.nodes, self.edges = graphs, nodes, edges - self.graph = TensorGraph( - node_features=buffer(graph.node_features, nodes), - edge_features=buffer(graph.edge_features, edges), - edge_src=buffer(graph.edge_src, edges), - edge_dst=buffer(graph.edge_dst, edges), - node_batch=buffer(graph.node_batch, nodes), - edge_batch=buffer(graph.edge_batch, edges), - node_counts=buffer(graph.node_counts, graphs + PADDING_GRAPHS), - edge_counts=buffer(graph.edge_counts, graphs + PADDING_GRAPHS), - ) - self._positions = torch.arange( - max(nodes, edges), device=graph.node_features.device - ) - self.batch = GraphBatch( - graph=self.graph, - labels=buffer(example.labels, graphs), - weight=buffer(example.weight, graphs), - global_features=None - if example.global_features is None - else buffer(example.global_features, graphs + PADDING_GRAPHS), - event_index=buffer(example.event_index, 0), - ) - - def load(self, batch: GraphBatch) -> None: - """Copy ``batch`` (already on the buffers' device) into the buffers.""" - source, target = batch.graph, self.graph - nodes = int(source.node_features.shape[0]) - edges = int(source.edge_features.shape[0]) - if int(source.node_counts.numel()) != self.graphs: - raise ValueError("batch has a different number of graphs") - padding_nodes, padding_edges = self.nodes - nodes, self.edges - edges - if padding_nodes < PADDING_GRAPHS or padding_edges < 1: - raise ValueError("batch exceeds the buffer capacity") - for name in ("node_features", "node_batch"): - getattr(target, name)[:nodes].copy_(getattr(source, name)) - for name in ("edge_features", "edge_src", "edge_dst", "edge_batch"): - getattr(target, name)[:edges].copy_(getattr(source, name)) - # Padding edges are self-loops on padding nodes, assigned round-robin - # to the padding graphs; padding node ``k`` belongs to padding graph - # ``k % PADDING_GRAPHS``. Stale padding feature rows stay in padding. - node_slots = self._positions[:padding_nodes] - loops = self._positions[:padding_edges] % padding_nodes + nodes - target.node_batch[nodes:].copy_(node_slots % PADDING_GRAPHS + self.graphs) - target.edge_src[edges:].copy_(loops) - target.edge_dst[edges:].copy_(loops) - target.edge_batch[edges:].copy_((loops - nodes) % PADDING_GRAPHS + self.graphs) - # Padding-graph counts only reach discarded rows; any positive value - # keeps them finite without a host synchronization to count them. - target.node_counts[: self.graphs].copy_(source.node_counts) - target.node_counts[self.graphs :].fill_(1) - target.edge_counts[: self.graphs].copy_(source.edge_counts) - target.edge_counts[self.graphs :].fill_(1) - if batch.global_features is not None: - self.batch.global_features[: self.graphs].copy_(batch.global_features) - self.batch.labels.copy_(batch.labels) - self.batch.weight.copy_(batch.weight) diff --git a/src/gnn4colliders/inference.py b/src/gnn4colliders/inference.py deleted file mode 100644 index 33dcd9bfc2bf18e88d54fe409915d2bd5115a995..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/inference.py +++ /dev/null @@ -1,120 +0,0 @@ -"""Model outputs over a loader, and writing predictions.""" - -from __future__ import annotations - -from collections.abc import Callable -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import numpy as np -import torch -import uproot - -from .data import PackedEvents -from .distributed import DistributedContext, gather_tensor - - -@dataclass(frozen=True) -class Outputs: - """Logits and the matching labels, weights and event indices, in loader - order (concatenated over ranks under DDP).""" - - logits: torch.Tensor - labels: torch.Tensor - weight: torch.Tensor - event_index: torch.Tensor - - -def collect_outputs( - forward: Callable[[Any], torch.Tensor], - loader: Any, - *, - device: torch.device, - context: DistributedContext | None = None, -) -> Outputs: - """Run ``forward`` on every batch of ``loader`` without gradients.""" - parts: dict[str, list[torch.Tensor]] = { - name: [] for name in Outputs.__dataclass_fields__ - } - with torch.inference_mode(): - for batch in loader: - batch = batch.to(device, non_blocking=True) - parts["logits"].append(forward(batch).float()) - parts["labels"].append(batch.labels) - parts["weight"].append(batch.weight) - parts["event_index"].append(batch.event_index) - if not parts["logits"]: - raise ValueError("cannot evaluate an empty loader") - context = context or DistributedContext(device=device) - return Outputs( - **{ - name: gather_tensor(torch.cat(values), context) - for name, values in parts.items() - } - ) - - -def write_predictions( - path: str | Path, outputs: Outputs, task: Any, events: PackedEvents -) -> Path: - """Write logits, scores, predictions, labels, weights, folds and each - event's source file and entry to a compressed NumPy archive.""" - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - index = outputs.event_index.cpu() - scores = task.scores(outputs.logits).cpu() - np.savez_compressed( - path, - logits=outputs.logits.cpu().numpy(), - scores=scores.numpy(), - predictions=task.predictions(scores).numpy(), - labels=outputs.labels.cpu().numpy(), - weight=outputs.weight.cpu().numpy(), - fold=events.fold[index].numpy(), - source_file=np.asarray(events.source_files)[events.file_index[index].numpy()], - entry=events.entry_index[index].numpy(), - ) - return path - - -def write_root_scores( - scores: np.ndarray, - entries: np.ndarray, - *, - input_path: str | Path, - output_path: str | Path, - tree_name: str, - prefix: str = "score", -) -> Path: - """Copy a ROOT tree and add score branches aligned by entry. - - Entries without a score get NaN and ``selection_pass = 0``. Multiclass - scores become one ``_class_`` branch per class. - """ - with uproot.open(input_path) as root_file: - tree = root_file[tree_name] - columns, count = tree.arrays(library="ak"), int(tree.num_entries) - entries = np.asarray(entries, dtype=np.int64) - if entries.size and (entries.min() < 0 or entries.max() >= count): - raise ValueError("entries lie outside the ROOT tree") - scores = np.asarray(scores, dtype=np.float32) - if scores.ndim == 1: - scores = scores[:, None] - names = ( - [prefix] - if scores.shape[1] == 1 - else [f"{prefix}_class_{index}" for index in range(scores.shape[1])] - ) - output = {name: columns[name] for name in columns.fields} - output["selection_pass"] = np.zeros(count, dtype=np.int32) - output["selection_pass"][entries] = 1 - for column, name in enumerate(names): - branch = np.full(count, np.nan, dtype=np.float32) - branch[entries] = scores[:, column] - output[name] = branch - target = Path(output_path) - target.parent.mkdir(parents=True, exist_ok=True) - with uproot.recreate(target) as root_file: - root_file[tree_name] = output - return target diff --git a/src/gnn4colliders/jepa.py b/src/gnn4colliders/jepa.py deleted file mode 100644 index 8ebf307b0e784e862a0b532fae10b402676c9dfd..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/jepa.py +++ /dev/null @@ -1,288 +0,0 @@ -"""JEPA pretraining of an object encoder (I-JEPA, Assran et al. 2023). - -The encoder (``model.backbone``) is the object transformer or the edge -network, whose node states are the object encodings. Some objects of each -event are hidden. A context encoder sees the others; a -predictor, told only the type of each hidden object and its index among the -objects of that type (their pt order as stored), predicts the hidden -objects' encodings by a target encoder that sees the whole event. The target -encoder is an exponential moving average of the context encoder and gets no -gradient. The loss is the mean squared error between predicted and -layer-normalized target encodings; labels are not used. - -With ``model.jepa_target: features`` the predictor instead predicts the -hidden objects' kinematics, ``log(1 + pt)``, eta, ``cos phi``, ``sin phi``, -``log(1 + energy)``, b-tag and charge, as in masked particle modelling -(Golling et al. 2024); there is no target encoder. Both differ only in -whether the target is an encoding or the input. - -The pretrained context encoder initializes a classifier of the same -backbone (``model.name: transformer`` or ``edge_network`` with -``model.pretrained``). - -For the edge network, an object's encoding is the mean of its node states -after the encoder and after each processing step (as data2vec averages -layers). Its last node states alone are nearly the same for all objects of -an event at initialization (message passing over fully connected graphs -without residual connections smooths them), which makes predicting them -trivial. Its decoded global state, which its classifier reads, joins the -predictor's input as one more token, so it is pretrained too. -""" - -from __future__ import annotations - -import copy -from types import SimpleNamespace -from typing import Any - -import torch -from torch import nn - -from .graphs import build_graph -from .models import NODE_TYPE, EdgeNetwork, EncoderBlock, _mean_by_group -from .sequences import BTAG, CHARGE, ENERGY, ETA, PHI, PT - -#: Objects of one type beyond this index share the last index embedding. -MAX_TYPE_INDEX = 16 - - -def type_index( - node_type: torch.Tensor, mask: torch.Tensor, num_node_types: int -) -> torch.Tensor: - """Index of every object among the real objects of its type, in stored - order, capped at ``MAX_TYPE_INDEX - 1``.""" - one_hot = nn.functional.one_hot(node_type, num_node_types) - counts = (one_hot * mask.unsqueeze(-1)).cumsum(dim=1) - index = counts.gather(-1, node_type.unsqueeze(-1)).squeeze(-1) - 1 - return index.clamp(0, MAX_TYPE_INDEX - 1) - - -# Random sizes from a generator do not trace with dynamic shapes; the masking -# is cheap, so it runs outside torch.compile. -@torch.compiler.disable -def sample_targets( - mask: torch.Tensor, fraction: float, generator: torch.Generator | None = None -) -> torch.Tensor: - """Hide ``round(fraction * objects)`` random real objects of each event, - at least one, but always leave one visible (one-object events hide - none).""" - counts = mask.sum(dim=1) - hidden = torch.minimum((fraction * counts).round().clamp_min(1), counts - 1) - scores = torch.rand(mask.shape, generator=generator, device=mask.device) - rank = scores.masked_fill(~mask, 2.0).argsort(dim=1).argsort(dim=1) - return rank < hidden.unsqueeze(1) - - -# Data-dependent sizes do not compile (inductor has no aten.nonzero); the -# edge network itself compiles with dynamic sizes. -@torch.compiler.disable -def _graph_of(tokens: torch.Tensor, mask: torch.Tensor) -> tuple[Any, torch.Tensor]: - """Fully connected graphs of the objects in ``mask`` and their flat - positions in ``tokens``.""" - positions = mask.flatten().nonzero().squeeze(1) - nodes = tokens.flatten(0, 1).index_select(0, positions) - return SimpleNamespace(graph=build_graph(nodes, mask.sum(dim=1))), positions - - -def encode( - encoder: nn.Module, tokens: torch.Tensor, mask: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor | None]: - """Encodings ``[events, objects, hid_size]`` of the objects in ``mask`` - (zero elsewhere; see the module documentation for the edge network) and, - for the edge network, its decoded global state ``[events, hid_size]``.""" - if not isinstance(encoder, EdgeNetwork): - return encoder(tokens, mask), None - batch, positions = _graph_of(tokens, mask) - node_states, global_h = encoder.encode(batch) - node_h = torch.stack(node_states).mean(dim=0) - states = node_h.new_zeros((mask.numel(), node_h.shape[-1])) - states = states.index_copy(0, positions, node_h).view(*mask.shape, -1) - return states, encoder.global_decoder(global_h) - - -def event_embedding( - encoder: nn.Module, tokens: torch.Tensor, mask: torch.Tensor -) -> torch.Tensor: - """Embeddings ``[events, hid_size]`` of events by an object encoder: the - mean of the encodings of the objects in ``mask``, or the edge network's - decoded global state (what its classifier reads).""" - states, event = encode(encoder, tokens, mask) - if event is not None: - return event - return _mean_over_objects(states, mask) - - -def _mean_over_objects(states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - counts = mask.sum(dim=1, keepdim=True).clamp_min(1) - return (states * mask.unsqueeze(-1)).sum(dim=1) / counts - - -def layer_event_embeddings( - encoder: nn.Module, tokens: torch.Tensor, mask: torch.Tensor -) -> list[torch.Tensor]: - """Event embeddings ``[events, hid_size]`` after each layer of an object - encoder: the mean over objects after each transformer block and after its - final norm; the edge network's mean node states after its encoder and - each processing step, then its decoded global state. The last is - :func:`event_embedding`.""" - if not isinstance(encoder, EdgeNetwork): - layers = encoder.encode_layers(tokens, mask) - return [_mean_over_objects(states, mask) for states in layers] - batch, _ = _graph_of(tokens, mask) - node_states, global_h = encoder.encode(batch) - graph = batch.graph - pooled = [ - _mean_by_group(h, graph.node_batch, graph.node_counts) for h in node_states - ] - return [*pooled, encoder.global_decoder(global_h)] - - -#: Width of the kinematic targets of :func:`object_features`. -FEATURE_TARGETS = 7 - - -def object_features(tokens: torch.Tensor) -> torch.Tensor: - """Kinematic targets ``[events, objects, FEATURE_TARGETS]`` (see the - module documentation).""" - pt, energy, phi = tokens[..., PT], tokens[..., ENERGY], tokens[..., PHI] - btag, charge = tokens[..., BTAG], tokens[..., CHARGE] - return torch.stack( - ( - pt.clamp_min(0).log1p(), - tokens[..., ETA], - phi.cos(), - phi.sin(), - energy.clamp_min(0).log1p(), - btag, - charge, - ), - dim=-1, - ) - - -def _variance_penalty(encodings: torch.Tensor) -> torch.Tensor: - """VICReg variance term (Bardes et al. 2022): the mean shortfall of each - dimension's standard deviation over ``encodings [n, dims]`` below 1.""" - std = (encodings.var(dim=0) + 1e-4).sqrt() - return nn.functional.relu(1 - std).mean() - - -class JEPA(nn.Module): - """Context and target encoders and the predictor. - - ``forward`` returns ``[events, 3]``: each event's loss, whether it has - hidden objects, and the spread of the target encodings over the batch - (their standard deviation over objects, averaged over dimensions; it - falls towards zero if the representation collapses). With kinematic - targets, the spread is that of the context encodings. - """ - - def __init__( - self, - encoder: nn.Module, - *, - num_node_types: int, - n_heads: int, - ffn_size: int, - dropout: float = 0.0, - predictor_layers: int = 2, - mask_fraction: float = 0.4, - ema_momentum: float = 0.996, - variance_weight: float = 0.0, - target: str = "encoding", - ) -> None: - super().__init__() - if target not in ("encoding", "features"): - raise ValueError( - f"unknown jepa_target {target!r}; use encoding or features" - ) - hid_size = encoder.hid_size - self.mask_fraction, self.ema_momentum = mask_fraction, ema_momentum - self.variance_weight = variance_weight - self.context_encoder = encoder - self.target_encoder = ( - copy.deepcopy(encoder).requires_grad_(False) - if target == "encoding" - else None - ) - self.query_type = nn.Embedding(num_node_types, hid_size) - self.query_index = nn.Embedding(MAX_TYPE_INDEX, hid_size) - self.predictor = nn.ModuleList( - EncoderBlock(hid_size, n_heads, ffn_size, dropout) - for _ in range(predictor_layers) - ) - self.predictor_norm = nn.LayerNorm(hid_size) - self.predictor_out = nn.Linear( - hid_size, hid_size if target == "encoding" else FEATURE_TARGETS - ) - #: Optimizer steps taken, and the steps over which the momentum - #: reaches 1 (set before training; see :meth:`update_target`). - self.register_buffer("step", torch.zeros((), dtype=torch.long)) - self.total_steps = 1 - - def forward(self, batch: Any) -> torch.Tensor: - tokens, mask = batch.tokens, batch.token_mask - # Validation hides the same objects every epoch. - generator = ( - None if self.training else torch.Generator(tokens.device).manual_seed(0) - ) - hidden = sample_targets(mask, self.mask_fraction, generator) - with torch.no_grad(): - if self.target_encoder is None: - target = object_features(tokens).float() - else: - target = encode(self.target_encoder, tokens, mask)[0].float() - spread = target[mask].std(dim=0).mean() - target = nn.functional.layer_norm(target, target.shape[-1:]) - predicted, context, event = self.predict(tokens, mask, hidden) - if self.target_encoder is None: - spread = context[mask & ~hidden].detach().float().std(dim=0).mean() - error = (predicted - target).square().mean(dim=-1) * hidden - count = hidden.sum(dim=1) - loss = error.sum(dim=1) / count.clamp_min(1) - if self.variance_weight > 0: - encodings = [context[mask & ~hidden]] - if event is not None: - encodings.append(event) - penalty = sum(_variance_penalty(h.float()) for h in encodings) - loss = loss + self.variance_weight * penalty - return torch.stack( - (loss, (count > 0).to(loss.dtype), spread.expand_as(loss)), dim=1 - ) - - def predict( - self, tokens: torch.Tensor, mask: torch.Tensor, hidden: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - """Predicted float32 encodings ``[events, objects, hid_size]`` (only - those of ``hidden`` objects are meaningful) and the context encodings - of :func:`encode`.""" - context, event = encode(self.context_encoder, tokens, mask & ~hidden) - node_type = tokens[..., NODE_TYPE].long() - query = self.query_type(node_type) + self.query_index( - type_index(node_type, mask, self.query_type.num_embeddings) - ) - # Hidden objects enter the predictor only through their queries. - x = torch.where(hidden.unsqueeze(-1), query.to(context.dtype), context) - keys = mask - if event is not None: # the edge network's global state, first - x = torch.cat((event.unsqueeze(1).to(x.dtype), x), dim=1) - keys = torch.cat((mask.new_ones((mask.shape[0], 1)), mask), dim=1) - for block in self.predictor: - x = block(x, keys[:, None, None, :]) - x = x[:, -mask.shape[1] :] - return self.predictor_out(self.predictor_norm(x)).float(), context, event - - @torch.no_grad() - def update_target(self) -> None: - """Move the target encoder towards the context encoder, with a - momentum rising linearly from ``ema_momentum`` to 1 over - ``total_steps``. Called by the trainer after each optimizer step.""" - if self.target_encoder is None: - return - progress = (self.step.float() / self.total_steps).clamp(max=1) - momentum = self.ema_momentum + (1 - self.ema_momentum) * progress - for target, context in zip( - self.target_encoder.parameters(), self.context_encoder.parameters() - ): - target.lerp_(context, 1 - momentum) - self.step += 1 diff --git a/src/gnn4colliders/loader.py b/src/gnn4colliders/loader.py deleted file mode 100644 index ffab46d317ac424b95e6784943b78c13d7453359..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/loader.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Deterministic batching of packed events.""" - -from __future__ import annotations - -from collections.abc import Callable, Iterator -from typing import Any - -import torch - -from .data import PackedEvents - -Collate = Callable[[PackedEvents, torch.Tensor], Any] - - -class _Batches(torch.utils.data.Dataset): - """Batch ``i`` of one epoch's fixed event order, for DataLoader workers.""" - - def __init__(self, events, collate, order, batch_size) -> None: - self.events, self.collate, self.order = events, collate, order - self.batch_size = batch_size - - def __len__(self) -> int: - return -(-len(self.order) // self.batch_size) - - def __getitem__(self, index: int) -> Any: - start = index * self.batch_size - return self.collate(self.events, self.order[start : start + self.batch_size]) - - -class EventLoader: - """Batch the events ``indices`` (all by default) with ``collate``. - - Shuffling permutes the selection with a generator seeded by - ``seed + epoch``, so iteration never touches global random state. Under - DDP each of ``world_size`` ranks takes every ``world_size``-th event of - that order; shuffled (training) orders are padded by repeating events so - all ranks run the same number of steps, evaluation orders are not. - - With ``num_workers > 0`` batches are collated ahead of time in forked - worker processes and yielded in the same order. ``pin_memory`` returns - page-locked batches for asynchronous copies to the GPU. - """ - - def __init__( - self, - events: PackedEvents, - collate: Collate, - batch_size: int, - *, - indices: torch.Tensor | None = None, - shuffle: bool = False, - seed: int = 0, - num_workers: int = 0, - pin_memory: bool = False, - rank: int = 0, - world_size: int = 1, - ) -> None: - if batch_size < 1: - raise ValueError("batch_size must be positive") - self.events, self.collate, self.batch_size = events, collate, batch_size - self.indices = torch.arange(len(events)) if indices is None else indices - self.shuffle, self.seed, self.epoch = shuffle, seed, 0 - self.num_workers, self.pin_memory = num_workers, pin_memory - self.rank, self.world_size = rank, world_size - - def set_epoch(self, epoch: int) -> None: - self.epoch = epoch - - def order(self) -> torch.Tensor: - """The events of this rank and epoch, in batch order.""" - order = self.indices - if self.shuffle: - generator = torch.Generator().manual_seed(self.seed + self.epoch) - order = order[torch.randperm(len(order), generator=generator)] - if self.world_size > 1: - extra = -len(order) % self.world_size - if self.shuffle and extra: - order = torch.cat((order, order[:extra])) - order = order[self.rank :: self.world_size] - return order - - def __len__(self) -> int: - return -(-len(self.order()) // self.batch_size) - - def __iter__(self) -> Iterator[Any]: - order = self.order() - if self.num_workers == 0: - for start in range(0, len(order), self.batch_size): - batch = self.collate( - self.events, order[start : start + self.batch_size] - ) - yield batch.pin_memory() if self.pin_memory else batch - return - yield from torch.utils.data.DataLoader( - _Batches(self.events, self.collate, order, self.batch_size), - batch_size=None, - num_workers=self.num_workers, - pin_memory=self.pin_memory, - prefetch_factor=4, - # Forked workers share the events copy-on-write; Python 3.14's - # default (forkserver) would pickle them to every worker. Workers - # only collate on the CPU, so forking a CUDA process is safe. - multiprocessing_context="fork", - ) diff --git a/src/gnn4colliders/models.py b/src/gnn4colliders/models.py deleted file mode 100644 index d75d271a270dca76b62b5e084f87d0d0bcf73e8c..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/models.py +++ /dev/null @@ -1,720 +0,0 @@ -"""Models: the ROOT-GNN edge network (Ho et al. 2026), a transformer and Deep Sets. - -Every model takes a batch (:class:`~gnn4colliders.graphs.GraphBatch` or -:class:`~gnn4colliders.sequences.SequenceBatch`) and returns raw logits. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any - -import torch -from torch import nn - -from .data import NODE_FEATURE_NAMES -from .graphs import EDGE_FEATURE_NAMES, TensorGraph -from .sequences import ETA, PHI, PT - -if TYPE_CHECKING: - from .config import ModelConfig - -NODE_TYPE = NODE_FEATURE_NAMES.index("node_type") - - -def make_mlp( - in_size: int, hidden_size: int, out_size: int, n_layers: int, dropout: float = 0.0 -) -> nn.Sequential: - """Linear, ReLU, dropout per layer (the last included), then LayerNorm. - - Identity stands in for disabled dropout so module indices, and with them - the state-dict keys of the pretrained checkpoints, stay the same. - """ - sizes = [in_size] + [hidden_size] * (n_layers - 1) + [out_size] - layers: list[nn.Module] = [] - for size_in, size_out in zip(sizes, sizes[1:]): - layers += [ - nn.Linear(size_in, size_out), - nn.ReLU(), - nn.Dropout(dropout) if dropout > 0 else nn.Identity(), - ] - return nn.Sequential(*layers, nn.LayerNorm(out_size)) - - -def _sum_by_group( - values: torch.Tensor, groups: torch.Tensor, count: int -) -> torch.Tensor: - # Sums accumulate in at least float32, also under bfloat16 autocast. - dtype = torch.promote_types(values.dtype, torch.float32) - result = values.new_zeros((count, values.shape[1]), dtype=dtype) - return result.index_add_(0, groups, values.to(dtype)) - - -def _mean_by_group( - values: torch.Tensor, groups: torch.Tensor, sizes: torch.Tensor -) -> torch.Tensor: - # ``shape[0]`` stays symbolic under torch.compile (no recompile per size). - sums = _sum_by_group(values, groups, sizes.shape[0]) - return sums / sizes.to(sums.dtype)[:, None].clamp_min(1) - - -def _gather(values: torch.Tensor, index: torch.Tensor) -> torch.Tensor: - """``values[index]`` whose gradient is a scatter-add, which is much faster - on GPUs than the sort-based backward of advanced indexing.""" - return values.index_select(0, index) - - -class EdgeNetwork(nn.Module): - """Encode nodes, edges and a global vector, then ``n_proc_steps`` rounds - of edge, node and global updates; classify the decoded global vector. - - Without global features the global input is the node count. Parameter - names match the original DGL implementation, so the pretrained - checkpoints load unchanged (see :func:`load_weights`). - """ - - def __init__( - self, - node_size: int, - edge_size: int, - global_size: int, - hid_size: int, - out_size: int, - n_layers: int, - n_proc_steps: int, - dropout: float = 0.0, - ) -> None: - super().__init__() - self.has_global = global_size > 0 - self.hid_size, self.n_proc_steps = hid_size, n_proc_steps - - def mlp(in_size: int) -> nn.Sequential: - return make_mlp(in_size, hid_size, hid_size, n_layers, dropout) - - self.node_encoder = mlp(node_size) - self.edge_encoder = mlp(edge_size) - self.global_encoder = mlp(global_size or 1) - self.node_update = mlp(3 * hid_size) - self.edge_update = mlp(4 * hid_size) - self.global_update = mlp(3 * hid_size) - self.global_decoder = mlp(hid_size) - self.classifier = nn.Linear(hid_size, out_size) - - def encode(self, batch: Any) -> tuple[list[torch.Tensor], torch.Tensor]: - """Node states after the encoder and after each processing step, and - the global state after the last step.""" - graph: TensorGraph = batch.graph - if self.has_global: - global_input = batch.global_features - else: - global_input = graph.node_counts[:, None].to(torch.float32) - node_h = self.node_encoder(graph.node_features) - edge_h = self.edge_encoder(graph.edge_features) - global_h = self.global_encoder(global_input) - node_states = [node_h] - for _ in range(self.n_proc_steps): - edge_h = self.edge_update( - torch.cat( - ( - edge_h, - _gather(node_h, graph.edge_src), - _gather(node_h, graph.edge_dst), - _gather(global_h, graph.edge_batch), - ), - dim=1, - ) - ) - node_h = self.node_update( - torch.cat( - ( - node_h, - _sum_by_group(edge_h, graph.edge_dst, node_h.shape[0]), - _gather(global_h, graph.node_batch), - ), - dim=1, - ) - ) - node_states.append(node_h) - global_h = self.global_update( - torch.cat( - ( - global_h, - _mean_by_group(node_h, graph.node_batch, graph.node_counts), - _mean_by_group(edge_h, graph.edge_batch, graph.edge_counts), - ), - dim=1, - ) - ) - return node_states, global_h - - def forward_features(self, batch: Any) -> torch.Tensor: - """The decoded global representation, before the classifier.""" - return self.global_decoder(self.encode(batch)[1]) - - def forward(self, batch: Any) -> torch.Tensor: - return self.classifier(self.forward_features(batch)) - - -class FineTunedEdgeNetwork(nn.Module): - """A (pretrained) :class:`EdgeNetwork` backbone with a new classifier.""" - - def __init__( - self, backbone: EdgeNetwork, out_size: int, *, freeze_backbone: bool = False - ) -> None: - super().__init__() - backbone.classifier = nn.Identity() - self.backbone = backbone - self.classifier = nn.Linear(backbone.hid_size, out_size) - for parameter in backbone.parameters(): - parameter.requires_grad = not freeze_backbone - - def forward_features(self, batch: Any) -> torch.Tensor: - return self.backbone.forward_features(batch) - - def forward(self, batch: Any) -> torch.Tensor: - return self.classifier(self.forward_features(batch)) - - -class ObjectEmbedding(nn.Module): - """Object tokens ``[..., 7]`` to ``[..., hid_size]``: a linear map of the - kinematic columns (``phi`` as its cosine and sine) plus a learned - embedding of ``node_type``. Padding tokens are embedded too; models mask - them.""" - - def __init__(self, input_size: int, num_node_types: int, hid_size: int) -> None: - super().__init__() - # Minus the node type column, plus the second phi column. - self.linear = nn.Linear(input_size, hid_size) - self.type_embedding = nn.Embedding(num_node_types, hid_size) - - def forward(self, tokens: torch.Tensor) -> torch.Tensor: - phi = tokens[..., PHI] - kinematics = torch.cat( - ( - tokens[..., :PHI], - torch.cos(phi).unsqueeze(-1), - torch.sin(phi).unsqueeze(-1), - tokens[..., PHI + 1 : NODE_TYPE], - tokens[..., NODE_TYPE + 1 :], - ), - dim=-1, - ) - return self.linear(kinematics) + self.type_embedding( - tokens[..., NODE_TYPE].long() - ) - - -#: Pair feature columns of :func:`pair_features`. -PAIR_FEATURE_NAMES = ("deta", "dphi", "dR", "log1p_m2") - - -def pair_features(tokens: torch.Tensor) -> torch.Tensor: - """``[events, objects, objects, 4]`` features of every object pair: the - edge features of the edge network and ``log(1 + m^2)`` of the pair's - invariant mass for massless objects, in the units of the tokens.""" - pt, eta, phi = tokens[..., PT], tokens[..., ETA], tokens[..., PHI] - deta = eta.unsqueeze(2) - eta.unsqueeze(1) - dphi = phi.unsqueeze(2) - phi.unsqueeze(1) - dphi = torch.atan2(torch.sin(dphi), torch.cos(dphi)) - # m^2 = 2 pt1 pt2 (cosh deta - cos dphi), written without the cancellation - # of nearly collinear pairs; sinh through exp (no ONNX sinh at opset 17). - sinh_half = (torch.exp(deta / 2) - torch.exp(-deta / 2)) / 2 - m2 = ( - 4 - * pt.unsqueeze(2) - * pt.unsqueeze(1) - * (sinh_half.square() + torch.sin(dphi / 2).square()) - ) - return torch.stack( - (deta, dphi, torch.sqrt(deta.square() + dphi.square()), torch.log1p(m2)), - dim=-1, - ) - - -class EncoderBlock(nn.Module): - """Pre-norm transformer block: multi-head attention with an additive or - boolean ``attn_mask``, then a ReLU feed-forward network.""" - - def __init__(self, hid_size: int, n_heads: int, ffn_size: int, dropout: float): - super().__init__() - if hid_size % n_heads: - raise ValueError("model.hid_size must be a multiple of model.n_heads") - self.n_heads, self.dropout = n_heads, dropout - self.attention_norm = nn.LayerNorm(hid_size) - self.qkv = nn.Linear(hid_size, 3 * hid_size) - self.out = nn.Linear(hid_size, hid_size) - self.ffn_norm = nn.LayerNorm(hid_size) - self.ffn = nn.Sequential( - nn.Linear(hid_size, ffn_size), - nn.ReLU(), - nn.Dropout(dropout), - nn.Linear(ffn_size, hid_size), - ) - self.residual_dropout = nn.Dropout(dropout) - - def forward(self, x: torch.Tensor, attn_mask: torch.Tensor) -> torch.Tensor: - return self.attend(x, *self.project(x).unbind(0), attn_mask) - - def project(self, x: torch.Tensor) -> torch.Tensor: - """Queries, keys and values ``[3, events, heads, objects, head size]``.""" - events, objects, size = x.shape - return ( - self.qkv(self.attention_norm(x)) - .view(events, objects, 3, self.n_heads, size // self.n_heads) - .permute(2, 0, 3, 1, 4) - ) - - def attend( - self, - x: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - attn_mask: torch.Tensor | None = None, - ) -> torch.Tensor: - """``x`` updated by attention of its queries ``q`` to ``k`` and - ``v``, then by the feed-forward network.""" - if attn_mask is not None and attn_mask.is_floating_point(): - attn_mask = attn_mask.to(q.dtype) - attended = nn.functional.scaled_dot_product_attention( - q, k, v, attn_mask, dropout_p=self.dropout if self.training else 0.0 - ) - attended = attended.transpose(1, 2).reshape(x.shape) - x = x + self.residual_dropout(self.out(attended)) - return x + self.residual_dropout(self.ffn(self.ffn_norm(x))) - - -class ObjectTransformer(nn.Module): - """Per-object encodings ``[events, objects, hid_size]`` of a stack of - pre-norm attention blocks over the objects whose ``mask`` is true. - - There is no position embedding, so the encodings do not depend on the - order of the objects. With ``pair_bias`` (as in the Particle Transformer, - Qu et al. 2022) an MLP of :func:`pair_features` adds a per-head bias to - the attention logits of every block. - """ - - def __init__( - self, - input_size: int, - num_node_types: int, - hid_size: int = 64, - n_layers: int = 2, - n_heads: int = 4, - ffn_size: int | None = None, - dropout: float = 0.0, - pair_bias: bool = False, - ) -> None: - super().__init__() - self.hid_size = hid_size - self.embedding = ObjectEmbedding(input_size, num_node_types, hid_size) - self.pair_encoder = ( - nn.Sequential( - nn.Linear(len(PAIR_FEATURE_NAMES), hid_size), - nn.ReLU(), - nn.Linear(hid_size, n_heads), - ) - if pair_bias - else None - ) - self.blocks = nn.ModuleList( - EncoderBlock(hid_size, n_heads, ffn_size or 4 * hid_size, dropout) - for _ in range(n_layers) - ) - self.norm = nn.LayerNorm(hid_size) - - def encode_layers( - self, tokens: torch.Tensor, mask: torch.Tensor - ) -> list[torch.Tensor]: - """The encodings ``[events, objects, hid_size]`` after each block, - then after the final norm (the output of :meth:`forward`).""" - encoded = self.embedding(tokens) - # [events, heads, queries, keys]; objects outside mask are not attended to. - attn_mask = mask[:, None, None, :] - if self.pair_encoder is not None: - bias = self.pair_encoder(pair_features(tokens)).permute(0, 3, 1, 2) - attn_mask = bias.masked_fill(~attn_mask, float("-inf")) - layers = [] - for block in self.blocks: - encoded = block(encoded, attn_mask) - layers.append(encoded) - return [*layers, self.norm(encoded)] - - def forward(self, tokens: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: - return self.encode_layers(tokens, mask)[-1] - - -class TransformerClassifier(nn.Module): - """An :class:`ObjectTransformer`, a mean over real objects and a linear - classifier. ``freeze_encoder`` trains only the classifier (and the - global-feature encoder), e.g. on a JEPA-pretrained encoder.""" - - def __init__( - self, - input_size: int, - out_size: int, - num_node_types: int, - hid_size: int = 64, - n_layers: int = 2, - n_heads: int = 4, - ffn_size: int | None = None, - dropout: float = 0.0, - global_size: int = 0, - pair_bias: bool = False, - freeze_encoder: bool = False, - ) -> None: - super().__init__() - self.encoder = ObjectTransformer( - input_size, - num_node_types, - hid_size, - n_layers, - n_heads, - ffn_size, - dropout, - pair_bias, - ) - self.encoder.requires_grad_(not freeze_encoder) - self.global_encoder = nn.Linear(global_size, hid_size) if global_size else None - self.classifier = nn.Linear(hid_size, out_size) - - def forward(self, batch: Any) -> torch.Tensor: - mask = batch.token_mask - encoded = self.encoder(batch.tokens, mask) - counts = mask.sum(dim=1, keepdim=True).clamp_min(1) - pooled = (encoded * mask.unsqueeze(-1)).sum(dim=1) / counts - if self.global_encoder is not None: - pooled = pooled + self.global_encoder(batch.global_features) - return self.classifier(pooled) - - -class DeepSets(nn.Module): - """Deep Sets (Zaheer et al. 2017; the particle flow network of Komiske et - al. 2019): an MLP per object, a sum over real objects, then an MLP of the - sum. The simplest model that does not depend on object order and has no - interactions between objects.""" - - def __init__( - self, - input_size: int, - out_size: int, - num_node_types: int, - hid_size: int = 64, - n_layers: int = 3, - dropout: float = 0.0, - global_size: int = 0, - ) -> None: - super().__init__() - self.embedding = ObjectEmbedding(input_size, num_node_types, hid_size) - self.object_mlp = make_mlp(hid_size, hid_size, hid_size, n_layers, dropout) - self.event_mlp = make_mlp( - hid_size + global_size, hid_size, hid_size, n_layers, dropout - ) - self.classifier = nn.Linear(hid_size, out_size) - - def forward(self, batch: Any) -> torch.Tensor: - mask = batch.token_mask - objects = self.object_mlp(self.embedding(batch.tokens)) - pooled = (objects * mask.unsqueeze(-1)).sum(dim=1) - if batch.global_features is not None: - pooled = torch.cat((pooled, batch.global_features), dim=1) - return self.classifier(self.event_mlp(pooled)) - - -#: Models that take :class:`~gnn4colliders.sequences.SequenceBatch` inputs. -SEQUENCE_MODELS = ( - "transformer", - "deep_sets", - "jepa", - "contrastive", - "pfn", - "tabpfn", - "tabicl", - "tabdpt", -) -#: Models that classify events given labelled context events (see cli.py). -IN_CONTEXT_MODELS = ("pfn", "tabpfn", "tabicl", "tabdpt") - - -def state_dict_of(checkpoint: Mapping[str, Any]) -> dict[str, Any]: - """Model weights of a checkpoint, with historical key spellings mapped. - - Accepts a bare state dict or a checkpoint holding ``model_state_dict`` - (this package's and the pretrained ``model_epoch_N.pt`` files). DDP and - ``torch.compile`` prefixes are removed and the original ``classify`` - layer is renamed ``classifier``. - """ - state = checkpoint.get("model_state_dict", checkpoint) - result = {} - for key, value in state.items(): - key = key.removeprefix("module.").removeprefix("_orig_mod.") - if key.startswith("classify."): - key = "classifier." + key.removeprefix("classify.") - result[key] = value - return result - - -def load_weights(model: nn.Module, checkpoint: Mapping[str, Any]) -> None: - model.load_state_dict(state_dict_of(checkpoint), strict=True) - - -def _strip_prefix( - state: dict[str, Any], prefixes: tuple[str, ...] -) -> dict[str, Any] | None: - """The entries of ``state`` under the first of ``prefixes`` it has, the - prefix removed, or None.""" - for prefix in prefixes: - part = { - key.removeprefix(prefix): value - for key, value in state.items() - if key.startswith(prefix) - } - if part: - return part - return None - - -def _load_state(path: str) -> dict[str, Any]: - return state_dict_of(torch.load(path, map_location="cpu", weights_only=False)) - - -def pretrained_state(path: str) -> dict[str, Any]: - """Weights of a pretrained backbone: those of a supervised checkpoint, or - the context encoder of a JEPA or contrastive checkpoint.""" - state = _load_state(path) - return _strip_prefix(state, ("context_encoder.",)) or state - - -#: Prefixes of an encoder's weights in the checkpoints this package writes: -#: pretraining (JEPA, contrastive), a supervised transformer, a fine-tuned -#: edge network and the encoder of a meta-trained tabular model. -ENCODER_PREFIXES = ("context_encoder.", "encoder.", "backbone.", "inputs.encoder.") - - -def encoder_state(path: str) -> dict[str, torch.Tensor]: - """The encoder weights of a checkpoint (see ``ENCODER_PREFIXES``; else a - supervised edge network), without a classifier.""" - state = _load_state(path) - state = _strip_prefix(state, ENCODER_PREFIXES) or state - return { - key: value for key, value in state.items() if not key.startswith("classifier.") - } - - -def _load_encoder(encoder: ObjectTransformer, path: str) -> None: - state = pretrained_state(path) - if "embedding.linear.weight" not in state: - raise ValueError(f"{path} holds no JEPA transformer encoder") - encoder.load_state_dict(state, strict=True) - - -def _transformer_sizes( - config: ModelConfig, num_features: int, num_node_types: int -) -> dict[str, Any]: - return { - "input_size": num_features, - "num_node_types": num_node_types, - "hid_size": config.hid_size, - "n_layers": config.n_layers, - "n_heads": config.n_heads, - "ffn_size": config.ffn_size, - "dropout": config.dropout, - "pair_bias": config.pair_bias, - } - - -def _edge_network_sizes(config: ModelConfig, num_features: int) -> dict[str, Any]: - return { - "node_size": num_features, - "edge_size": len(EDGE_FEATURE_NAMES), - "hid_size": config.hid_size, - "n_layers": config.n_layers, - "n_proc_steps": config.n_proc_steps, - "dropout": config.dropout, - } - - -def build_encoder( - config: ModelConfig, backbone: str, *, num_features: int, num_node_types: int -) -> nn.Module: - """The object encoder ``backbone`` (``transformer`` or ``edge_network``) - with the sizes of ``config``, as pretrained (JEPA, contrastive) and - probed: the edge network without global features or classifier.""" - if backbone == "transformer": - return ObjectTransformer( - **_transformer_sizes(config, num_features, num_node_types) - ) - if backbone == "edge_network": - encoder = EdgeNetwork( - global_size=0, out_size=1, **_edge_network_sizes(config, num_features) - ) - encoder.classifier = nn.Identity() - return encoder - raise ValueError(f"unknown backbone {backbone!r}; use transformer or edge_network") - - -#: Seed of the initialization of encoders without pretrained weights, so -#: that a frozen random encoder is rebuilt identically at evaluation and on -#: every rank. -RANDOM_ENCODER_SEED = 0 - - -def pretrained_encoder( - config: ModelConfig, - *, - num_features: int, - num_node_types: int, - backbone: str | None = None, - freeze: bool | None = None, -) -> nn.Module: - """The object encoder ``backbone`` (default ``config.backbone``) with - the weights of ``config.pretrained``, or randomly initialized with - ``RANDOM_ENCODER_SEED`` without; frozen if ``freeze`` (default - ``config.freeze_backbone``).""" - with torch.random.fork_rng(devices=[]): - torch.manual_seed(RANDOM_ENCODER_SEED) - encoder = build_encoder( - config, - backbone or config.backbone, - num_features=num_features, - num_node_types=num_node_types, - ) - if config.pretrained: - encoder.load_state_dict(encoder_state(config.pretrained), strict=True) - freeze = config.freeze_backbone if freeze is None else freeze - return encoder.requires_grad_(not freeze) - - -def build_model( - config: ModelConfig, - *, - num_features: int, - num_node_types: int, - num_global_features: int = 0, - load_pretrained: bool = True, -) -> nn.Module: - """The model of ``config`` for inputs of the given widths. - - With ``config.pretrained`` the edge network is a - :class:`FineTunedEdgeNetwork` whose backbone is loaded from a supervised - or JEPA checkpoint, and the transformer's encoder is loaded from a JEPA - checkpoint, as is the PFN's (all skipped with ``load_pretrained=False``, - e.g. when a trained checkpoint is loaded next). - """ - transformer = _transformer_sizes(config, num_features, num_node_types) - edge_network = _edge_network_sizes(config, num_features) - if config.name == "jepa": - from .jepa import JEPA - - # Global features are not used in pretraining. - encoder = build_encoder( - config, - config.backbone, - num_features=num_features, - num_node_types=num_node_types, - ) - return JEPA( - encoder, - num_node_types=num_node_types, - n_heads=config.n_heads, - ffn_size=config.ffn_size or 4 * config.hid_size, - dropout=config.dropout, - predictor_layers=config.predictor_layers, - mask_fraction=config.mask_fraction, - ema_momentum=config.ema_momentum, - variance_weight=config.variance_weight, - target=config.jepa_target, - ) - if config.name == "contrastive": - from .contrastive import Contrastive - - encoder = build_encoder( - config, - config.backbone, - num_features=num_features, - num_node_types=num_node_types, - ) - return Contrastive( - encoder, - temperature=config.temperature, - scale_jitter=config.scale_jitter, - object_dropout=config.object_dropout, - ) - if config.name == "transformer": - model = TransformerClassifier( - out_size=config.out_size, - global_size=num_global_features, - freeze_encoder=config.freeze_backbone, - **transformer, - ) - if config.pretrained and load_pretrained: - _load_encoder(model.encoder, config.pretrained) - return model - if config.name == "pfn": - from .pfn import EventPFN - - encoder = ObjectTransformer(**transformer) - if config.pretrained and load_pretrained: - _load_encoder(encoder, config.pretrained) - return EventPFN( - encoder, - max_classes=config.max_classes, - n_layers=config.pfn_layers, - n_heads=config.n_heads, - ffn_size=config.ffn_size or 4 * config.hid_size, - dropout=config.dropout, - ) - if config.name in ("tabpfn", "tabicl"): - from . import tabular - - # The released weights load whatever load_pretrained says. - if config.name == "tabpfn": - network = tabular.load_tabpfn_network( - config.tabpfn_version, config.tabpfn_checkpoint - ) - meta = tabular.TabularPFN - else: - network = tabular.load_tabicl_network(config.tabicl_checkpoint) - meta = tabular.TabularICL - # So do those of a frozen encoder of the inputs. - inputs = tabular.EventInputs.of( - config, num_features=num_features, num_node_types=num_node_types - ) - return meta( - network, - inputs=inputs, - max_classes=config.max_classes, - lora_rank=config.lora_rank, - ) - if config.name == "tabdpt": - raise ValueError("tabdpt is not trained; run evaluate or predict") - if config.name == "deep_sets": - return DeepSets( - input_size=num_features, - out_size=config.out_size, - num_node_types=num_node_types, - hid_size=config.hid_size, - n_layers=config.n_layers, - dropout=config.dropout, - global_size=num_global_features, - ) - if config.name != "edge_network": - raise ValueError( - f"unknown model {config.name!r}; use edge_network, transformer, " - "deep_sets, jepa, contrastive, pfn, tabpfn, tabicl or tabdpt" - ) - sizes = {"global_size": num_global_features, **edge_network} - if not config.pretrained: - return EdgeNetwork(out_size=config.out_size, **sizes) - backbone = EdgeNetwork(out_size=1, **sizes) - if load_pretrained: - state = pretrained_state(config.pretrained) - if "classifier.weight" in state: # supervised; the classifier is replaced - out_size = state["classifier.weight"].shape[0] - backbone = EdgeNetwork(out_size=out_size, **sizes) - else: # a JEPA context encoder - backbone.classifier = nn.Identity() - backbone.load_state_dict(state, strict=True) - return FineTunedEdgeNetwork( - backbone, config.out_size, freeze_backbone=config.freeze_backbone - ) diff --git a/src/gnn4colliders/pfn.py b/src/gnn4colliders/pfn.py deleted file mode 100644 index 3567cf37c425fc03243f7a653cf6209b91aea584..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/pfn.py +++ /dev/null @@ -1,341 +0,0 @@ -"""A prior-fitted network (PFN; Müller et al. 2022, as in TabPFN of -Hollmann et al. 2023) that classifies collider events in context. - -A task is a set of context events with labels and query events without. -The PFN embeds every event with the object transformer (mean over objects), -adds a label embedding to the context events, and runs attention blocks over -the task's events in which every event attends to the context events only. -Queries therefore never see each other, and the context's keys and values -can be computed once for any number of queries (:func:`in_context_forward`). - -Meta-training (``task.name: pfn``) draws tasks from a labelled cache, the -prior: each task takes ``task.task_classes`` random classes of -``task.prior_classes``, numbered in random order, with random class -proportions and a log-uniform number of context events. Training tasks are -augmented by a random azimuthal rotation and longitudinal flip of each event -and a scale of the task's momenta and energies (``task.energy_jitter``). -A new dataset then needs no training: its labelled training events are the -context (``data.train_events``). -""" - -from __future__ import annotations - -import math -from collections.abc import Callable, Iterator -from dataclasses import dataclass -from typing import Any - -import torch -from torch import nn - -from .config import TaskConfig -from .data import PackedEvents, map_tensors -from .models import EncoderBlock, ObjectTransformer -from .sequences import SequenceBatch, augment_tokens, collate_sequence_batch - -#: Logit of the classes a task does not have. -ABSENT_LOGIT = -1e4 - - -@dataclass(frozen=True) -class TaskBatch: - """Tasks of equal size. ``events`` holds each task's context events, - then its query events; ``labels``, ``weight`` and ``event_index`` are the - queries' (labels numbered within each task).""" - - events: SequenceBatch - context_labels: torch.Tensor # [tasks, context events] - num_classes: torch.Tensor # [tasks] - labels: torch.Tensor - weight: torch.Tensor - event_index: torch.Tensor - - def to(self, device: torch.device | str, *, non_blocking: bool = False): - return map_tensors(self, lambda t: t.to(device, non_blocking=non_blocking)) - - def pin_memory(self) -> TaskBatch: - return map_tensors(self, torch.Tensor.pin_memory) - - -class EventPFN(nn.Module): - """Query logits ``[tasks * queries, max_classes]`` of a :class:`TaskBatch`; - the logits of classes a task does not have are ``ABSENT_LOGIT``.""" - - def __init__( - self, - encoder: ObjectTransformer, - *, - max_classes: int, - n_layers: int, - n_heads: int, - ffn_size: int, - dropout: float = 0.0, - ) -> None: - super().__init__() - hid_size = encoder.hid_size - self.max_classes = max_classes - self.encoder = encoder - self.label_embedding = nn.Embedding(max_classes, hid_size) - self.blocks = nn.ModuleList( - EncoderBlock(hid_size, n_heads, ffn_size, dropout) for _ in range(n_layers) - ) - self.norm = nn.LayerNorm(hid_size) - self.classifier = nn.Linear(hid_size, max_classes) - - def embed(self, batch: SequenceBatch) -> torch.Tensor: - """``[events, hid_size]``: the mean encoding of each event's objects.""" - mask = batch.token_mask - encoded = self.encoder(batch.tokens, mask) - counts = mask.sum(dim=1, keepdim=True).clamp_min(1) - return (encoded * mask.unsqueeze(-1)).sum(dim=1) / counts - - def context( - self, embedded: torch.Tensor, labels: torch.Tensor - ) -> list[tuple[torch.Tensor, torch.Tensor]]: - """Each block's keys and values of the context events - ``embedded [tasks, events, hid_size]`` with ``labels``.""" - x = embedded + self.label_embedding(labels) - memory = [] - for block in self.blocks: - q, k, v = block.project(x) - memory.append((k, v)) - x = block.attend(x, q, k, v) - return memory - - def queries( - self, - embedded: torch.Tensor, - memory: list[tuple[torch.Tensor, torch.Tensor]], - num_classes: torch.Tensor, - ) -> torch.Tensor: - """Logits ``[tasks, events, max_classes]`` of the query events - ``embedded`` given the context ``memory`` of :meth:`context`.""" - x = embedded - for block, (k, v) in zip(self.blocks, memory): - x = block.attend(x, block.project(x)[0], k, v) - logits = self.classifier(self.norm(x)) - classes = torch.arange(self.max_classes, device=logits.device) - absent = classes >= num_classes.reshape(-1, 1, 1) - return logits.masked_fill(absent, ABSENT_LOGIT) - - def forward(self, batch: TaskBatch) -> torch.Tensor: - tasks, size = batch.context_labels.shape - embedded = self.embed(batch.events).view(tasks, -1, self.encoder.hid_size) - memory = self.context(embedded[:, :size], batch.context_labels) - logits = self.queries(embedded[:, size:], memory, batch.num_classes) - return logits.flatten(0, 1) - - -@torch.inference_mode() -def in_context_forward( - model: EventPFN, - events: PackedEvents, - context: torch.Tensor, - *, - num_classes: int, - batch_size: int, - device: torch.device, -) -> Callable[[SequenceBatch], torch.Tensor]: - """A function from a batch of events to their logits ``[events, - num_classes]``, given the labelled ``context`` events (one task).""" - if num_classes > model.max_classes: - raise ValueError(f"the PFN tells apart at most {model.max_classes} classes") - model.eval() - embedded = torch.cat( - [ - model.embed(collate_sequence_batch(events, part).to(device)) - for part in context.split(batch_size) - ] - ) - labels = events.labels[context].to(device, torch.long) - memory = model.context(embedded.unsqueeze(0), labels.unsqueeze(0)) - classes = torch.tensor([num_classes], device=device) - - def forward(batch: SequenceBatch) -> torch.Tensor: - logits = model.queries(model.embed(batch).unsqueeze(0), memory, classes) - return logits[0, :, :num_classes] - - return forward - - -# --- The prior ---------------------------------------------------------------- - - -@dataclass(frozen=True) -class _Step: - """The sampled events and augmentations of one meta-training step.""" - - indices: torch.Tensor # [tasks * (context + queries)] - labels: torch.Tensor # [tasks, context + queries], within each task - num_classes: torch.Tensor # [tasks] - context: int - scale: torch.Tensor # [events] - rotation: torch.Tensor # [events] - flip: torch.Tensor # [events], 1 or -1 - - -def _augment(batch: SequenceBatch, step: _Step) -> SequenceBatch: - """Scale momenta and energies, rotate in phi and flip eta, per event.""" - mask = batch.token_mask - tokens = augment_tokens( - batch.tokens, mask, scale=step.scale, rotation=step.rotation, flip=step.flip - ) - return SequenceBatch( - tokens=tokens, - token_mask=mask, - labels=batch.labels, - weight=batch.weight, - global_features=None, - event_index=batch.event_index, - ) - - -class _Steps(torch.utils.data.Dataset): - """Collation of sampled steps, for DataLoader workers.""" - - def __init__(self, events: PackedEvents, steps: list[_Step]) -> None: - self.events, self.steps = events, steps - - def __len__(self) -> int: - return len(self.steps) - - def __getitem__(self, index: int) -> TaskBatch: - step = self.steps[index] - batch = _augment(collate_sequence_batch(self.events, step.indices), step) - tasks = step.labels.shape[0] - queries = step.indices.view(tasks, -1)[:, step.context :].flatten() - return TaskBatch( - events=batch, - context_labels=step.labels[:, : step.context], - num_classes=step.num_classes, - labels=step.labels[:, step.context :].flatten(), - weight=torch.ones(len(queries)), - event_index=queries, - ) - - -class TaskLoader: - """Meta-training batches of tasks drawn from the events ``indices``. - - A batch holds about ``batch_size`` events: tasks with the same numbers - of context and query events. Training tasks are drawn with a generator - seeded by the seed, epoch and rank and are augmented; validation tasks - (``train=False``) are the same every epoch and are not augmented. - Interface of :class:`~gnn4colliders.loader.EventLoader`. - """ - - def __init__( - self, - events: PackedEvents, - indices: torch.Tensor, - config: TaskConfig, - batch_size: int, - *, - train: bool, - seed: int = 0, - num_workers: int = 0, - pin_memory: bool = False, - rank: int = 0, - ) -> None: - self.events, self.indices, self.config = events, indices, config - self.batch_size, self.train, self.seed, self.rank = ( - batch_size, - train, - seed, - rank, - ) - self.num_workers, self.pin_memory, self.epoch = num_workers, pin_memory, 0 - labels = events.labels[indices] - allowed = config.prior_classes - self.pools = [ - indices[labels == label] - for label in labels.unique().tolist() - if allowed is None or label in allowed - ] - if len(self.pools) < 2: - raise ValueError("the PFN prior needs events of at least two classes") - - def set_epoch(self, epoch: int) -> None: - self.epoch = epoch - - def __len__(self) -> int: - config = self.config - return config.steps_per_epoch if self.train else config.validation_steps - - def _generator(self) -> torch.Generator: - epoch = self.epoch if self.train else -1 - return torch.Generator().manual_seed(hash((self.seed, epoch, self.rank))) - - def sample(self, generator: torch.Generator) -> _Step: - """One step's tasks (see the module documentation).""" - config = self.config - - def uniform(low: float, high: float, size: int = 1) -> torch.Tensor: - return low + (high - low) * torch.rand(size, generator=generator) - - low, high = config.context_events - context = int(math.exp(float(uniform(math.log(low), math.log(high + 1))))) - context = min(max(context, low), high) - size = context + config.query_events - tasks = max(1, self.batch_size // size) - largest = min(config.task_classes[1], len(self.pools), context) - smallest = min(config.task_classes[0], largest) - indices, labels, num_classes = [], [], [] - for _ in range(tasks): - k = int(torch.randint(smallest, largest + 1, (), generator=generator)) - chosen = torch.randperm(len(self.pools), generator=generator)[:k] - # Dirichlet(1) class proportions; every class has a context event. - share = -torch.log(torch.rand(k, generator=generator)) - share = share / share.sum() - rest = torch.multinomial(share, size - k, True, generator=generator) - task_labels = torch.cat((torch.arange(k), rest)) - task_indices = torch.empty(size, dtype=torch.long) - for label in range(k): - pool = self.pools[int(chosen[label])] - where = task_labels == label - draw = torch.randint( - len(pool), (int(where.sum()),), generator=generator - ) - task_indices[where] = pool[draw] - indices.append(task_indices) - labels.append(task_labels) - num_classes.append(k) - events = tasks * size - if self.train: - jitter = config.energy_jitter - scale = uniform(1 - jitter, 1 + jitter, tasks).repeat_interleave(size) - rotation = uniform(-math.pi, math.pi, events) - flip = torch.randint(2, (events,), generator=generator) * 2.0 - 1 - else: - scale, rotation, flip = ( - torch.ones(events), - torch.zeros(events), - torch.ones(events), - ) - return _Step( - indices=torch.cat(indices), - labels=torch.stack(labels), - num_classes=torch.tensor(num_classes), - context=context, - scale=scale, - rotation=rotation, - flip=flip, - ) - - def __iter__(self) -> Iterator[Any]: - generator = self._generator() - steps = _Steps(self.events, [self.sample(generator) for _ in range(len(self))]) - if self.num_workers == 0: - for index in range(len(steps)): - batch = steps[index] - yield batch.pin_memory() if self.pin_memory else batch - return - yield from torch.utils.data.DataLoader( - steps, - batch_size=None, - num_workers=self.num_workers, - pin_memory=self.pin_memory, - prefetch_factor=4, - # As in EventLoader: forked workers share the events. - multiprocessing_context="fork", - ) diff --git a/src/gnn4colliders/probe.py b/src/gnn4colliders/probe.py deleted file mode 100644 index 5bbf6252b3c6f8fcfc6c17dbf3fa890a9b86595d..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/probe.py +++ /dev/null @@ -1,206 +0,0 @@ -"""k-NN and linear probes of a frozen object encoder (as for jBOT, Tsoi and -Rankin 2026), to compare pretraining objectives without fine-tuning. - -The encoder of ``model.pretrained`` embeds every event: the mean of its -object encodings, or the edge network's decoded global state (what its -classifier reads). The checkpoint can be a JEPA, contrastive or masked-object -one (its context encoder) or a supervised transformer or edge network; with -none, the encoder is randomly initialized (with a fixed seed), the -reference every pretraining should beat. The config can be that of the -pretraining or of a classifier with the encoder's sizes. Two classifiers of -the frozen embeddings, of all the dataset's classes, are trained on -``data.train_events`` labelled training events and evaluated on ``split`` -(without those events): - -* k-NN: the class fractions among the ``model.probe_neighbours`` training - embeddings of largest cosine similarity; -* linear: multinomial logistic regression of the standardized embeddings - (full-batch L-BFGS, L2 penalty ``model.probe_l2``). - -:func:`linear_cka` compares two encoders' representations of the same events -(centered kernel alignment, Kornblith et al. 2019, as in Ho et al. 2026), -layer by layer with :func:`layer_embeddings` (``scripts/compare_encoders.py``). -""" - -from __future__ import annotations - -import math -from typing import Any - -import torch -import torch.nn.functional as F -from torch import nn - -from .config import Config -from .data import PackedEvents -from .jepa import event_embedding, layer_event_embeddings -from .models import pretrained_encoder -from .sequences import collate_sequence_batch - - -def backbone_of(config: Config) -> str: - """The encoder architecture of the model of ``config``.""" - name = config.model.name - if name in ("jepa", "contrastive"): - return config.model.backbone - if name in ("transformer", "edge_network"): - return name - raise ValueError(f"cannot probe model.name {name!r}") - - -def load_encoder(config: Config, events: PackedEvents) -> nn.Module: - """The (pretrained) frozen encoder of ``config``.""" - encoder = pretrained_encoder( - config.model, - num_features=events.num_features, - num_node_types=len(config.data.object_types), - backbone=backbone_of(config), - freeze=True, - ) - return encoder.eval() - - -@torch.inference_mode() -def embed( - encoder: nn.Module, - events: PackedEvents, - indices: torch.Tensor, - *, - batch_size: int, - device: torch.device, -) -> torch.Tensor: - """Event embeddings ``[events, hid_size]`` of ``indices`` - (:func:`~gnn4colliders.jepa.event_embedding`).""" - parts = [] - for part in indices.split(batch_size): - batch = collate_sequence_batch(events, part).to(device) - parts.append(event_embedding(encoder, batch.tokens, batch.token_mask).float()) - return torch.cat(parts) - - -def knn_logits( - train: torch.Tensor, - labels: torch.Tensor, - queries: torch.Tensor, - *, - num_classes: int, - neighbours: int, - chunk: int = 4096, -) -> torch.Tensor: - """Log class fractions ``[queries, num_classes]`` among each query's - ``neighbours`` most cosine-similar ``train`` embeddings.""" - train, queries = F.normalize(train, dim=1), F.normalize(queries, dim=1) - neighbours = min(neighbours, len(train)) - one_hot = F.one_hot(labels, num_classes).float() - parts = [] - for part in queries.split(chunk): - nearest = (part @ train.T).topk(neighbours, dim=1).indices - fractions = one_hot[nearest].mean(dim=1) - parts.append((fractions + 1e-3).log()) - return torch.cat(parts) - - -def linear_logits( - train: torch.Tensor, - labels: torch.Tensor, - queries: torch.Tensor, - *, - num_classes: int, - l2: float, - steps: int = 200, -) -> torch.Tensor: - """Logits ``[queries, num_classes]`` of a multinomial logistic regression - fitted on the standardized ``train`` embeddings.""" - mean, std = train.mean(dim=0), train.std(dim=0).clamp_min(1e-6) - x = (train - mean) / std - linear = nn.Linear(x.shape[1], num_classes).to(x.device) - optimizer = torch.optim.LBFGS( - linear.parameters(), max_iter=steps, line_search_fn="strong_wolfe" - ) - - def closure() -> torch.Tensor: - optimizer.zero_grad() - loss = F.cross_entropy(linear(x), labels) + l2 * linear.weight.square().sum() - loss.backward() - return loss - - with torch.enable_grad(): - optimizer.step(closure) - with torch.no_grad(): - return linear((queries - mean) / std) - - -def probe( - config: Config, - events: PackedEvents, - train: torch.Tensor, - queries: torch.Tensor, - *, - num_classes: int, - device: torch.device, -) -> dict[str, Any]: - """Logits ``[queries, num_classes]`` of the k-NN and linear probes on the - ``queries`` events, given the labelled ``train`` events.""" - encoder = load_encoder(config, events).to(device) - batch_size = config.data.batch_size - embedded = { - name: embed(encoder, events, indices, batch_size=batch_size, device=device) - for name, indices in (("train", train), ("queries", queries)) - } - labels = events.labels[train].to(device, torch.long) - return { - "knn": knn_logits( - embedded["train"], - labels, - embedded["queries"], - num_classes=num_classes, - neighbours=config.model.probe_neighbours, - ), - "linear": linear_logits( - embedded["train"], - labels, - embedded["queries"], - num_classes=num_classes, - l2=config.model.probe_l2, - ), - } - - -def linear_cka(x: torch.Tensor, y: torch.Tensor) -> float: - """Linear CKA of two representations ``[events, d1]`` and ``[events, - d2]`` of the same events: 1 when they are equal up to a rotation and an - isotropic scale, near 0 when unrelated. NaN if either is the same for - all events (e.g. a collapsed encoder), for which it is undefined.""" - x, y = _centered(x), _centered(y) - if x is None or y is None: - return math.nan - norm = (x.T @ x).norm() * (y.T @ y).norm() - return float((y.T @ x).square().sum() / norm) - - -def _centered(x: torch.Tensor) -> torch.Tensor | None: - """``x`` minus its mean over events, in float64; None if what is left is - rounding (the same representation for all events).""" - x = x.detach().cpu().double() - centered = x - x.mean(dim=0) - scale = x.abs().max().clamp_min(1.0) - return None if centered.abs().max() <= 1e-9 * scale else centered - - -@torch.inference_mode() -def layer_embeddings( - encoder: nn.Module, - events: PackedEvents, - indices: torch.Tensor, - *, - batch_size: int, - device: torch.device, -) -> list[torch.Tensor]: - """Event embeddings ``[events, hid_size]`` of ``indices`` after each - layer (:func:`~gnn4colliders.jepa.layer_event_embeddings`).""" - parts = [] - for part in indices.split(batch_size): - batch = collate_sequence_batch(events, part).to(device) - layers = layer_event_embeddings(encoder, batch.tokens, batch.token_mask) - parts.append([layer.float().cpu() for layer in layers]) - return [torch.cat(layer) for layer in zip(*parts)] diff --git a/src/gnn4colliders/sequences.py b/src/gnn4colliders/sequences.py deleted file mode 100644 index cff71761284dd8243dcd93f4c00b09c63c39f96e..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/sequences.py +++ /dev/null @@ -1,77 +0,0 @@ -"""Padded object sequences for the transformer and Deep Sets.""" - -from __future__ import annotations - -import math -from dataclasses import dataclass - -import torch - -from .data import NODE_FEATURE_NAMES, PackedEvents, map_tensors - -PT, ETA, PHI, ENERGY, BTAG, CHARGE = ( - NODE_FEATURE_NAMES.index(name) - for name in ("pt", "eta", "phi", "energy", "btag", "charge") -) - - -@dataclass(frozen=True) -class SequenceBatch: - """``[events, objects, features]`` tokens with a mask of real objects.""" - - tokens: torch.Tensor - token_mask: torch.Tensor - labels: torch.Tensor - weight: torch.Tensor - global_features: torch.Tensor | None - event_index: torch.Tensor - - def to(self, device: torch.device | str, *, non_blocking: bool = False): - return map_tensors(self, lambda t: t.to(device, non_blocking=non_blocking)) - - def pin_memory(self) -> SequenceBatch: - return map_tensors(self, torch.Tensor.pin_memory) - - -def collate_sequence_batch( - events: PackedEvents, indices: torch.Tensor -) -> SequenceBatch: - rows, counts = events.gather_objects(indices) - batch_index = torch.repeat_interleave( - torch.arange(counts.numel()), counts, output_size=rows.shape[0] - ) - position = torch.arange(rows.shape[0]) - (counts.cumsum(0) - counts)[batch_index] - tokens = rows.new_zeros((counts.numel(), int(counts.max()), rows.shape[1])) - tokens[batch_index, position] = rows - token_mask = torch.zeros(tokens.shape[:2], dtype=torch.bool) - token_mask[batch_index, position] = True - return SequenceBatch( - tokens=tokens, - token_mask=token_mask, - labels=events.labels[indices], - weight=events.weight[indices], - global_features=None - if events.global_features is None - else events.global_features[indices], - event_index=indices, - ) - - -def augment_tokens( - tokens: torch.Tensor, - mask: torch.Tensor, - *, - scale: torch.Tensor, - rotation: torch.Tensor, - flip: torch.Tensor, -) -> torch.Tensor: - """``tokens`` with each event's momenta and energies scaled, azimuths - rotated and pseudorapidities multiplied by ``flip`` (1 or -1), all - ``[events]``; padding is left unchanged.""" - augmented = tokens.clone() - for column in (PT, ENERGY): - augmented[..., column] *= scale[:, None] - augmented[..., ETA] *= flip[:, None] - phi = augmented[..., PHI] + rotation[:, None] - augmented[..., PHI] = torch.remainder(phi + math.pi, 2 * math.pi) - math.pi - return torch.where(mask.unsqueeze(-1), augmented, tokens) diff --git a/src/gnn4colliders/tabular.py b/src/gnn4colliders/tabular.py deleted file mode 100644 index edca837c5fc04ab38e4a9496f8f0dcd55b42fb31..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/tabular.py +++ /dev/null @@ -1,705 +0,0 @@ -"""Tabular in-context models on events as fixed-width rows: TabPFN -(Hollmann et al. 2025; Prior Labs), TabICL (Qu et al. 2025, 2026; Inria) -and TabDPT (Ma et al. 2024; Layer 6), which is trained on real tables and -can retrieve each query's nearest context rows. - -Both classify the rows of a table in context, given labelled rows, in one -forward pass and without training. A row of an event (:class:`RowLayout`) -holds the leading ``model.slots[t]`` objects of each object type ``t`` in -stored (pt) order, with their six kinematic columns (``pt, eta, phi, energy, -btag, charge``; NaN for absent objects, which both models handle), then the -number of objects of each type. Optionally (``model.relative_phi``) every -azimuth is measured from that of the event's leading object, which removes -the azimuthal rotation the models cannot learn from a few events, and -(``model.pair_objects``) the row ends with the invariant mass and Delta R of -every pair of the event's leading objects. The context is -``data.train_events`` labelled training events. - -The ``tabpfn``, ``tabicl`` and ``tabdpt`` packages are optional -dependencies (``uv sync --extra tabpfn --extra tabicl --extra tabdpt``). -TabICL's (BSD-3) and TabDPT's (Apache-2.0) weights download from Hugging -Face (``jingang/TabICL``, ``Layer6/TabDPT``). TabPFN-2.5 and later weights -are under a non-commercial license and download after a one-time login (a -token in ``~/.cache/tabpfn`` or ``TABPFN_TOKEN``); ``TABPFN_MODEL_CACHE_DIR`` -sets where they are stored. A checkpoint already in that directory loads -without the login, so the weights can also be downloaded from Hugging Face -(``Prior-Labs/tabpfn_3_5``). - -:class:`TabularPFN` and :class:`TabularICL` meta-train the TabPFN-3.5 and -TabICL networks on the same rows, like the PFN of :mod:`gnn4colliders.pfn` -(``task.name: pfn``, the same tasks and augmentations), and classify in -context with the network's own forward pass instead of the classifier (no -preprocessing beyond TabICL's standardization, no ensembling). By default -the network is frozen and its in-context blocks learn low-rank updates -(``model.lora_rank``). -""" - -from __future__ import annotations - -import dataclasses -import math -import os -import sys -from collections.abc import Callable -from typing import Any - -import torch -import torch.nn.functional as F -from torch import nn -from torch.nn.utils import parametrize - -from .config import ModelConfig -from .data import NODE_FEATURE_NAMES, PackedEvents -from .jepa import event_embedding -from .models import pretrained_encoder -from .pfn import ABSENT_LOGIT, TaskBatch -from .sequences import ENERGY, ETA, PHI, PT, SequenceBatch, collate_sequence_batch - -NODE_TYPE = NODE_FEATURE_NAMES.index("node_type") -#: Log-probability of the classes a classifier gives no probability. -MIN_LOG_PROBABILITY = math.log(1e-12) - - -@dataclasses.dataclass(frozen=True) -class RowLayout: - """What a row of an event holds (see the module documentation).""" - - slots: tuple[int, ...] - relative_phi: bool = False - pair_objects: int = 0 - - def __post_init__(self) -> None: - if not 0 <= self.pair_objects <= sum(self.slots): - raise ValueError("model.pair_objects must be between 0 and sum(slots)") - - @classmethod - def of(cls, config: ModelConfig) -> RowLayout: - return cls(tuple(config.slots), config.relative_phi, config.pair_objects) - - @property - def width(self) -> int: - pairs = self.pair_objects * (self.pair_objects - 1) // 2 - return sum(self.slots) * NODE_TYPE + len(self.slots) + 2 * pairs - - -def _wrap(phi: torch.Tensor) -> torch.Tensor: - return torch.remainder(phi + math.pi, 2 * math.pi) - math.pi - - -def _pair_features(table: torch.Tensor, objects: int) -> torch.Tensor: - """Invariant mass and Delta R ``[events, 2 * pairs]`` of every pair of - the ``objects`` leading slots of ``table [events, slots, columns]`` by pt - (NaN for pairs with an absent object).""" - order = table[..., PT].nan_to_num(-1.0).argsort(dim=1, descending=True) - leading = table.gather(1, order[:, :objects, None].expand(-1, -1, NODE_TYPE)) - first, second = torch.triu_indices(objects, objects, 1, device=table.device) - a, b = leading[:, first], leading[:, second] - - def momentum(x: torch.Tensor) -> torch.Tensor: - pt, eta, phi = x[..., PT], x[..., ETA], x[..., PHI] - return torch.stack( - (pt * torch.cos(phi), pt * torch.sin(phi), pt * torch.sinh(eta)), dim=-1 - ) - - energy = a[..., ENERGY] + b[..., ENERGY] - squared = energy.square() - (momentum(a) + momentum(b)).square().sum(dim=-1) - mass = squared.clamp_min(0).sqrt() - delta_r = torch.hypot(a[..., ETA] - b[..., ETA], _wrap(a[..., PHI] - b[..., PHI])) - return torch.cat((mass, delta_r), dim=1) - - -def object_table( - rows: torch.Tensor, counts: torch.Tensor, layout: RowLayout -) -> torch.Tensor: - """Rows ``[events, layout.width]`` of the events with ``counts`` - consecutive object ``rows`` each.""" - if NODE_TYPE != len(NODE_FEATURE_NAMES) - 1: - raise ValueError("object_table expects node_type as the last column") - slots = layout.slots - device, count, types = rows.device, len(counts), len(slots) - event = torch.repeat_interleave( - torch.arange(count, device=device), counts, output_size=len(rows) - ) - node_type = rows[:, NODE_TYPE].long() - one_hot = F.one_hot(node_type, types) - # Index of every object among its event's objects of the same type. - seen = one_hot.cumsum(0) - first = counts.cumsum(0) - counts - before = (seen - one_hot)[first[event]] - rank = (seen - before).gather(1, node_type[:, None]).squeeze(1) - 1 - capacity = torch.tensor(slots, device=device) - offset = capacity.cumsum(0) - capacity - keep = rank < capacity[node_type] - table = torch.full((count, sum(slots), NODE_TYPE), math.nan, device=device) - table[event[keep], (offset[node_type] + rank)[keep]] = rows[keep, :NODE_TYPE] - if layout.relative_phi: - leading = table[..., PT].nan_to_num(-1.0).argmax(dim=1) - reference = table[torch.arange(count, device=device), leading, PHI] - table[..., PHI] = _wrap(table[..., PHI] - reference[:, None]) - multiplicity = torch.zeros(count, types, device=device) - multiplicity.index_add_(0, event, one_hot.float()) - parts = [table.flatten(1), multiplicity] - if layout.pair_objects > 1: - parts.append(_pair_features(table, layout.pair_objects)) - return torch.cat(parts, dim=1) - - -def table_features( - events: PackedEvents, indices: torch.Tensor, layout: RowLayout -) -> torch.Tensor: - """Rows ``[events, layout.width]`` of ``indices``.""" - return object_table(*events.gather_objects(indices), layout) - - -def sequence_table(batch: SequenceBatch, layout: RowLayout) -> torch.Tensor: - """The rows of :func:`table_features` of a (possibly augmented) batch.""" - mask = batch.token_mask - return object_table(batch.tokens[mask], mask.sum(dim=1), layout) - - -class EventInputs(nn.Module): - """The inputs ``[events, columns]`` of a tabular model: each event's row - (:class:`RowLayout`), its embedding by an object ``encoder`` - (:func:`~gnn4colliders.jepa.event_embedding`), or both - (``model.tabular_inputs``: ``rows``, ``embeddings`` or ``both``). The - encoder is that of ``model.backbone`` with the weights of - ``model.pretrained`` (randomly initialized without), trained with the - LoRA updates in meta-training unless ``model.freeze_backbone``.""" - - def __init__( - self, layout: RowLayout, encoder: nn.Module | None, *, rows: bool - ) -> None: - super().__init__() - if not rows and encoder is None: - raise ValueError("tabular inputs need rows or an encoder") - self.layout, self.encoder, self.rows = layout, encoder, rows - - @classmethod - def of( - cls, config: ModelConfig, *, num_features: int, num_node_types: int - ) -> EventInputs: - if len(config.slots) != num_node_types: - raise ValueError("model.slots needs one count per object type") - inputs = config.tabular_inputs - if inputs not in ("rows", "embeddings", "both"): - raise ValueError( - f"unknown tabular_inputs {inputs!r}; use rows, embeddings or both" - ) - encoder = None - if inputs != "rows": - encoder = pretrained_encoder( - config, num_features=num_features, num_node_types=num_node_types - ) - return cls(RowLayout.of(config), encoder, rows=inputs != "embeddings") - - def forward(self, batch: SequenceBatch) -> torch.Tensor: - parts = [] - if self.rows: - parts.append(sequence_table(batch, self.layout)) - if self.encoder is not None: - mask = batch.token_mask - parts.append(event_embedding(self.encoder, batch.tokens, mask).float()) - return torch.cat(parts, dim=1) - - def of_events( - self, - events: PackedEvents, - indices: torch.Tensor, - *, - batch_size: int, - device: torch.device, - ) -> torch.Tensor: - """The inputs of the events ``indices``, on ``device``.""" - return torch.cat( - [ - self(collate_sequence_batch(events, part).to(device)) - for part in indices.split(batch_size) - ] - ) - - -def _classifier_forward( - classifier: Any, - inputs: EventInputs, - events: PackedEvents, - context: torch.Tensor, - *, - num_classes: int, - batch_size: int, - device: torch.device, - predict: Callable[[Any], Any] | None = None, -) -> Callable[[SequenceBatch], torch.Tensor]: - """Logits (log-probabilities) ``[events, num_classes]`` of a batch by a - scikit-learn ``classifier`` fitted on the ``inputs`` of the labelled - ``context``; its probabilities are ``predict(inputs)`` (default - ``predict_proba``).""" - predict = predict or classifier.predict_proba - inputs = inputs.to(device).eval() - # The context's classes, numbered 0, 1, ... for the classifier. - labels = events.labels[context].long() - classes = labels.unique() - with torch.no_grad(): - rows = inputs.of_events(events, context, batch_size=batch_size, device=device) - classifier.fit(rows.cpu().numpy(), torch.searchsorted(classes, labels).numpy()) - - def forward(batch: SequenceBatch) -> torch.Tensor: - with torch.no_grad(): - rows = inputs(batch).cpu() - probability = torch.as_tensor(predict(rows.numpy())) - logits = torch.full((len(rows), num_classes), MIN_LOG_PROBABILITY) - logits[:, classes] = probability.float().clamp_min(1e-12).log() - return logits.to(device) - - return forward - - -def tabpfn_forward( - config: ModelConfig, - events: PackedEvents, - context: torch.Tensor, - *, - num_classes: int, - num_node_types: int, - batch_size: int, - device: torch.device, - seed: int, -) -> Callable[[SequenceBatch], torch.Tensor]: - """A function from a batch of events to their logits (log-probabilities) - ``[events, num_classes]`` by TabPFN given the labelled ``context``.""" - try: - from tabpfn import TabPFNClassifier - from tabpfn.constants import ModelVersion - except ImportError as error: - raise ImportError("tabpfn is not installed: uv sync --extra tabpfn") from error - checkpoint = config.tabpfn_checkpoint - classifier = TabPFNClassifier.create_default_for_version( - ModelVersion(config.tabpfn_version), - n_estimators=config.tabpfn_estimators, - device=str(device), - random_state=seed, - # The context is processed once for all batches of events. - fit_mode="fit_with_cache", - **({} if checkpoint is None else {"model_path": checkpoint}), - ) - return _classifier_forward( - classifier, - EventInputs.of( - config, num_features=events.num_features, num_node_types=num_node_types - ), - events, - context, - num_classes=num_classes, - batch_size=batch_size, - device=device, - ) - - -def tabicl_forward( - config: ModelConfig, - events: PackedEvents, - context: torch.Tensor, - *, - num_classes: int, - num_node_types: int, - batch_size: int, - device: torch.device, - seed: int, -) -> Callable[[SequenceBatch], torch.Tensor]: - """A function from a batch of events to their logits (log-probabilities) - ``[events, num_classes]`` by TabICL given the labelled ``context``.""" - try: - from tabicl import TabICLClassifier - except ImportError as error: - raise ImportError("tabicl is not installed: uv sync --extra tabicl") from error - classifier = TabICLClassifier( - n_estimators=config.tabicl_estimators, - checkpoint_version=config.tabicl_checkpoint, - device=str(device), - random_state=seed, - # The context is processed once for all batches of events. - kv_cache=True, - ) - return _classifier_forward( - classifier, - EventInputs.of( - config, num_features=events.num_features, num_node_types=num_node_types - ), - events, - context, - num_classes=num_classes, - batch_size=batch_size, - device=device, - ) - - -def tabdpt_forward( - config: ModelConfig, - events: PackedEvents, - context: torch.Tensor, - *, - num_classes: int, - num_node_types: int, - batch_size: int, - device: torch.device, - seed: int, -) -> Callable[[SequenceBatch], torch.Tensor]: - """A function from a batch of events to their logits (log-probabilities) - ``[events, num_classes]`` by TabDPT given the labelled ``context``; with - ``model.tabdpt_context_size`` each query's context is its nearest rows. - - TabDPT computes in float64 in places, so on MPS it runs on the CPU. Its - retrieval (faiss) and PyTorch load separate OpenMP runtimes on macOS, - which crash together unless ``OMP_NUM_THREADS=1``.""" - try: - from tabdpt import TabDPTClassifier - except ImportError as error: - raise ImportError("tabdpt is not installed: uv sync --extra tabdpt") from error - size = config.tabdpt_context_size - if size is not None and sys.platform == "darwin": - if os.environ.get("OMP_NUM_THREADS") != "1": - raise ValueError( - "TabDPT retrieval on macOS needs OMP_NUM_THREADS=1 (faiss and " - "torch load separate OpenMP runtimes)" - ) - classifier = TabDPTClassifier( - # Absent objects are NaN: imputed, with indicator columns. - missing_indicators=True, - context_reduction="subsample" if size is None else "retrieval", - device="cpu" if device.type == "mps" else str(device), - verbose=False, - ) - - def predict(rows: Any) -> Any: - return classifier.ensemble_predict_proba( - rows, n_ensembles=config.tabdpt_estimators, context_size=size, seed=seed - ) - - return _classifier_forward( - classifier, - EventInputs.of( - config, num_features=events.num_features, num_node_types=num_node_types - ), - events, - context, - num_classes=num_classes, - batch_size=batch_size, - device=device, - predict=predict, - ) - - -# --- Meta-training of the tabular networks ------------------------------------ - - -def load_tabpfn_network(version: str, checkpoint: str | None) -> nn.Module: - """The TabPFN-3.5 network of ``model.tabpfn_version`` and - ``model.tabpfn_checkpoint`` (default: that version's).""" - try: - from tabpfn.model_loading import load_model_criterion_config - except ImportError as error: - raise ImportError("tabpfn is not installed: uv sync --extra tabpfn") from error - if version not in ("v3.5", "v3.5-fast"): - raise ValueError("TabularPFN trains a TabPFN-3.5 network (v3.5, v3.5-fast)") - models, *_ = load_model_criterion_config( - checkpoint, - check_bar_distribution_criterion=False, - estimator_type="classifier", - version=version, - download_if_not_exists=True, - ) - return models[0] - - -def load_tabicl_network(checkpoint: str) -> nn.Module: - """The TabICL network of ``model.tabicl_checkpoint`` (a path, or a file - name in the Hugging Face repository ``jingang/TabICL``), recomputing - activations in the backward pass (gradient checkpointing).""" - try: - from huggingface_hub import hf_hub_download - from tabicl._model.tabicl import TabICL - except ImportError as error: - raise ImportError("tabicl is not installed: uv sync --extra tabicl") from error - path = checkpoint if os.path.exists(checkpoint) else None - path = path or hf_hub_download(repo_id="jingang/TabICL", filename=checkpoint) - saved = torch.load(path, map_location="cpu", weights_only=True) - network = TabICL(**{**saved["config"], "recompute": True}) - network.load_state_dict(saved["state_dict"]) - return network - - -class LoRA(nn.Module): - """A parametrization that adds a trainable low-rank update ``up @ down`` - (LoRA; Hu et al. 2022), zero at initialization, to a frozen weight. As a - parametrization it also reaches weights the code reads directly, such as - TabICL's fused attention projections.""" - - def __init__(self, weight: torch.Tensor, rank: int) -> None: - super().__init__() - out_features, in_features = weight.shape - down = torch.randn(rank, in_features) / math.sqrt(in_features) - self.down = nn.Parameter(down.to(weight)) - self.up = nn.Parameter(weight.new_zeros(out_features, rank)) - - def forward(self, weight: torch.Tensor) -> torch.Tensor: - return weight + (self.up @ self.down).to(weight.dtype) - - -def add_lora(network: nn.Module, weights: list[str], rank: int) -> None: - """Freeze ``network`` and give its ``weights`` (dotted names) LoRA - updates of ``rank``.""" - network.requires_grad_(False) - for name in weights: - owner_name, _, attribute = name.rpartition(".") - owner = network.get_submodule(owner_name) - weight = getattr(owner, attribute).detach() - parametrize.register_parametrization(owner, attribute, LoRA(weight, rank)) - - -def merge_lora(network: nn.Module) -> None: - """Fold the LoRA updates into the weights, for inference (the network - can no longer be trained or saved as LoRA updates).""" - for module in list(network.modules()): - if parametrize.is_parametrized(module): - for name in list(module.parametrizations): - parametrize.remove_parametrizations( - module, name, leave_parametrized=True - ) - - -class MetaTabular(nn.Module): - """Query logits ``[tasks * queries, max_classes]`` of a - :class:`~gnn4colliders.pfn.TaskBatch` by a tabular in-context - ``network`` on the events' :class:`EventInputs` (interface of - :class:`~gnn4colliders.pfn.EventPFN`); the logits of classes a task does - not have are ``ABSENT_LOGIT``. - - With ``lora_rank > 0`` only the LoRA updates of the network's - :attr:`lora_weights` train (:func:`add_lora`), else all weights do; an - encoder of the inputs trains too unless frozen. The state dict holds only - the trainable weights (and buffers): the rest are those of the released - network and of ``model.pretrained``. - """ - - def __init__( - self, - network: nn.Module, - *, - inputs: EventInputs, - max_classes: int, - lora_rank: int, - ) -> None: - super().__init__() - self.network, self.inputs, self.max_classes = network, inputs, max_classes - if lora_rank > 0: - add_lora(network, self.lora_weights(), lora_rank) - self.register_state_dict_post_hook(_drop_frozen) - self.register_load_state_dict_post_hook(_allow_missing_frozen) - - def lora_weights(self) -> list[str]: - """The weights of the network's in-context blocks that LoRA updates.""" - raise NotImplementedError - - def task_logits(self, rows: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: - """Logits ``[tasks, queries, classes]`` of tasks of ``rows [tasks, - context + queries, columns]`` given the context ``labels [tasks, - context]``.""" - raise NotImplementedError - - def in_context( - self, rows: torch.Tensor, labels: torch.Tensor - ) -> Callable[[torch.Tensor], torch.Tensor]: - """A function from query rows to their logits ``[queries, classes]`` - given the context ``rows`` and ``labels`` (one task, computed once).""" - raise NotImplementedError - - def frozen(self) -> set[str]: - return { - name for name, value in self.named_parameters() if not value.requires_grad - } - - def forward(self, batch: TaskBatch) -> torch.Tensor: - tasks = batch.context_labels.shape[0] - rows = self.inputs(batch.events) - rows = rows.view(tasks, -1, rows.shape[1]) - logits = self.task_logits(rows, batch.context_labels) - logits = logits[..., : self.max_classes].float() - classes = torch.arange(self.max_classes, device=logits.device) - absent = classes >= batch.num_classes.reshape(-1, 1, 1) - return logits.masked_fill(absent, ABSENT_LOGIT).flatten(0, 1) - - -class TabularPFN(MetaTabular): - """:class:`MetaTabular` of the TabPFN-3.5 network: LoRA updates of the - attention projections and MLP layers of its 24 in-context blocks.""" - - def lora_weights(self) -> list[str]: - names = [] - for index in range(len(self.network.icl_blocks)): - block = f"icl_blocks.{index}" - names += [ - f"{block}.icl_attention.{name}_projection.weight" - for name in ("q", "k", "v", "out") - ] - names += [f"{block}.mlp.0.weight", f"{block}.mlp.2.weight"] - return names - - def options(self) -> Any: - """TabPFN's performance options; with gradients, activations are - recomputed in the backward pass (gradient checkpointing).""" - options = self.network.get_default_performance_options() - recompute = torch.is_grad_enabled() - return dataclasses.replace(options, force_recompute_layer=recompute) - - def task_logits(self, rows: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: - x = rows.transpose(0, 1) # TabPFN: [rows, tasks, columns] - y = labels.T.to(x.dtype) - output = self.network(x, y, "multiclass", performance_options=self.options()) - return output.transpose(0, 1) - - def in_context( - self, rows: torch.Tensor, labels: torch.Tensor - ) -> Callable[[torch.Tensor], torch.Tensor]: - y = labels.to(rows.dtype)[:, None] - options = self.options() - _, cache = self.network( - rows[:, None], - y, - "multiclass", - performance_options=options, - return_kv_cache=True, - ) - - def forward(queries: torch.Tensor) -> torch.Tensor: - output = self.network( - queries[:, None], - y, - "multiclass", - performance_options=options, - kv_cache=cache, - x_is_test_only=True, - ) - return output[:, 0] - - return forward - - -#: Standardized TabICL inputs are clipped to this many standard deviations. -TABICL_CLIP = 8.0 - - -def _standardizer(context: torch.Tensor) -> Callable[[torch.Tensor], torch.Tensor]: - """Standardization of rows ``[..., rows, columns]`` by the mean and - standard deviation of each task's ``context`` rows, absent (NaN) values - at the mean, clipped at ``TABICL_CLIP``: TabICL's network takes no NaN - (its classifier preprocesses similarly, without ensembling).""" - mean = context.nanmean(dim=-2, keepdim=True) - std = (context - mean).square().nanmean(dim=-2, keepdim=True).sqrt() - mean, std = mean.nan_to_num(0.0), std.nan_to_num(1.0).clamp_min(1e-6) - - def standardize(rows: torch.Tensor) -> torch.Tensor: - z = ((rows - mean) / std).nan_to_num(0.0) - return z.clamp(-TABICL_CLIP, TABICL_CLIP) - - return standardize - - -class TabularICL(MetaTabular): - """:class:`MetaTabular` of the TabICL network: LoRA updates of the - attention projections and feed-forward layers of its 12 in-context - blocks. The rows are standardized by :func:`_standardizer`.""" - - def lora_weights(self) -> list[str]: - names = [] - for index in range(len(self.network.icl_predictor.tf_icl.blocks)): - block = f"icl_predictor.tf_icl.blocks.{index}" - names += [ - f"{block}.attn.in_proj_weight", - f"{block}.attn.out_proj.weight", - f"{block}.linear1.weight", - f"{block}.linear2.weight", - ] - return names - - def task_logits(self, rows: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: - context = labels.shape[1] - x = _standardizer(rows[:, :context])(rows) - # The training forward, in validation too: its logits cover - # max_classes whatever the context's classes. In eval mode TabICL's - # modules switch to their inference path, which chooses its own - # device; its dropout is 0, so training mode changes nothing else. - training = self.network.training - self.network.train() - try: - return self.network._train_forward(x, labels.to(x.dtype)) - finally: - self.network.train(training) - - def in_context( - self, rows: torch.Tensor, labels: torch.Tensor - ) -> Callable[[torch.Tensor], torch.Tensor]: - from tabicl import InferenceConfig - - # On the model's device (TabICL picks one itself by default) and in - # float32, as in meta-training: mixed precision on the CPU changes - # the predictions substantially. - inference = InferenceConfig() - for stage in (inference.COL_CONFIG, inference.ROW_CONFIG, inference.ICL_CONFIG): - stage.update({"device": rows.device, "use_amp": False}) - standardize = _standardizer(rows) - self.network.forward_with_cache( - X_train=standardize(rows)[None], - y_train=labels.to(rows.dtype)[None], - store_cache=True, - inference_config=inference, - ) - - def forward(queries: torch.Tensor) -> torch.Tensor: - logits = self.network.forward_with_cache( - X_test=standardize(queries)[None], - use_cache=True, - store_cache=False, - return_logits=True, - inference_config=inference, - ) - return logits[0] - - return forward - - -def _drop_frozen( - module: MetaTabular, state: dict[str, Any], prefix: str, metadata: Any -) -> None: - for name in module.frozen(): - del state[prefix + name] - - -def _allow_missing_frozen(module: MetaTabular, keys: Any) -> None: - frozen = module.frozen() - keys.missing_keys[:] = [key for key in keys.missing_keys if key not in frozen] - - -@torch.inference_mode() -def meta_tabular_forward( - model: MetaTabular, - events: PackedEvents, - context: torch.Tensor, - *, - num_classes: int, - batch_size: int, - device: torch.device, -) -> Callable[[SequenceBatch], torch.Tensor]: - """A function from a batch of events to their logits ``[events, - num_classes]``, given the labelled ``context`` events (one task). The - context is processed once.""" - if num_classes > model.max_classes: - raise ValueError(f"the model tells apart at most {model.max_classes} classes") - model.eval() - merge_lora(model.network) - rows = model.inputs.of_events(events, context, batch_size=batch_size, device=device) - labels = events.labels[context].to(device) - logits = model.in_context(rows, labels) - - def forward(batch: SequenceBatch) -> torch.Tensor: - return logits(model.inputs(batch))[:, :num_classes] - - return forward diff --git a/src/gnn4colliders/tasks.py b/src/gnn4colliders/tasks.py deleted file mode 100644 index eac092db6faefdc1ce6992f31a98ec762d67f3ec..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/tasks.py +++ /dev/null @@ -1,235 +0,0 @@ -"""Classification losses and metrics. - -Losses are the ROOT-GNN class-balanced form: a weighted mean per class -present in the batch, averaged equally over those classes (not ordinary -weighted cross-entropy). Metrics exclude events with non-positive weight. -""" - -from __future__ import annotations - -import torch -import torch.nn.functional as F - - -def per_class_weighted_mean( - elementwise_loss: torch.Tensor, - labels: torch.Tensor, - weights: torch.Tensor, - num_classes: int, -) -> torch.Tensor: - """Mean over present classes of each class's weighted mean loss. - - Dense per-class sums need no host synchronization, so the loss can be - captured in a CUDA graph. - """ - one_hot = F.one_hot(labels, num_classes).to(elementwise_loss.dtype) - present = one_hot.sum(dim=0) > 0 - numerator = (one_hot * (weights * elementwise_loss)[:, None]).sum(dim=0) - denominator = (one_hot * weights[:, None]).sum(dim=0) - # Absent classes would divide 0 by 0; keep NaN out of the gradient. - means = numerator / torch.where(present, denominator, torch.ones_like(denominator)) - return torch.where(present, means, torch.zeros_like(means)).sum() / present.sum() - - -def float64_device(tensor: torch.Tensor) -> torch.device: - """Where float64 metrics of ``tensor`` run: the CPU for MPS, which has no - float64, else its own device. Move, then cast: on PyTorch 2.14 a combined - ``.to(device="cpu", dtype=torch.float64)`` of an MPS tensor gives zeros.""" - return torch.device("cpu") if tensor.device.type == "mps" else tensor.device - - -def weighted_binary_auc( - labels: torch.Tensor, scores: torch.Tensor, weights: torch.Tensor -) -> float: - """Weighted ROC AUC with half credit for tied scores, in ``O(n log n)``. - - Equals the weighted fraction of (positive, negative) pairs ordered - correctly. Returns NaN without positive or negative weight. - """ - device = float64_device(scores) - labels = labels.detach().to(device).to(torch.bool).reshape(-1) - scores = scores.detach().to(device).to(torch.float64).reshape(-1) - weights = weights.detach().to(device).to(torch.float64).reshape(-1) - zero = weights.new_zeros(()) - positive = torch.where(labels, weights, zero) - negative = torch.where(labels, zero, weights) - positive_weight, negative_weight = positive.sum(), negative.sum() - if positive_weight <= 0 or negative_weight <= 0: - return float("nan") - order = torch.argsort(scores) - sorted_scores = scores[order] - # Cumulative weights at the last element of every group of equal scores. - group_end = torch.ones_like(sorted_scores, dtype=torch.bool) - group_end[:-1] = sorted_scores[1:] != sorted_scores[:-1] - cumulative_positive = torch.cumsum(positive[order], 0)[group_end] - cumulative_negative = torch.cumsum(negative[order], 0)[group_end] - group_positive = torch.diff(cumulative_positive, prepend=zero.reshape(1)) - group_negative = torch.diff(cumulative_negative, prepend=zero.reshape(1)) - below = cumulative_negative - group_negative - pairs = (group_positive * (below + 0.5 * group_negative)).sum() - return float(pairs / (positive_weight * negative_weight)) - - -def macro_scores( - labels: torch.Tensor, predictions: torch.Tensor, num_classes: int -) -> dict[str, float]: - """Unweighted macro precision, recall and F1 (absent classes count 0).""" - confusion = ( - torch.bincount( - labels.to(torch.long) * num_classes + predictions.to(torch.long), - minlength=num_classes * num_classes, - ) - .reshape(num_classes, num_classes) - .to(float64_device(labels)) - .to(torch.float64) - ) - hits, predicted, actual = confusion.diagonal(), confusion.sum(0), confusion.sum(1) - precision = torch.where(predicted > 0, hits / predicted, 0.0) - recall = torch.where(actual > 0, hits / actual, 0.0) - total = precision + recall - f1 = torch.where(total > 0, 2 * precision * recall / total, 0.0) - return { - "precision": float(precision.mean()), - "recall": float(recall.mean()), - "f1": float(f1.mean()), - } - - -class MulticlassClassificationTask: - """``[events, classes]`` logits with integer labels.""" - - def __init__(self, *, absolute_weights: bool = False) -> None: - self.absolute_weights = absolute_weights - - def loss( - self, logits: torch.Tensor, labels: torch.Tensor, weights: torch.Tensor - ) -> torch.Tensor: - labels = labels.to(torch.long) - weights = weights.abs() if self.absolute_weights else weights - elementwise = F.cross_entropy(logits, labels, reduction="none") - return per_class_weighted_mean(elementwise, labels, weights, logits.shape[1]) - - @staticmethod - def scores(logits: torch.Tensor) -> torch.Tensor: - return torch.softmax(logits, dim=1) - - @staticmethod - def predictions(scores: torch.Tensor) -> torch.Tensor: - return scores.argmax(dim=1) - - def metrics( - self, logits: torch.Tensor, labels: torch.Tensor, weights: torch.Tensor - ) -> dict[str, float]: - """Accuracy, one-vs-rest AUC averaged over classes, macro P/R/F1.""" - scores = self.scores(logits) - predictions = self.predictions(scores) - labels = labels.to(torch.long) - mask = weights > 0 - metric_weights = (weights.abs() if self.absolute_weights else weights)[mask] - aucs = [ - weighted_binary_auc( - labels[mask] == index, scores[mask, index], metric_weights - ) - for index in range(scores.shape[1]) - ] - return { - "accuracy": float((predictions == labels).float().mean()), - "auc": sum(aucs) / len(aucs), - **macro_scores(labels, predictions, scores.shape[1]), - } - - -class BinaryClassificationTask: - """One logit per event with 0/1 labels.""" - - def __init__( - self, *, absolute_weights: bool = False, threshold: float = 0.5 - ) -> None: - self.absolute_weights = absolute_weights - self.threshold = threshold - - def loss( - self, logits: torch.Tensor, labels: torch.Tensor, weights: torch.Tensor - ) -> torch.Tensor: - logits = logits.reshape(-1) - labels = labels.reshape(-1) - weights = weights.abs() if self.absolute_weights else weights - elementwise = F.binary_cross_entropy_with_logits( - logits, labels.to(logits.dtype), reduction="none" - ) - return per_class_weighted_mean(elementwise, labels.to(torch.long), weights, 2) - - @staticmethod - def scores(logits: torch.Tensor) -> torch.Tensor: - return torch.sigmoid(logits.reshape(-1)) - - def predictions(self, scores: torch.Tensor) -> torch.Tensor: - return scores >= self.threshold - - def metrics( - self, logits: torch.Tensor, labels: torch.Tensor, weights: torch.Tensor - ) -> dict[str, float]: - scores = self.scores(logits) - labels = labels.reshape(-1) == 1 - mask = weights > 0 - metric_weights = weights.abs() if self.absolute_weights else weights - return { - "accuracy": float((self.predictions(scores) == labels).float().mean()), - "auc": weighted_binary_auc( - labels[mask], scores[mask], metric_weights[mask] - ), - } - - -class PretrainingTask: - """The loss of JEPA or contrastive pretraining. The model's outputs per - event are ``[loss, counts, spread]`` (:class:`~gnn4colliders.jepa.JEPA`, - :class:`~gnn4colliders.contrastive.Contrastive`), the spread reported as - ``spread_name``; labels and weights are not used.""" - - def __init__(self, spread_name: str) -> None: - self.spread_name = spread_name - - @staticmethod - def loss( - outputs: torch.Tensor, labels: torch.Tensor, weight: torch.Tensor - ) -> torch.Tensor: - has_targets = outputs[:, 1] - return (outputs[:, 0] * has_targets).sum() / has_targets.sum().clamp_min(1) - - def metrics( - self, outputs: torch.Tensor, labels: torch.Tensor, weight: torch.Tensor - ) -> dict[str, float]: - return {self.spread_name: float(outputs[:, 2].mean())} - - -class PfnTask(MulticlassClassificationTask): - """Meta-training of the PFN (:mod:`gnn4colliders.pfn`): the multiclass - loss of the query events. Their labels are numbered within each task, so - only the accuracy is meaningful over a batch of tasks.""" - - def metrics( - self, logits: torch.Tensor, labels: torch.Tensor, weights: torch.Tensor - ) -> dict[str, float]: - predictions = logits.argmax(dim=1) - return { - "accuracy": float((predictions == labels.to(torch.long)).float().mean()) - } - - -def build_task(name: str, *, absolute_weights: bool = False, threshold: float = 0.5): - if name == "multiclass": - return MulticlassClassificationTask(absolute_weights=absolute_weights) - if name == "binary": - return BinaryClassificationTask( - absolute_weights=absolute_weights, threshold=threshold - ) - if name == "jepa": - return PretrainingTask("target_spread") - if name == "contrastive": - return PretrainingTask("embedding_spread") - if name == "pfn": - return PfnTask() - raise ValueError( - f"unknown task {name!r}; use multiclass, binary, jepa, contrastive or pfn" - ) diff --git a/src/gnn4colliders/train.py b/src/gnn4colliders/train.py deleted file mode 100644 index c8dc5c2ba3afc0217cb57c4e2e97d00f4f208043..0000000000000000000000000000000000000000 --- a/src/gnn4colliders/train.py +++ /dev/null @@ -1,328 +0,0 @@ -"""Training loop, early stopping and checkpoints.""" - -from __future__ import annotations - -import logging -import math -import os -import random -from collections.abc import Callable -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -import numpy as np -import torch -from torch import nn - -from .distributed import DistributedContext, broadcast_bool, gather_tensor, wrap_model -from .inference import collect_outputs -from .models import load_weights - -logger = logging.getLogger(__name__) - -#: ``train.precision`` values and their autocast dtype (``None``: float32). -PRECISIONS = {"float32": None, "bfloat16": torch.bfloat16} -#: Eager steps before the first CUDA graph capture (lazy CUDA and optimizer -#: initialization must not happen during capture). -CUDA_GRAPH_WARMUP_STEPS = 3 - - -def seed_everything(seed: int) -> None: - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - - -@dataclass(frozen=True) -class EpochResult: - loss: float - metrics: dict[str, float] = field(default_factory=dict) - num_samples: int = 0 - - -class EarlyStopping: - """Stop after ``patience`` epochs without improvement by ``min_delta``. - - ``monitor`` is ``train_loss``, ``loss`` (validation loss) or a - validation metric name; losses are minimized and metrics maximized. - """ - - def __init__( - self, monitor: str = "loss", patience: int = 5, min_delta: float = 0.0 - ): - self.monitor, self.patience, self.min_delta = monitor, patience, min_delta - self.maximize = monitor not in ("loss", "train_loss") - self.best = -math.inf if self.maximize else math.inf - self.num_bad_epochs = 0 - - def update(self, value: float) -> bool: - """Record one epoch's value; return whether to stop.""" - if self.maximize: - improved = value > self.best + self.min_delta - else: - improved = value < self.best - self.min_delta - if improved: - self.best, self.num_bad_epochs = value, 0 - else: - self.num_bad_epochs += 1 - return self.num_bad_epochs >= self.patience - - def value(self, train: EpochResult, validation: EpochResult | None) -> float: - if self.monitor == "train_loss": - return train.loss - if validation is None: - raise ValueError( - f"early stopping on {self.monitor!r} needs validation data" - ) - return ( - validation.loss - if self.monitor == "loss" - else validation.metrics[self.monitor] - ) - - def state_dict(self) -> dict[str, Any]: - return {"best": self.best, "num_bad_epochs": self.num_bad_epochs} - - def load_state_dict(self, state: dict[str, Any]) -> None: - self.best, self.num_bad_epochs = ( - float(state["best"]), - int(state["num_bad_epochs"]), - ) - - -class Trainer: - """Train ``model`` on ``task`` with ``optimizer`` and an optional - per-epoch ``scheduler``. - - Opt-in speedups: ``compile_model`` (``torch.compile`` with dynamic - shapes), ``cuda_graphs`` (replayed training steps, single CUDA process) - and ``precision="bfloat16"`` (autocast forward; loss, graph sums and - optimizer stay float32). - """ - - def __init__( - self, - model: nn.Module, - task: Any, - optimizer: torch.optim.Optimizer, - scheduler: Any | None = None, - *, - device: torch.device | str = "cpu", - context: DistributedContext | None = None, - early_stopping: EarlyStopping | None = None, - compile_model: bool = False, - cuda_graphs: bool = False, - precision: str = "float32", - ) -> None: - if precision not in PRECISIONS: - raise ValueError(f"precision must be one of {sorted(PRECISIONS)}") - self.context = context or DistributedContext(device=torch.device(device)) - self.device = self.context.device - self.model = model.to(self.device) - self.task, self.optimizer, self.scheduler = task, optimizer, scheduler - self.early_stopping = early_stopping - self.epoch, self.global_step = -1, 0 # last completed epoch - self._autocast_dtype = PRECISIONS[precision] - wrapped = wrap_model(self.model, self.context) - self._module = ( - torch.compile(wrapped, dynamic=True) if compile_model else wrapped - ) - self._graph_step = None - if cuda_graphs: - if self.context.enabled or self.device.type != "cuda": - raise ValueError("CUDA graph training needs a single CUDA process") - from .cuda_graphs import CudaGraphTrainStep - - self._graph_step = CudaGraphTrainStep(self.forward, task, optimizer) - - def forward(self, batch: Any) -> torch.Tensor: - """Float32 logits of ``batch``, computed under autocast if requested.""" - if self._autocast_dtype is None: - return self._module(batch) - with torch.autocast(device_type=self.device.type, dtype=self._autocast_dtype): - return self._module(batch).float() - - def _step(self, batch: Any) -> tuple[torch.Tensor, torch.Tensor]: - """One optimizer step. Nothing here waits for the device, so the host - can prepare the next batch while this one runs.""" - if ( - self._graph_step is not None - and self.global_step >= CUDA_GRAPH_WARMUP_STEPS - and hasattr(batch, "graph") - ): - return self._graph_step(batch) - self.optimizer.zero_grad(set_to_none=True) - logits = self.forward(batch) - loss = self.task.loss(logits, batch.labels, batch.weight) - loss.backward() - self.optimizer.step() - if hasattr(self.model, "update_target"): # JEPA's EMA target encoder - self.model.update_target() - return loss.detach(), logits.detach() - - def train_epoch(self, loader: Any) -> EpochResult: - """One pass over ``loader``. The loss is the mean of batch losses - (single process) or the loss over all ranks' events (DDP); metrics - cover every event of the epoch.""" - self.model.train() - losses, logits, labels, weights = [], [], [], [] - for batch in loader: - # Pinned batches (EventLoader pin_memory) copy asynchronously. - batch = batch.to(self.device, non_blocking=True) - loss, batch_logits = self._step(batch) - self.global_step += 1 - losses.append(loss) - logits.append(batch_logits) - labels.append(batch.labels) - weights.append(batch.weight) - if not losses: - raise ValueError("cannot train on an empty loader") - # Batch losses are checked once per epoch to avoid a per-step sync. - values = torch.stack(losses).cpu().tolist() - bad = [index for index, value in enumerate(values) if not math.isfinite(value)] - if bad: - raise FloatingPointError(f"non-finite training loss at batch {bad[0]}") - logits, labels, weights = ( - gather_tensor(torch.cat(parts), self.context) - for parts in (logits, labels, weights) - ) - loss = ( - float(self.task.loss(logits, labels, weights)) - if self.context.enabled - else sum(values) / len(values) - ) - metrics = self.task.metrics(logits, labels, weights) - return EpochResult(loss, metrics, int(labels.shape[0])) - - def evaluate(self, loader: Any) -> EpochResult: - self.model.eval() - outputs = collect_outputs( - self.forward, loader, device=self.device, context=self.context - ) - loss = self.task.loss(outputs.logits, outputs.labels, outputs.weight) - metrics = self.task.metrics(outputs.logits, outputs.labels, outputs.weight) - return EpochResult(float(loss), metrics, int(outputs.labels.shape[0])) - - def fit( - self, - train_loader: Any, - validation_loader: Any | None = None, - *, - max_epochs: int, - on_epoch_end: Callable[[EpochResult, EpochResult | None], None] | None = None, - ) -> list[tuple[EpochResult, EpochResult | None]]: - """Train until epoch ``max_epochs - 1`` or early stopping. - - A resumed trainer continues after its last completed epoch. - ``on_epoch_end(train, validation)`` runs on every rank once the - epoch, early stopping and scheduler reflect it, so a checkpoint - saved there is resumable. - """ - history = [] - for epoch in range(self.epoch + 1, max_epochs): - for loader in (train_loader, validation_loader): - if loader is not None: - loader.set_epoch(epoch) - train = self.train_epoch(train_loader) - validation = ( - None if validation_loader is None else self.evaluate(validation_loader) - ) - history.append((train, validation)) - if self.context.is_main_process: - logger.info( - "epoch %d: train loss %.5f %s, validation %s", - epoch, - train.loss, - train.metrics, - None - if validation is None - else (validation.loss, validation.metrics), - ) - self.epoch = epoch - stop = False - if self.early_stopping is not None: - value = self.early_stopping.value(train, validation) - stop = broadcast_bool(self.early_stopping.update(value), self.context) - # The scheduler is not stepped after the final, stopping epoch. - if self.scheduler is not None and not stop: - self.scheduler.step() - if on_epoch_end is not None: - on_epoch_end(train, validation) - if stop: - break - return history - - -# --- Checkpoints -------------------------------------------------------------- - - -def _portable_optimizer_state(optimizer: torch.optim.Optimizer) -> dict[str, Any]: - """Optimizer state without CUDA-graph settings (tensor learning rates and - ``capturable``), so any training mode can resume from it.""" - state = optimizer.state_dict() - groups = [] - for group in state["param_groups"]: - group = dict(group) - if torch.is_tensor(group.get("lr")): - group["lr"] = float(group["lr"]) - if "capturable" in group: - group["capturable"] = False - groups.append(group) - return {**state, "param_groups": groups} - - -def save_checkpoint(path: str | Path, trainer: Trainer, **extra: Any) -> Path: - """Write a resumable checkpoint atomically, with ``extra`` entries.""" - path = Path(path) - path.parent.mkdir(parents=True, exist_ok=True) - payload = { - "epoch": trainer.epoch, - "global_step": trainer.global_step, - "model_state_dict": { - k: v.detach().cpu() for k, v in trainer.model.state_dict().items() - }, - "optimizer_state_dict": _portable_optimizer_state(trainer.optimizer), - "scheduler_state_dict": None - if trainer.scheduler is None - else trainer.scheduler.state_dict(), - "early_stopping_state": None - if trainer.early_stopping is None - else trainer.early_stopping.state_dict(), - "rng_state": { - "python": random.getstate(), - "numpy": np.random.get_state(), - "torch": torch.get_rng_state(), - "cuda": torch.cuda.get_rng_state_all() - if torch.cuda.is_available() - else None, - }, - **extra, - } - temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp") - torch.save(payload, temporary) - os.replace(temporary, path) - return path - - -def load_checkpoint(path: str | Path) -> dict[str, Any]: - return torch.load(path, map_location="cpu", weights_only=False) - - -def restore(trainer: Trainer, checkpoint: dict[str, Any]) -> None: - """Resume ``trainer`` from a :func:`save_checkpoint` payload.""" - load_weights(trainer.model, checkpoint) - trainer.optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) - if trainer.scheduler is not None and checkpoint.get("scheduler_state_dict"): - trainer.scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) - if trainer.early_stopping is not None and checkpoint.get("early_stopping_state"): - trainer.early_stopping.load_state_dict(checkpoint["early_stopping_state"]) - trainer.epoch = int(checkpoint["epoch"]) - trainer.global_step = int(checkpoint.get("global_step", 0)) - rng = checkpoint.get("rng_state") - if rng: - random.setstate(rng["python"]) - np.random.set_state(rng["numpy"]) - torch.set_rng_state(rng["torch"]) - if rng.get("cuda") is not None and torch.cuda.is_available(): - torch.cuda.set_rng_state_all(rng["cuda"]) diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index b830537d7428890b4616040a82198fad6acb5ecf..0000000000000000000000000000000000000000 --- a/tests/conftest.py +++ /dev/null @@ -1,68 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import pytest -import torch - -from gnn4colliders.config import Config, load_config -from gnn4colliders.data import PackedEvents, pack_events, prepare - -ROOT = Path(__file__).resolve().parents[1] -FIXTURES = ROOT / "tests/fixtures" -#: 64-event fixtures of the five GamGam samples, labels 0..4 (ttH..ZH). -GAMGAM_FIXTURES = sorted((FIXTURES / "atlas_opendata_gamgam").glob("mc_*_64.root")) -TTH_FIXTURE = FIXTURES / "ttH_NLO_64.root" -PRETRAINED = ROOT / "Pretrained_GNN" -#: ``train.device`` of the fixture runs: ``pytest --device mps`` (default cpu). -DEVICE = "cpu" - - -def pytest_addoption(parser: pytest.Parser) -> None: - parser.addoption("--device", default="cpu", help="train.device of fixture runs") - - -def pytest_configure(config: pytest.Config) -> None: - global DEVICE - DEVICE = config.getoption("--device") - - -def random_events(count: int = 40, *, seed: int = 0, classes: int = 3) -> PackedEvents: - generator = torch.Generator().manual_seed(seed) - objects = [ - torch.rand((int(n), 7), generator=generator) - for n in torch.randint(1, 6, (count,), generator=generator) - ] - return pack_events(objects, torch.arange(count) % classes) - - -def fixture_config( - name: str, files: list[Path], labels: list[int], tmp_path: Path, *overrides: str -) -> Config: - """``configs/.yaml`` on local fixture files, on ``DEVICE``, in tmp_path.""" - entries = [ - {"path": str(path), "label": label} for path, label in zip(files, labels) - ] - return load_config( - ROOT / "configs" / f"{name}.yaml", - [ - f"data.files={entries}", - "data.hf_repo=null", - f"data.cache={tmp_path}/events.pt", - "data.prepare_workers=0", - "data.loader_workers=0", - "data.pin_memory=false", - "data.batch_size=64", - f"train.device={DEVICE}", - f"run_dir={tmp_path}/run", - *overrides, - ], - ) - - -@pytest.fixture -def gamgam(tmp_path: Path) -> Config: - """The GamGam config on its prepared fixtures.""" - config = fixture_config("gamgam", GAMGAM_FIXTURES, list(range(5)), tmp_path) - prepare(config.data) - return config diff --git a/tests/fixtures/atlas_opendata_gamgam/mc_341081.ttH125_gamgam_64.root b/tests/fixtures/atlas_opendata_gamgam/mc_341081.ttH125_gamgam_64.root deleted file mode 100644 index 0fe7847b2621fe321603e1504d75a97d32db7c12..0000000000000000000000000000000000000000 Binary files a/tests/fixtures/atlas_opendata_gamgam/mc_341081.ttH125_gamgam_64.root and /dev/null differ diff --git a/tests/fixtures/atlas_opendata_gamgam/mc_343981.ggH125_gamgam_64.root b/tests/fixtures/atlas_opendata_gamgam/mc_343981.ggH125_gamgam_64.root deleted file mode 100644 index a508cd81dc249e0d0f64b43674669baf87f9f0ac..0000000000000000000000000000000000000000 Binary files a/tests/fixtures/atlas_opendata_gamgam/mc_343981.ggH125_gamgam_64.root and /dev/null differ diff --git a/tests/fixtures/atlas_opendata_gamgam/mc_345041.VBFH125_gamgam_64.root b/tests/fixtures/atlas_opendata_gamgam/mc_345041.VBFH125_gamgam_64.root deleted file mode 100644 index 41a1a32eb556847204ae671f6270268ee309926a..0000000000000000000000000000000000000000 Binary files a/tests/fixtures/atlas_opendata_gamgam/mc_345041.VBFH125_gamgam_64.root and /dev/null differ diff --git a/tests/fixtures/atlas_opendata_gamgam/mc_345318.WpH125J_Wincl_gamgam_64.root b/tests/fixtures/atlas_opendata_gamgam/mc_345318.WpH125J_Wincl_gamgam_64.root deleted file mode 100644 index 0ee51ae3833ff2966ee9c24e88573d26d6c51be5..0000000000000000000000000000000000000000 Binary files a/tests/fixtures/atlas_opendata_gamgam/mc_345318.WpH125J_Wincl_gamgam_64.root and /dev/null differ diff --git a/tests/fixtures/atlas_opendata_gamgam/mc_345319.ZH125J_Zincl_gamgam_64.root b/tests/fixtures/atlas_opendata_gamgam/mc_345319.ZH125J_Zincl_gamgam_64.root deleted file mode 100644 index f2c04a179339b1e0d9f7d4feb055908129e271e3..0000000000000000000000000000000000000000 Binary files a/tests/fixtures/atlas_opendata_gamgam/mc_345319.ZH125J_Zincl_gamgam_64.root and /dev/null differ diff --git a/tests/fixtures/ttH_NLO_64.root b/tests/fixtures/ttH_NLO_64.root deleted file mode 100644 index e7a78f9ceb2f93b3812a126e4252bf03c7eff703..0000000000000000000000000000000000000000 Binary files a/tests/fixtures/ttH_NLO_64.root and /dev/null differ diff --git a/tests/test_config.py b/tests/test_config.py deleted file mode 100644 index 0df68eef9d8efe2edabe7c1690db6d516033987f..0000000000000000000000000000000000000000 --- a/tests/test_config.py +++ /dev/null @@ -1,45 +0,0 @@ -from pathlib import Path - -import pytest - -from gnn4colliders.config import load_config, save_config - -CONFIGS = sorted((Path(__file__).resolve().parents[1] / "configs").glob("*.yaml")) - - -@pytest.mark.parametrize("path", CONFIGS, ids=lambda path: path.stem) -def test_every_shipped_config_loads_and_round_trips(path, tmp_path): - config = load_config(path) - assert config.run_dir == f"outputs/{config.name}" - saved = save_config(config, tmp_path / "config.yaml") - assert load_config(saved) == config - - -def test_overrides_parse_yaml_values_and_fast_expands(): - path = CONFIGS[0].parent / "gamgam.yaml" - config = load_config( - path, - [ - "train.fast=true", - "train.learning_rate=1e-3", - "data.files=[{path: a, label: 1}]", - ], - ) - train = config.train - assert train.compile and train.cuda_graphs and train.allow_tf32 - assert train.precision == "bfloat16" and train.learning_rate == 1e-3 - assert config.data.files == [{"path": "a", "label": 1}] - - -@pytest.mark.parametrize( - "override, message", - [ - ("train.learnig_rate=1", "unknown config keys in train"), - ("data.validation_folds=[0]", "disjoint"), - ("task.name=binary", "out_size"), - ("split", "key=value"), - ], -) -def test_invalid_configs_are_rejected(override, message): - with pytest.raises(ValueError, match=message): - load_config(CONFIGS[0].parent / "gamgam.yaml", [override]) diff --git a/tests/test_contrastive.py b/tests/test_contrastive.py deleted file mode 100644 index 5d4ee4b99a58ab3c951d02ed83b77fa06f7a4ff2..0000000000000000000000000000000000000000 --- a/tests/test_contrastive.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Contrastive pretraining of both backbones: the random views, the NT-Xent -loss, and pretraining then fine-tuning and probing.""" - -import json -import math - -import pytest -import torch -from conftest import GAMGAM_FIXTURES, TTH_FIXTURE, fixture_config, random_events - -from gnn4colliders.cli import main -from gnn4colliders.config import ModelConfig, save_config -from gnn4colliders.contrastive import random_view -from gnn4colliders.models import build_model -from gnn4colliders.sequences import ETA, PHI, PT, collate_sequence_batch - -BACKBONES = { - "transformer": {"hid_size": 16, "n_layers": 1, "n_heads": 2, "pair_bias": True}, - "edge_network": {"hid_size": 16, "n_layers": 2, "n_proc_steps": 2, "n_heads": 2}, -} -backbones = pytest.mark.parametrize("backbone", BACKBONES) - - -def _batch(): - events = random_events() - events.object_features[:, 6] = torch.randint( - 0, 5, (events.object_features.shape[0],) - ).float() - return collate_sequence_batch(events, torch.arange(16)) - - -def test_views_rotate_flip_scale_and_drop_objects_but_keep_the_first(): - batch = _batch() - tokens, mask = batch.tokens, batch.token_mask - generator = torch.Generator().manual_seed(0) - view, kept = random_view( - tokens, mask, scale_jitter=0.2, object_dropout=0.5, generator=generator - ) - assert kept[:, 0].all() and not (kept & ~mask).any() and (mask & ~kept).any() - # One factor per event on pt, one sign on eta, azimuths in [-pi, pi). - ratio = view[..., PT][mask] / tokens[..., PT][mask] - assert ((ratio > 0.8 - 1e-6) & (ratio < 1.2 + 1e-6)).all() - sign = view[..., ETA][mask] / tokens[..., ETA][mask] - torch.testing.assert_close(sign.abs(), torch.ones_like(sign)) - assert (view[..., PHI].abs() <= math.pi).all() - torch.testing.assert_close(view[~mask], tokens[~mask]) - - -@backbones -def test_loss_is_finite_and_trains_every_encoder_parameter(backbone): - torch.manual_seed(0) - config = ModelConfig(name="contrastive", backbone=backbone, **BACKBONES[backbone]) - model = build_model(config, num_features=7, num_node_types=5).train() - outputs = model(_batch()) - assert outputs.shape == (16, 3) and torch.isfinite(outputs).all() - # A loss above log(1) and at most about log(2 * 16 - 1) at initialization. - assert 0 < outputs[:, 0].mean() < math.log(31) + 1 - outputs[:, 0].mean().backward() - for name, parameter in model.context_encoder.named_parameters(): - assert parameter.grad is not None and parameter.grad.abs().sum() > 0, name - - -@pytest.mark.parametrize( - ("backbone", "pretrain_config", "finetune_config", "overrides"), - [ - ("transformer", "delphes_contrastive", "gamgam_jepa_finetune", []), - ( - "edge_network", - "delphes_contrastive", - "gamgam_finetune", - ["model.backbone=edge_network"], - ), - ("transformer", "delphes_masked", "gamgam_jepa_finetune", []), - ], -) -def test_pretrain_then_fine_tune_and_probe_on_gamgam( - backbone, pretrain_config, finetune_config, overrides, tmp_path -): - small = [f"model.{key}={value}" for key, value in BACKBONES[backbone].items()] - pretrain = fixture_config( - pretrain_config, - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "pretrain", - "train.max_epochs=2", - *small, - *overrides, - ) - path = str(save_config(pretrain, tmp_path / "pretrain.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - history = [ - json.loads(line) - for line in (pretrain.path / "history.jsonl").read_text().splitlines() - ] - spread = ( - "embedding_spread" if pretrain.model.name == "contrastive" else "target_spread" - ) - assert len(history) == 2 and history[1]["validation"][spread] > 0 - # The pretraining config probes its own encoder, on its labels. - probe = ["data.train_events=40", "model.probe_neighbours=5", "split=validation"] - probe.append(f"model.pretrained={pretrain.path}/last.pt") - assert main(["probe", path, *probe]) == 0 - summary = json.loads((pretrain.path / "validation_probe.json").read_text()) - assert math.isfinite(summary["knn"]["accuracy"]) - - finetune = fixture_config( - finetune_config, - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "finetune", - "train.max_epochs=1", - f"model.pretrained={pretrain.path}/last.pt", - "data.train_events=100", - "model.probe_neighbours=5", - *small, - ) - path = str(save_config(finetune, tmp_path / "finetune.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - assert main(["probe", path]) == 0 - summary = json.loads((finetune.path / "test_probe.json").read_text()) - assert math.isfinite(summary["linear"]["auc"]) diff --git a/tests/test_data.py b/tests/test_data.py deleted file mode 100644 index 6493cfe590212a9eb68de7e5b269847553e0d220..0000000000000000000000000000000000000000 --- a/tests/test_data.py +++ /dev/null @@ -1,218 +0,0 @@ -import hashlib - -import awkward as ak -import numpy as np -import pytest -import torch -import uproot -from conftest import GAMGAM_FIXTURES, TTH_FIXTURE, fixture_config, random_events - -from gnn4colliders import data -from gnn4colliders.data import ( - apply_derived_branches, - build_object_features, - load_events, - parse_derived_branches, - prepare, - save_events, -) -from gnn4colliders.graphs import collate_graph_batch -from gnn4colliders.loader import EventLoader - -BTAG_THRESHOLD = 0.8244273 - - -def test_object_features_are_grouped_by_type_scaled_and_typed(): - arrays = { - "jet_pt": ak.Array([[30.0, 20.0], []]), - "jet_eta": ak.Array([[0.5, -1.0], []]), - "met": np.array([10.0, 40.0]), - } - features, counts = build_object_features( - arrays, - 2, - [["jet_pt", "met"], ["jet_eta", 0], "CALC_E", "NODE_TYPE"], - ["vector", "single"], - [0.1, 1, 0.1, 1], - ) - assert counts.tolist() == [3, 1] - torch.testing.assert_close(features[:, 0], torch.tensor([3.0, 2.0, 1.0, 4.0])) - energy = torch.tensor([30 * np.cosh(0.5), 20 * np.cosh(-1.0), 10, 40]) * 0.1 - torch.testing.assert_close(features[:, 2], energy.float()) - assert features[:, 3].tolist() == [0, 0, 1, 1] - - -def test_derived_branches_split_collections_and_threshold_scores(): - arrays = ak.Array( - { - "lep_pt": [[1.0, 2.0, 3.0]], - "lep_type": [[11, 13, 11]], - "jet_score": [[0.9, 0.1]], - } - ) - derived = parse_derived_branches( - { - "ele_pt": {"source": "lep_pt", "where": {"lep_type": 11}}, - "jet_btag": {"source": "jet_score", "above": 0.5}, - } - ) - result = apply_derived_branches(arrays, derived) - assert result["ele_pt"].tolist() == [[1.0, 3.0]] - assert result["jet_btag"].tolist() == [[1.0, 0.0]] - with pytest.raises(ValueError, match="depends on a derived"): - parse_derived_branches( - {"a": {"source": "x"}, "b": {"source": "a"}}, - ) - - -def test_gamgam_preparation_follows_the_open_data_schema(gamgam): - events = load_events(gamgam.data.cache) - assert len(events) == 5 * 64 - assert events.labels.tolist() == [label for label in range(5) for _ in range(64)] - assert torch.equal(events.weight, torch.ones(len(events))) - offset = 0 - for path in GAMGAM_FIXTURES: - raw = uproot.open(path)["mini"].arrays(library="ak") - folds = events.fold[offset : offset + 64].tolist() - assert folds == (np.asarray(raw.eventNumber, dtype=np.int64) % 10).tolist() - for local in (0, 17, 63): - event = raw[local] - first = int(events.object_offsets[offset + local]) - nodes = events.object_features[ - first : int(events.object_offsets[offset + local + 1]) - ] - electron = np.asarray(event.lep_type) == 11 - muon = np.asarray(event.lep_type) == 13 - pt = np.concatenate( - [ - np.asarray(event.jet_pt), - np.asarray(event.lep_pt)[electron], - np.asarray(event.lep_pt)[muon], - np.asarray(event.photon_pt), - [event.met_et], - ] - ) - sizes = [len(event.jet_pt), electron.sum(), muon.sum()] - sizes += [len(event.photon_pt), 1] - # MeV -> units of 10 GeV, as in pretraining. - np.testing.assert_allclose(nodes[:, 0], pt * 1e-4, rtol=1e-6) - assert nodes[:, 6].tolist() == np.repeat(np.arange(5), sizes).tolist() - jets = len(event.jet_pt) - btag = np.asarray(event.jet_MV2c10) > BTAG_THRESHOLD - assert nodes[:jets, 4].tolist() == btag.astype(float).tolist() - charge = np.concatenate( - [ - np.asarray(event.lep_charge)[electron], - np.asarray(event.lep_charge)[muon], - ] - ) - leptons = nodes[jets : jets + len(charge), 5] - assert leptons.tolist() == charge.astype(float).tolist() - assert nodes[-1, 1] == 0.0 # MET has no pseudorapidity - offset += 64 - - -def test_parallel_and_chunked_preparation_match_serial(gamgam, tmp_path, monkeypatch): - serial = load_events(gamgam.data.cache) - monkeypatch.setattr(data, "CHUNK_EVENTS", 24) - gamgam.data.cache = str(tmp_path / "parallel.pt") - parallel = load_events(prepare(gamgam.data, workers=2)) - for name in ("object_features", "object_offsets", "labels", "fold", "weight"): - assert torch.equal(getattr(parallel, name), getattr(serial, name)) - assert torch.equal(parallel.entry_index, serial.entry_index) - - -def test_delphes_preparation_reads_signed_weights_and_number_folds(tmp_path): - config = fixture_config("tth_cp_even_odd", [TTH_FIXTURE], [0], tmp_path) - events = load_events(prepare(config.data)) - raw = uproot.open(TTH_FIXTURE)["output"].arrays(["Number", "weight"], library="np") - assert events.fold.tolist() == (raw["Number"].astype(np.int64) % 10).tolist() - np.testing.assert_array_equal(events.weight.numpy(), raw["weight"]) - assert (events.weight < 0).any() - - -def test_huggingface_files_are_downloaded_at_the_revision_and_verified( - tmp_path, monkeypatch -): - root_file = tmp_path / "sample.root" - root_file.write_bytes(b"fixture") - requests = [] - - def download(**kwargs): - requests.append(kwargs) - return str(root_file) - - monkeypatch.setattr("huggingface_hub.hf_hub_download", download) - config = fixture_config("gamgam", GAMGAM_FIXTURES[:1], [0], tmp_path).data - config.hf_repo, config.hf_revision = "org/data", "abc123" - good = hashlib.sha256(b"fixture").hexdigest() - config.files = [{"path": "sample.root", "label": 4, "sha256": good}] - assert data.resolve_files(config) == [(root_file, 4)] - assert requests[0]["revision"] == "abc123" - config.files[0]["sha256"] = hashlib.sha256(b"wrong").hexdigest() - with pytest.raises(ValueError, match="SHA-256 mismatch"): - data.resolve_files(config) - - -def test_cache_round_trip_and_version_check(tmp_path): - events = random_events() - loaded = load_events(save_events(events, tmp_path / "e.pt")) - assert torch.equal(loaded.object_features, events.object_features) - payload = torch.load(tmp_path / "e.pt") - payload["cache_schema_version"] = 1 - torch.save(payload, tmp_path / "old.pt") - with pytest.raises(ValueError, match="cache version"): - load_events(tmp_path / "old.pt") - - -def test_shuffled_order_depends_only_on_seed_and_epoch(): - events = random_events() - loader = EventLoader(events, collate_graph_batch, 8, shuffle=True, seed=3) - first = loader.order() - assert torch.equal(first, loader.order()) - loader.set_epoch(1) - assert not torch.equal(first, loader.order()) - assert sorted(loader.order().tolist()) == list(range(len(events))) - - -def test_ranks_split_training_with_padding_and_evaluation_without(): - events = random_events(41) - - def orders(shuffle): - return [ - EventLoader( - events, - collate_graph_batch, - 8, - shuffle=shuffle, - rank=rank, - world_size=2, - ).order() - for rank in range(2) - ] - - train = orders(True) - assert len(train[0]) == len(train[1]) == 21 - assert set(torch.cat(train).tolist()) == set(range(41)) - evaluation = orders(False) - assert sorted(torch.cat(evaluation).tolist()) == list(range(41)) - - -def test_loader_workers_yield_identical_pinned_batches(): - events = random_events(seed=4) - serial = EventLoader(events, collate_graph_batch, 7, shuffle=True, seed=3) - workers = EventLoader( - events, - collate_graph_batch, - 7, - shuffle=True, - seed=3, - num_workers=2, - pin_memory=torch.cuda.is_available(), - ) - pairs = list(zip(serial, workers, strict=True)) - assert len(pairs) == len(serial) == 6 - for left, right in pairs: - assert torch.equal(left.graph.edge_features, right.graph.edge_features) - assert torch.equal(left.labels, right.labels) - assert torch.equal(left.event_index, right.event_index) diff --git a/tests/test_graphs.py b/tests/test_graphs.py deleted file mode 100644 index 6d58ddbdf2c5e5d04071e665812f5c5fc668198b..0000000000000000000000000000000000000000 --- a/tests/test_graphs.py +++ /dev/null @@ -1,67 +0,0 @@ -import math - -import pytest -import torch -from conftest import random_events - -from gnn4colliders.data import pack_events -from gnn4colliders.graphs import ( - PADDING_GRAPHS, - StaticGraphBatch, - collate_graph_batch, - edge_features, - fully_connected_edges, - graph_batch_capacity, -) -from gnn4colliders.models import EdgeNetwork - - -def test_fully_connected_edges_are_source_major_without_self_loops(): - src, dst, counts = fully_connected_edges(torch.tensor([3, 1, 2])) - assert counts.tolist() == [6, 1, 2] - pairs = list(zip(src.tolist(), dst.tolist())) - assert pairs[:6] == [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)] - # A one-node graph keeps its single self-loop. - assert pairs[6:] == [(3, 3), (4, 5), (5, 4)] - - -def test_edge_features_wrap_delta_phi(): - nodes = torch.zeros((2, 7)) - nodes[:, 1] = torch.tensor([0.5, -0.5]) - nodes[:, 2] = torch.tensor([3.0, -3.0]) - features = edge_features(nodes, torch.tensor([0]), torch.tensor([1])) - deta, dphi, dr = features[0].tolist() - assert deta == pytest.approx(1.0) - assert dphi == pytest.approx(6.0 - 2 * math.pi) - assert dr == pytest.approx(math.hypot(deta, dphi)) - - -def test_padded_batch_reproduces_real_graph_outputs(): - events = random_events(12) - batch = collate_graph_batch(events, torch.arange(len(events))) - torch.manual_seed(0) - model = EdgeNetwork(7, 3, 0, 16, 3, 2, 2).eval() - static = StaticGraphBatch( - batch, graph_batch_capacity(batch, node_multiple=32, edge_multiple=64) - ) - # Garbage left in padding rows by an earlier, larger batch must not leak. - static.graph.node_features.normal_() - static.graph.edge_features.normal_() - static.load(batch) - with torch.no_grad(): - expected, padded = model(batch), model(static.batch) - assert padded.shape[0] == len(events) + PADDING_GRAPHS - torch.testing.assert_close(padded[: len(events)], expected, rtol=1e-5, atol=1e-6) - - -def test_static_batch_rejects_batches_beyond_capacity(): - def batch(nodes: int): - events = pack_events([torch.rand((nodes, 7)) for _ in range(4)], [0] * 4) - return collate_graph_batch(events, torch.arange(4)) - - small = batch(1) - static = StaticGraphBatch( - small, graph_batch_capacity(small, node_multiple=1, edge_multiple=1) - ) - with pytest.raises(ValueError, match="capacity"): - static.load(batch(5)) diff --git a/tests/test_jepa.py b/tests/test_jepa.py deleted file mode 100644 index e034b2557c4fa4b27711e223343b32e16616b964..0000000000000000000000000000000000000000 --- a/tests/test_jepa.py +++ /dev/null @@ -1,214 +0,0 @@ -"""JEPA pretraining of both backbones: masking, the predictor, the EMA target, -kinematic targets (masked objects) and fine-tuning.""" - -import json -import math - -import pytest -import torch -from conftest import GAMGAM_FIXTURES, TTH_FIXTURE, fixture_config, random_events - -from gnn4colliders.cli import main -from gnn4colliders.config import ModelConfig, save_config -from gnn4colliders.jepa import ( - FEATURE_TARGETS, - _variance_penalty, - object_features, - sample_targets, - type_index, -) -from gnn4colliders.models import build_model -from gnn4colliders.sequences import collate_sequence_batch -from gnn4colliders.train import load_checkpoint - -SMALL = {"hid_size": 16, "n_layers": 1, "n_heads": 2, "pair_bias": True} -#: Encoder sizes of each backbone. -BACKBONES = { - "transformer": SMALL, - "edge_network": {"hid_size": 16, "n_layers": 2, "n_proc_steps": 2, "n_heads": 2}, -} -backbones = pytest.mark.parametrize("backbone", BACKBONES) - - -def _jepa(backbone="transformer", target="encoding"): - torch.manual_seed(0) - config = ModelConfig( - name="jepa", backbone=backbone, jepa_target=target, **BACKBONES[backbone] - ) - return build_model(config, num_features=7, num_node_types=5) - - -def _batch(seed=0): - events = random_events(seed=seed) - events.object_features[:, 6] = torch.randint( - 0, 5, (events.object_features.shape[0],) - ).float() - return collate_sequence_batch(events, torch.arange(16)) - - -def test_sample_targets_hides_real_objects_and_keeps_one_visible(): - mask = torch.tensor([[1, 1, 1, 1, 1], [1, 1, 0, 0, 0], [1, 0, 0, 0, 0]]).bool() - hidden = sample_targets(mask, 0.4, torch.Generator().manual_seed(0)) - assert hidden.sum(dim=1).tolist() == [2, 1, 0] - assert not (hidden & ~mask).any() - - -def test_type_index_counts_objects_of_each_type_in_order(): - node_type = torch.tensor([[0, 0, 3, 0, 3, 4, 0]]) - mask = torch.tensor([[1, 1, 1, 1, 1, 1, 0]]).bool() - index = type_index(node_type, mask, 5) - assert index[mask].tolist() == [0, 1, 0, 2, 1, 0] - - -@backbones -def test_predictions_do_not_see_the_hidden_objects(backbone): - model, batch = _jepa(backbone).eval(), _batch() - hidden = sample_targets(batch.token_mask, 0.4, torch.Generator().manual_seed(1)) - changed = batch.tokens.clone() - changed[..., :6] = torch.where( - hidden.unsqueeze(-1), torch.rand_like(changed[..., :6]), changed[..., :6] - ) - with torch.no_grad(): - expected = model.predict(batch.tokens, batch.token_mask, hidden)[0] - actual = model.predict(changed, batch.token_mask, hidden)[0] - torch.testing.assert_close(actual[hidden], expected[hidden]) - - -@backbones -def test_loss_is_finite_and_trains_every_context_parameter(backbone): - model, batch = _jepa(backbone).train(), _batch() - outputs = model(batch) - assert outputs.shape == (16, 3) and torch.isfinite(outputs).all() - outputs[:, 0].mean().backward() - assert all(p.grad is None for p in model.target_encoder.parameters()) - # The edge network's global state reaches the loss through the predictor. - for name, parameter in model.context_encoder.named_parameters(): - assert parameter.grad is not None and parameter.grad.abs().sum() > 0, name - - -def test_kinematic_targets_of_objects(): - tokens = torch.tensor([[[1.0, 0.5, math.pi / 2, 2.0, 1.0, -1.0, 0.0]]]) - torch.testing.assert_close( - object_features(tokens), - torch.tensor([[[math.log(2), 0.5, 0.0, 1.0, math.log(3), 1.0, -1.0]]]), - ) - - -@backbones -def test_masked_objects_predict_kinematics_without_a_target_encoder(backbone): - model, batch = _jepa(backbone, "features").train(), _batch() - assert model.target_encoder is None - model.update_target() # nothing to update - outputs = model(batch) - assert outputs.shape == (16, 3) and torch.isfinite(outputs).all() - assert model.predictor_out.out_features == FEATURE_TARGETS - outputs[:, 0].mean().backward() - for name, parameter in model.context_encoder.named_parameters(): - assert parameter.grad is not None and parameter.grad.abs().sum() > 0, name - - -def test_variance_term_penalizes_only_collapsed_dimensions(): - spread = torch.randn(4096, 8, generator=torch.Generator().manual_seed(0)) - assert _variance_penalty(spread) < 0.03 - spread[:, :4] = 0.5 # half the dimensions collapsed - torch.testing.assert_close( - _variance_penalty(spread), torch.tensor(0.5), atol=0.02, rtol=0 - ) - - -def test_target_is_an_exponential_moving_average(): - model = _jepa() - model.total_steps = 2 - with torch.no_grad(): - for parameter in model.context_encoder.parameters(): - parameter.add_(1.0) - before = [p.clone() for p in model.target_encoder.parameters()] - context = list(model.context_encoder.parameters()) - model.update_target() # momentum 0.996 - for old, new, target in zip(before, model.target_encoder.parameters(), context): - torch.testing.assert_close(new, 0.996 * old + 0.004 * target) - model.update_target() # momentum 0.998 - model.update_target() # momentum 1: unchanged - frozen = [p.clone() for p in model.target_encoder.parameters()] - model.update_target() - assert all( - torch.equal(a, b) for a, b in zip(frozen, model.target_encoder.parameters()) - ) - assert int(model.step) == 4 - - -@backbones -def test_classifier_fine_tunes_from_the_context_encoder(backbone, tmp_path): - model = _jepa(backbone) - path = tmp_path / "jepa.pt" - torch.save({"model_state_dict": model.state_dict()}, path) - config = ModelConfig( - name=backbone, - out_size=3, - pretrained=str(path), - freeze_backbone=True, - **BACKBONES[backbone], - ) - classifier = build_model(config, num_features=7, num_node_types=5) - encoder = getattr(classifier, "encoder", None) or classifier.backbone - expected = model.context_encoder.state_dict() - assert encoder.state_dict().keys() == expected.keys() - for name, value in encoder.state_dict().items(): - assert torch.equal(value, expected[name]) - trainable = [n for n, p in classifier.named_parameters() if p.requires_grad] - assert trainable == ["classifier.weight", "classifier.bias"] - - -def test_transformer_needs_a_transformer_encoder(tmp_path): - path = tmp_path / "jepa.pt" - torch.save({"model_state_dict": _jepa("edge_network").state_dict()}, path) - config = ModelConfig(name="transformer", pretrained=str(path), **SMALL) - with pytest.raises(ValueError, match="no JEPA transformer encoder"): - build_model(config, num_features=7, num_node_types=5) - - -@pytest.mark.parametrize( - ("backbone", "pretrain_config", "finetune_config"), - [ - ("transformer", "delphes_jepa", "gamgam_jepa_finetune"), - ("edge_network", "delphes_jepa_edge", "gamgam_finetune"), - ], -) -def test_pretrain_on_delphes_then_fine_tune_on_gamgam( - backbone, pretrain_config, finetune_config, tmp_path -): - small = [f"model.{key}={value}" for key, value in BACKBONES[backbone].items()] - pretrain = fixture_config( - pretrain_config, - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "jepa", - "train.max_epochs=2", - *small, - ) - path = str(save_config(pretrain, tmp_path / "jepa.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - history = [ - json.loads(line) - for line in (pretrain.path / "history.jsonl").read_text().splitlines() - ] - assert len(history) == 2 - assert history[1]["validation"]["target_spread"] > 0 - assert ( - int(load_checkpoint(pretrain.path / "last.pt")["model_state_dict"]["step"]) > 0 - ) - - finetune = fixture_config( - finetune_config, - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "finetune", - "train.max_epochs=1", - f"model.pretrained={pretrain.path}/last.pt", - *small, - ) - path = str(save_config(finetune, tmp_path / "finetune.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - assert main(["evaluate", path]) == 0 diff --git a/tests/test_models.py b/tests/test_models.py deleted file mode 100644 index d232a680003034b3cf108a327b1e3be570373936..0000000000000000000000000000000000000000 --- a/tests/test_models.py +++ /dev/null @@ -1,130 +0,0 @@ -import dataclasses - -import pytest -import torch -from conftest import PRETRAINED, random_events - -from gnn4colliders.config import ModelConfig -from gnn4colliders.graphs import collate_graph_batch, edge_features -from gnn4colliders.models import ( - EdgeNetwork, - FineTunedEdgeNetwork, - build_model, - pair_features, - state_dict_of, -) -from gnn4colliders.sequences import collate_sequence_batch - -CHECKPOINTS = [ - PRETRAINED / "multiclass_pretrained_model_12/model_epoch_71.pt", - PRETRAINED / "multilabel_pretrained_model_41/model_epoch_44.pt", -] - - -@pytest.mark.parametrize("checkpoint", CHECKPOINTS, ids=lambda path: path.parent.name) -def test_pretrained_checkpoints_fine_tune(checkpoint): - if checkpoint.read_bytes()[:7] == b"version": - pytest.fail(f"{checkpoint} is an unfetched Git LFS pointer") - events = random_events(seed=1) - batch = collate_graph_batch(events, torch.arange(len(events))) - model = build_model( - ModelConfig(out_size=5, pretrained=str(checkpoint), freeze_backbone=True), - num_features=7, - num_node_types=5, - ) - assert isinstance(model, FineTunedEdgeNetwork) - pretrained = state_dict_of( - torch.load(checkpoint, map_location="cpu", weights_only=False) - ) - for name, value in model.backbone.state_dict().items(): - assert torch.equal(value, pretrained[name]) - assert [name for name, p in model.named_parameters() if p.requires_grad] == [ - "classifier.weight", - "classifier.bias", - ] - logits = model(batch) - assert logits.shape == (len(events), 5) and torch.isfinite(logits).all() - - -def test_state_dict_keys_of_wrapped_and_original_models_are_mapped(): - state = {"module.classify.weight": 1, "_orig_mod.node_encoder.0.bias": 2} - assert state_dict_of({"model_state_dict": state}) == { - "classifier.weight": 1, - "node_encoder.0.bias": 2, - } - - -def test_edge_network_uses_global_features_when_configured(): - events = random_events(seed=2) - events = type(events)(**{**events.__dict__, "global_features": torch.rand(40, 2)}) - batch = collate_graph_batch(events, torch.arange(8)) - model = EdgeNetwork(7, 3, 2, 8, 3, 2, 1) - assert model(batch).shape == (8, 3) - - -SEQUENCE_CONFIGS = { - "transformer": ModelConfig(name="transformer", hid_size=16, out_size=3, n_layers=1), - "pair_bias": ModelConfig( - name="transformer", hid_size=16, out_size=3, n_layers=1, pair_bias=True - ), - "deep_sets": ModelConfig(name="deep_sets", hid_size=16, out_size=3, n_layers=2), -} - - -def _sequence_model(name): - torch.manual_seed(0) - return build_model(SEQUENCE_CONFIGS[name], num_features=7, num_node_types=5).eval() - - -def _typed_events(seed): - events = random_events(seed=seed) - events.object_features[:, 6] = torch.randint( - 0, 5, (events.object_features.shape[0],) - ).float() - return events - - -@pytest.mark.parametrize("name", SEQUENCE_CONFIGS) -def test_sequence_models_ignore_padding_tokens(name): - events, model = _typed_events(3), _sequence_model(name) - alone = collate_sequence_batch(events, torch.tensor([0])) - batch = collate_sequence_batch(events, torch.arange(8)) - with torch.no_grad(): - torch.testing.assert_close(model(batch)[:1], model(alone), rtol=1e-5, atol=1e-5) - - -@pytest.mark.parametrize("name", SEQUENCE_CONFIGS) -def test_sequence_models_do_not_depend_on_object_order(name): - events, model = _typed_events(4), _sequence_model(name) - batch = collate_sequence_batch(events, torch.arange(8)) - permutation = torch.randperm(batch.tokens.shape[1]) - shuffled = dataclasses.replace( - batch, - tokens=batch.tokens[:, permutation], - token_mask=batch.token_mask[:, permutation], - ) - with torch.no_grad(): - torch.testing.assert_close(model(shuffled), model(batch), rtol=1e-5, atol=1e-5) - - -def test_pair_features_are_the_edge_features_and_the_pair_mass(): - events = _typed_events(5) - batch = collate_sequence_batch(events, torch.tensor([0])) - tokens = batch.tokens[0] - pairs = pair_features(batch.tokens)[0] - src, dst = torch.meshgrid( - torch.arange(tokens.shape[0]), torch.arange(tokens.shape[0]), indexing="ij" - ) - expected = edge_features(tokens, src.flatten(), dst.flatten()) - torch.testing.assert_close(pairs[..., :3].reshape(-1, 3), expected) - # The mass of two massless four-vectors. - pt, eta, phi = tokens[:, 0], tokens[:, 1], tokens[:, 2] - pt, eta, phi = pt.double(), eta.double(), phi.double() - p = torch.stack( - (pt * eta.cosh(), pt * phi.cos(), pt * phi.sin(), pt * eta.sinh()), dim=1 - ) - total = p[0] + p[1] - m2 = total[0] ** 2 - (total[1:] ** 2).sum() - torch.testing.assert_close( - pairs[0, 1, 3].double(), torch.log1p(m2), rtol=1e-5, atol=0 - ) diff --git a/tests/test_pfn.py b/tests/test_pfn.py deleted file mode 100644 index 773ed948a10211a4d89ef97ef63875dda1692724..0000000000000000000000000000000000000000 --- a/tests/test_pfn.py +++ /dev/null @@ -1,180 +0,0 @@ -"""The prior-fitted network: the task prior, context-only attention, in-context -inference and the Delphes-to-GamGam workflow.""" - -import json -from dataclasses import replace - -import torch -from conftest import GAMGAM_FIXTURES, TTH_FIXTURE, fixture_config, random_events - -from gnn4colliders.cli import main -from gnn4colliders.config import ModelConfig, TaskConfig, save_config -from gnn4colliders.models import build_model -from gnn4colliders.pfn import ABSENT_LOGIT, TaskLoader, in_context_forward -from gnn4colliders.sequences import collate_sequence_batch - -SMALL = {"hid_size": 16, "n_layers": 1, "n_heads": 2, "pair_bias": True} -PRIOR = { - "name": "pfn", - "task_classes": [2, 3], - "context_events": [8, 16], - "query_events": 4, - "steps_per_epoch": 3, - "validation_steps": 2, -} - - -def _events(): - events = random_events(200, classes=4) - events.object_features[:, 6] = torch.randint( - 0, 5, (events.object_features.shape[0],) - ).float() - return events - - -def _loader(events, train=True, **prior): - config = TaskConfig(**{**PRIOR, **prior}) - return TaskLoader(events, torch.arange(len(events)), config, 64, train=train) - - -def _pfn(): - torch.manual_seed(0) - config = ModelConfig(name="pfn", pfn_layers=2, max_classes=4, **SMALL) - return build_model(config, num_features=7, num_node_types=5).eval() - - -def test_tasks_follow_the_prior(): - events = _events() - for batch in _loader(events, prior_classes=[1, 2, 3]): - tasks, context = batch.context_labels.shape - assert 8 <= context <= 16 and tasks == 64 // (context + 4) - assert ((batch.num_classes >= 2) & (batch.num_classes <= 3)).all() - index = batch.events.event_index.view(tasks, -1) - queries = batch.labels.view(tasks, -1) - for task in range(tasks): - k = int(batch.num_classes[task]) - labels = torch.cat((batch.context_labels[task], queries[task])) - # Every class has a context event, and each task label is one - # original class out of prior_classes. - assert set(batch.context_labels[task].tolist()) == set(range(k)) - original = events.labels[index[task]] - for label in range(k): - assert len(set(original[labels == label].tolist())) == 1 - assert not set(original.tolist()) - {1, 2, 3} - - -def test_training_tasks_vary_by_epoch_and_validation_tasks_do_not(): - events = _events() - - def first(loader, epoch): - loader.set_epoch(epoch) - return next(iter(loader)).events.event_index - - train, validation = _loader(events), _loader(events, train=False) - assert torch.equal(first(train, 0), first(train, 0)) - assert not torch.equal(first(train, 0), first(train, 1)) - assert torch.equal(first(validation, 0), first(validation, 1)) - - -def test_queries_see_only_the_context(): - model, batch = _pfn(), next(iter(_loader(_events(), train=False))) - tasks, context = batch.context_labels.shape - tokens = batch.events.tokens.clone() - first_query = context # of the first task - tokens[first_query, :, :6] += 1.0 - changed = replace(batch, events=replace(batch.events, tokens=tokens)) - with torch.no_grad(): - expected = model(batch).view(tasks, -1, 4) - actual = model(changed).view(tasks, -1, 4) - assert not torch.allclose(actual[0, 0], expected[0, 0]) - torch.testing.assert_close(actual[0, 1:], expected[0, 1:]) - torch.testing.assert_close(actual[1:], expected[1:]) - absent = (torch.arange(4) >= batch.num_classes[:, None, None]).expand_as(expected) - assert (expected[absent] == ABSENT_LOGIT).all() - - -def test_in_context_inference_matches_a_task(): - model, events = _pfn(), _events() - context, queries = torch.arange(0, 40), torch.arange(40, 60) - forward = in_context_forward( - model, events, context, num_classes=4, batch_size=16, device=torch.device("cpu") - ) - with torch.no_grad(): - actual = forward(collate_sequence_batch(events, queries)) - everything = collate_sequence_batch(events, torch.cat((context, queries))) - embedded = model.embed(everything).unsqueeze(0) - memory = model.context(embedded[:, :40], events.labels[context][None]) - expected = model.queries(embedded[:, 40:], memory, torch.tensor([4]))[0] - torch.testing.assert_close(actual, expected) - - -def test_encoder_starts_from_a_jepa_checkpoint(tmp_path): - jepa = build_model( - ModelConfig(name="jepa", **SMALL), num_features=7, num_node_types=5 - ) - path = tmp_path / "jepa.pt" - torch.save({"model_state_dict": jepa.state_dict()}, path) - config = ModelConfig(name="pfn", pretrained=str(path), **SMALL) - pfn = build_model(config, num_features=7, num_node_types=5) - expected = jepa.context_encoder.state_dict() - for name, value in pfn.encoder.state_dict().items(): - assert torch.equal(value, expected[name]) - - -def test_meta_train_on_delphes_then_classify_gamgam_in_context(tmp_path): - small = [f"model.{key}={value}" for key, value in SMALL.items()] - prior = [ - "task.context_events=[4,16]", - "task.query_events=4", - "task.steps_per_epoch=2", - "task.validation_steps=1", - ] - meta = fixture_config( - "delphes_pfn", - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "meta", - "train.max_epochs=2", - "model.pfn_layers=1", - *small, - *prior, - ) - path = str(save_config(meta, tmp_path / "meta.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - history = (meta.path / "history.jsonl").read_text().splitlines() - assert "accuracy" in json.loads(history[-1])["validation"] - - gamgam = fixture_config( - "gamgam_pfn", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "gamgam", - f"checkpoint={meta.path}/best.pt", - "data.train_events=20", - "model.pfn_layers=1", - "split=validation", - *small, - ) - path = str(save_config(gamgam, tmp_path / "gamgam.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 2 # nothing to train on GamGam - assert main(["evaluate", path]) == 0 - metrics = json.loads((gamgam.path / "validation_metrics.json").read_text()) - assert 0 <= metrics["metrics"]["auc"] <= 1 - assert main(["predict", path]) == 0 - assert main(["export", path]) == 2 - - -def test_in_context_models_need_a_context(tmp_path): - config = fixture_config( - "gamgam_tabpfn", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path, - "data.train_events=null", - ) - path = str(save_config(config, tmp_path / "tabpfn.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 2 - assert main(["evaluate", path]) == 2 diff --git a/tests/test_probe.py b/tests/test_probe.py deleted file mode 100644 index b628fc125eec4fa1de6f5eddaafa7f5a4d5d7aab..0000000000000000000000000000000000000000 --- a/tests/test_probe.py +++ /dev/null @@ -1,259 +0,0 @@ -"""k-NN and linear probes of frozen encoders: the classifiers, the encoder -weights of each kind of checkpoint, and the probe command.""" - -import json -import math -import runpy -import sys - -import pytest -import torch -from conftest import GAMGAM_FIXTURES, ROOT, TTH_FIXTURE, fixture_config - -from gnn4colliders.cli import main -from gnn4colliders.config import ModelConfig, save_config -from gnn4colliders.data import load_events -from gnn4colliders.models import encoder_state, pretrained_encoder -from gnn4colliders.probe import ( - knn_logits, - layer_embeddings, - linear_cka, - linear_logits, -) - -SMALL = {"hid_size": 16, "n_layers": 1, "n_heads": 2, "pair_bias": True} -BACKBONES = { - "transformer": SMALL, - "edge_network": {"hid_size": 16, "n_layers": 2, "n_proc_steps": 2, "n_heads": 2}, -} - - -def _clusters(count=60, seed=0): - """Two well separated classes of 8-dimensional embeddings.""" - generator = torch.Generator().manual_seed(seed) - labels = torch.arange(count) % 2 - centres = torch.tensor([[3.0] + [0.0] * 7, [0.0, 3.0] + [0.0] * 6]) - return centres[labels] + 0.3 * torch.randn(count, 8, generator=generator), labels - - -def test_knn_and_linear_probes_separate_clusters(): - train, labels = _clusters(seed=0) - queries, expected = _clusters(seed=1) - for logits in ( - knn_logits(train, labels, queries, num_classes=2, neighbours=5), - linear_logits(train, labels, queries, num_classes=2, l2=1e-4), - ): - assert logits.shape == (60, 2) - assert torch.equal(logits.argmax(dim=1), expected) - - -def test_encoder_weights_of_each_kind_of_checkpoint(tmp_path): - weight = torch.ones(2) - checkpoints = { - "jepa": {"context_encoder.norm.weight": weight, "target_encoder.x": weight}, - "transformer": {"encoder.norm.weight": weight, "classifier.weight": weight}, - "edge_network": {"node_encoder.0.weight": weight, "classifier.weight": weight}, - "fine_tuned": {"backbone.node_encoder.0.weight": weight, "head.weight": weight}, - "tabular": {"inputs.encoder.norm.weight": weight, "network.x.up": weight}, - } - expected = { - "jepa": ["norm.weight"], - "transformer": ["norm.weight"], - "edge_network": ["node_encoder.0.weight"], - "fine_tuned": ["node_encoder.0.weight"], - "tabular": ["norm.weight"], - } - for name, state in checkpoints.items(): - path = tmp_path / f"{name}.pt" - torch.save({"model_state_dict": state}, path) - assert list(encoder_state(str(path))) == expected[name] - - -@pytest.mark.parametrize( - ("backbone", "pretrain_config", "probe_config"), - [ - ("transformer", "delphes_jepa", "gamgam_jepa_finetune"), - ("edge_network", "delphes_jepa_edge", "gamgam_finetune"), - ], -) -def test_probe_random_and_jepa_encoders_on_gamgam( - backbone, pretrain_config, probe_config, tmp_path -): - small = [f"model.{key}={value}" for key, value in BACKBONES[backbone].items()] - pretrain = fixture_config( - pretrain_config, - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "jepa", - "train.max_epochs=1", - *small, - ) - path = str(save_config(pretrain, tmp_path / "jepa.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - - for name, encoder in (("random", "null"), ("jepa", f"{pretrain.path}/last.pt")): - config = fixture_config( - probe_config, - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / name, - f"model.pretrained={encoder}", - "data.train_events=100", - "model.probe_neighbours=5", - "split=validation", - *small, - ) - path = str(save_config(config, tmp_path / f"{name}.yaml")) - assert main(["prepare", path]) == 0 - assert main(["probe", path]) == 0 - summary = json.loads((config.path / "validation_probe.json").read_text()) - assert summary["train_events"] == 100 - for probe in ("knn", "linear"): - assert math.isfinite(summary[probe]["auc"]) - - -def test_probe_a_supervised_transformer(tmp_path): - small = [f"model.{key}={value}" for key, value in SMALL.items()] - trained = fixture_config( - "gamgam_transformer", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "supervised", - "train.max_epochs=1", - *small, - ) - path = str(save_config(trained, tmp_path / "supervised.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - config = fixture_config( - "gamgam_transformer", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "probe", - f"model.pretrained={trained.path}/last.pt", - "data.train_events=50", - *small, - ) - path = str(save_config(config, tmp_path / "probe.yaml")) - assert main(["prepare", path]) == 0 - assert main(["probe", path]) == 0 - summary = json.loads((config.path / "test_probe.json").read_text()) - assert summary["encoder"] == f"{trained.path}/last.pt" - # On the training split, the probes' own events are not queries. - assert main(["probe", path, "split=train"]) == 0 - summary = json.loads((config.path / "train_probe.json").read_text()) - events = load_events(config.data.cache) - train_fold = torch.isin(events.fold, torch.tensor(config.data.train_folds)) - assert summary["num_samples"] == int(train_fold.sum()) - 50 - - -def test_a_random_encoder_is_rebuilt_identically_without_touching_global_state(): - config = ModelConfig(name="transformer", **SMALL) - state = torch.random.get_rng_state() - first, second = ( - pretrained_encoder(config, num_features=7, num_node_types=5) for _ in range(2) - ) - assert torch.equal(torch.random.get_rng_state(), state) - for a, b in zip(first.parameters(), second.parameters()): - assert torch.equal(a, b) - - -def test_linear_cka_is_one_for_rotations_and_scales_and_small_for_noise(): - generator = torch.Generator().manual_seed(0) - x = torch.randn(500, 8, generator=generator) - rotation, _ = torch.linalg.qr(torch.randn(8, 8, generator=generator)) - assert linear_cka(x, x) == pytest.approx(1.0) - assert linear_cka(x, 3.0 * x @ rotation + 1.0) == pytest.approx(1.0) - assert linear_cka(x, torch.randn(500, 8, generator=generator)) < 0.1 - # A representation that drops half the dimensions is in between. - assert 0.3 < linear_cka(x, x[:, :4]) < 0.9 - - -def test_linear_cka_of_a_constant_representation_is_undefined(): - x = torch.randn(100, 8, generator=torch.Generator().manual_seed(0)) - # 0.1 is not exact in binary: centering leaves rounding, not zeros. - constant = torch.full((100, 8), 0.1) - assert math.isnan(linear_cka(x, constant)) - assert math.isnan(linear_cka(constant, constant)) - - -@pytest.mark.parametrize("backbone", BACKBONES) -def test_layer_embeddings_end_with_the_event_embedding(backbone): - from conftest import random_events - - from gnn4colliders.jepa import event_embedding - from gnn4colliders.sequences import collate_sequence_batch - - events = random_events(count=30) - events.object_features[:, 6] = torch.randint( - 0, 5, (events.object_features.shape[0],) - ).float() - config = ModelConfig(name=backbone, **BACKBONES[backbone]) - encoder = pretrained_encoder( - config, num_features=7, num_node_types=5, backbone=backbone - ).eval() - indices = torch.arange(30) - layers = layer_embeddings( - encoder, events, indices, batch_size=8, device=torch.device("cpu") - ) - # Transformer: 1 block and the final norm; edge network: encoder, 2 steps, - # global state. - assert len(layers) == {"transformer": 2, "edge_network": 4}[backbone] - assert all(layer.shape == (30, 16) for layer in layers) - batch = collate_sequence_batch(events, indices) - with torch.no_grad(): - expected = event_embedding(encoder, batch.tokens, batch.token_mask) - torch.testing.assert_close(layers[-1], expected) - - -def test_compare_encoders_script(tmp_path, monkeypatch, capsys): - small = [f"model.{key}={value}" for key, value in SMALL.items()] - pretrain = fixture_config( - "delphes_jepa", - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "jepa", - "train.max_epochs=1", - "model.ffn_size=64", - *small, - ) - path = str(save_config(pretrain, tmp_path / "jepa.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - config = fixture_config( - "gamgam_jepa_finetune", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "gamgam", - "model.ffn_size=64", - *small, - ) - path = str(save_config(config, tmp_path / "gamgam.yaml")) - assert main(["prepare", path]) == 0 - output = tmp_path / "cka.json" - encoders = ["random=random", f"jepa={pretrain.path}/last.pt"] - argv = [path, "--encoders", *encoders, "--events", "30", "--output", str(output)] - monkeypatch.setattr(sys, "argv", ["compare_encoders", *argv]) - runpy.run_path(str(ROOT / "scripts/compare_encoders.py"), run_name="__main__") - summary = json.loads(output.read_text()) - assert summary["events"] == 30 - embeddings = summary["embeddings"] - assert embeddings["random"][0] == pytest.approx(1.0) - assert embeddings["jepa"][1] == pytest.approx(1.0) - # The two encoders differ, and CKA is symmetric. - assert 0 < embeddings["jepa"][0] < 1 - 1e-6 - assert embeddings["jepa"][0] == pytest.approx(embeddings["random"][1]) - layerwise = summary["layerwise"]["jepa"] - assert len(layerwise) == 2 and all(0 < value <= 1 for value in layerwise) - assert layerwise[-1] == pytest.approx(embeddings["jepa"][0]) - assert "Layer by layer against random" in capsys.readouterr().out - - for encoders in (["random=random", "random=random"], ["last.pt"]): - argv = [path, "--encoders", *encoders, "--output", str(output)] - monkeypatch.setattr(sys, "argv", ["compare_encoders", *argv]) - with pytest.raises(SystemExit) as error: - runpy.run_path( - str(ROOT / "scripts/compare_encoders.py"), run_name="__main__" - ) - assert error.value.code == 2 diff --git a/tests/test_tabular.py b/tests/test_tabular.py deleted file mode 100644 index fe17be4ed7eafd96dde020d7abe3e368b0f67f10..0000000000000000000000000000000000000000 --- a/tests/test_tabular.py +++ /dev/null @@ -1,412 +0,0 @@ -"""Fixed-width event tables, balanced label subsets, LoRA, TabICL in -context, and TabPFN in context and meta-trained (when installed, with its -weights).""" - -import json -import math -import os -import sys - -import pytest -import torch -from conftest import GAMGAM_FIXTURES, TTH_FIXTURE, fixture_config -from torch import nn -from torch.nn.utils import parametrize - -from gnn4colliders.cli import main -from gnn4colliders.config import ModelConfig, save_config -from gnn4colliders.data import PackedEvents, pack_events -from gnn4colliders.sequences import collate_sequence_batch -from gnn4colliders.tabular import ( - EventInputs, - RowLayout, - add_lora, - merge_lora, - sequence_table, - table_features, -) -from gnn4colliders.train import load_checkpoint - - -def _objects(types): - rows = torch.arange(len(types) * 7, dtype=torch.float32).view(-1, 7) - rows[:, 6] = torch.tensor(types, dtype=torch.float32) - return rows - - -def test_table_keeps_the_leading_objects_of_each_type(): - first, second = _objects([0, 0, 0, 3, 3, 4]), _objects([1, 4]) - events = pack_events([first, second], [0, 1]) - layout = RowLayout((2, 1, 1, 1, 1)) - table = table_features(events, torch.tensor([0, 1]), layout) - assert table.shape == (2, layout.width) == (2, 6 * 6 + 5) - slots = table[:, :36].view(2, 6, 6) - # Slots: jet, jet, electron, muon, photon, MET. - torch.testing.assert_close(slots[0, :2], first[:2, :6]) - torch.testing.assert_close(slots[0, 4:], first[[3, 5], :6]) - assert slots[0, 2:4].isnan().all() - torch.testing.assert_close(slots[1, 2], second[0, :6]) - torch.testing.assert_close(slots[1, 5], second[1, :6]) - assert slots[1, [0, 1, 3, 4]].isnan().all() - assert table[:, 36:].tolist() == [[3, 0, 0, 2, 1], [0, 1, 0, 0, 1]] - - -def test_sequence_batches_give_the_same_table(): - events = pack_events( - [_objects([0, 0, 0, 3, 3, 4]), _objects([1, 4]), _objects([2, 2, 2, 0])], - [0, 1, 2], - ) - indices = torch.tensor([2, 0, 1]) - layout = RowLayout((2, 1, 1, 1, 1), relative_phi=True, pair_objects=3) - torch.testing.assert_close( - sequence_table(collate_sequence_batch(events, indices), layout), - table_features(events, indices, layout), - equal_nan=True, - ) - - -def _object(pt, eta, phi, node_type): - return [pt, eta, phi, pt * math.cosh(eta), 0.0, 0.0, node_type] - - -def test_relative_phi_and_the_masses_and_delta_r_of_leading_pairs(): - # Two back-to-back photons (mass 2, Delta R pi) and a softer jet; an - # event with one photon has no pair. - pair = torch.tensor( - [_object(1.0, 0.0, 0.5, 3), _object(1.0, 0.0, 0.5 - math.pi, 3)] - + [_object(0.5, 1.0, 2.0, 0)] - ) - single = torch.tensor([_object(1.0, 0.3, 1.0, 3)]) - events = pack_events([pair, single], [0, 1]) - layout = RowLayout((1, 0, 0, 2, 0), relative_phi=True, pair_objects=2) - table = table_features(events, torch.tensor([0, 1]), layout) - assert table.shape == (2, layout.width) == (2, 3 * 6 + 5 + 2) - slots = table[:, :18].view(2, 3, 6) - # Azimuths from the leading photon's. - torch.testing.assert_close(slots[0, :, 2], torch.tensor([1.5, 0.0, -math.pi])) - assert slots[1, 1, 2] == 0 - torch.testing.assert_close(table[0, -2:], torch.tensor([2.0, math.pi])) - assert table[1, -2:].isnan().all() - - -def test_lora_starts_at_the_frozen_weight_and_trains_only_its_update(): - network = nn.Sequential(nn.Linear(6, 4)) - weight, x = network[0].weight.detach().clone(), torch.randn(5, 6) - expected = network(x) - add_lora(network, ["0.weight"], rank=2) - torch.testing.assert_close(network(x), expected) - trainable = [name for name, p in network.named_parameters() if p.requires_grad] - assert sorted(trainable) == [ - "0.parametrizations.weight.0.down", - "0.parametrizations.weight.0.up", - ] - # Code that reads the weight directly sees the update too. - with torch.no_grad(): - network[0].parametrizations.weight[0].up.fill_(1.0) - lora = network[0].parametrizations.weight[0] - torch.testing.assert_close(network[0].weight, weight + lora.up @ lora.down) - # Merged for inference, the network computes the same. - expected = network(x) - merge_lora(network) - assert not parametrize.is_parametrized(network[0]) - torch.testing.assert_close(network(x), expected) - - -def test_random_evaluation_subsets_are_seeded_and_sorted(): - indices = torch.arange(100) * 3 - subset = PackedEvents.random_subset(indices, 20, seed=1) - assert len(subset) == 20 and set(subset.tolist()) <= set(indices.tolist()) - assert torch.equal(subset, subset.sort().values) - assert torch.equal(subset, PackedEvents.random_subset(indices, 20, seed=1)) - assert len(PackedEvents.random_subset(indices, 500, seed=1)) == 100 - - -def test_balanced_subsets_are_equal_per_class_and_nested(): - labels = torch.tensor([0] * 50 + [1] * 30 + [2] * 5) - events = pack_events([torch.zeros(1, 7)] * len(labels), labels) - indices = torch.arange(len(labels)) - small = events.balanced_subset(indices, 12, seed=3) - assert torch.bincount(labels[small]).tolist() == [4, 4, 4] - large = events.balanced_subset(indices, 30, seed=3) - assert torch.bincount(labels[large]).tolist() == [10, 10, 5] # 5 exist - assert set(small.tolist()) <= set(large.tolist()) - assert torch.equal(small, events.balanced_subset(indices, 12, seed=3)) - assert not torch.equal(small, events.balanced_subset(indices, 12, seed=4)) - - -def test_tabpfn_classifies_gamgam_in_context(tmp_path): - pytest.importorskip("tabpfn") - config = fixture_config( - "gamgam_tabpfn", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path, - "data.train_events=100", - # Version 2 weights download without a login. - "model.tabpfn_version=v2", - "model.tabpfn_estimators=1", - "split=validation", - ) - path = str(save_config(config, tmp_path / "tabpfn.yaml")) - assert main(["prepare", path]) == 0 - assert main(["evaluate", path]) == 0 - metrics = json.loads((config.path / "validation_metrics.json").read_text()) - assert not math.isnan(metrics["metrics"]["auc"]) - assert metrics["checkpoint"] == "TabPFN v2 (default checkpoint)" - - -def _tabpfn_3_5_weights() -> None: - pytest.importorskip("tabpfn") - from tabpfn.model_loading import get_cache_dir - - if not (get_cache_dir() / "tabpfn-v3.5-20260909.safetensors").exists(): - pytest.skip("TabPFN-3.5 weights are not in the TabPFN cache") - - -def test_meta_train_tabpfn_with_lora_then_classify_gamgam(tmp_path): - _tabpfn_3_5_weights() - meta = fixture_config( - "delphes_tabpfn", - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "meta", - "train.max_epochs=1", - "task.context_events=[4,16]", - "task.query_events=4", - "task.steps_per_epoch=2", - "task.validation_steps=1", - "model.relative_phi=true", - "model.pair_objects=4", - ) - path = str(save_config(meta, tmp_path / "meta.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - # Of the weights, only the LoRA updates are saved: 24 blocks of 4 - # projections and 2 MLP layers. - state = load_checkpoint(meta.path / "best.pt")["model_state_dict"] - weights = [key for key in state if key.endswith(".weight")] - updates = [key for key in state if key.endswith((".down", ".up"))] - assert not weights and len(updates) == 24 * 6 * 2 - - gamgam = fixture_config( - "gamgam_tabpfn", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "gamgam", - f"checkpoint={meta.path}/best.pt", - "data.train_events=40", - "split=validation", - "model.relative_phi=true", - "model.pair_objects=4", - ) - path = str(save_config(gamgam, tmp_path / "gamgam.yaml")) - assert main(["prepare", path]) == 0 - assert main(["evaluate", path]) == 0 - metrics = json.loads((gamgam.path / "validation_metrics.json").read_text()) - assert metrics["checkpoint"] == f"{meta.path}/best.pt" - assert math.isfinite(metrics["metrics"]["auc"]) - - -def test_tabicl_classifies_gamgam_in_context_with_physics_features(tmp_path): - pytest.importorskip("tabicl") - config = fixture_config( - "gamgam_tabicl", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path, - "data.train_events=100", - "model.tabicl_estimators=1", - "model.relative_phi=true", - "model.pair_objects=4", - "split=validation", - ) - path = str(save_config(config, tmp_path / "tabicl.yaml")) - assert main(["prepare", path]) == 0 - assert main(["evaluate", path]) == 0 - metrics = json.loads((config.path / "validation_metrics.json").read_text()) - assert metrics["checkpoint"] == "TabICL (tabicl-classifier-v2-20260212.ckpt)" - assert math.isfinite(metrics["metrics"]["auc"]) - - -#: faiss and torch crash together on macOS unless OpenMP has one thread. -RETRIEVAL = sys.platform != "darwin" or os.environ.get("OMP_NUM_THREADS") == "1" - - -@pytest.mark.parametrize( - "context_size", - [ - None, - pytest.param( - 50, - marks=pytest.mark.skipif( - not RETRIEVAL, reason="retrieval on macOS needs OMP_NUM_THREADS=1" - ), - ), - ], -) -def test_tabdpt_classifies_gamgam_in_context(context_size, tmp_path): - pytest.importorskip("tabdpt") - config = fixture_config( - "gamgam_tabdpt", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path, - "data.train_events=100", - "model.tabdpt_estimators=1", - f"model.tabdpt_context_size={context_size or 'null'}", - "split=validation", - ) - path = str(save_config(config, tmp_path / "tabdpt.yaml")) - assert main(["prepare", path]) == 0 - assert main(["evaluate", path]) == 0 - metrics = json.loads((config.path / "validation_metrics.json").read_text()) - assert ( - metrics["checkpoint"] == f"TabDPT (context {context_size or 'all'} per query)" - ) - assert math.isfinite(metrics["metrics"]["auc"]) - - -def test_meta_train_tabicl_with_lora_then_classify_gamgam(tmp_path): - pytest.importorskip("tabicl") - meta = fixture_config( - "delphes_tabicl", - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "meta", - "train.max_epochs=1", - "task.context_events=[4,16]", - "task.query_events=4", - "task.steps_per_epoch=2", - "task.validation_steps=1", - "model.pair_objects=4", - ) - path = str(save_config(meta, tmp_path / "meta.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - # Only the LoRA updates are saved: 12 blocks of 4 weights. - state = load_checkpoint(meta.path / "best.pt")["model_state_dict"] - assert len(state) == 12 * 4 * 2 - assert all(key.endswith((".down", ".up")) for key in state) - - gamgam = fixture_config( - "gamgam_tabicl", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path / "gamgam", - f"checkpoint={meta.path}/best.pt", - "data.train_events=40", - "split=validation", - "model.pair_objects=4", - ) - path = str(save_config(gamgam, tmp_path / "gamgam.yaml")) - assert main(["prepare", path]) == 0 - assert main(["evaluate", path]) == 0 - metrics = json.loads((gamgam.path / "validation_metrics.json").read_text()) - assert metrics["checkpoint"] == f"{meta.path}/best.pt" - assert math.isfinite(metrics["metrics"]["auc"]) - - -#: The encoder of the hybrid tests, as pretrained by JEPA on the fixtures. -ENCODER = { - "hid_size": 16, - "n_layers": 1, - "n_heads": 2, - "ffn_size": 64, - "pair_bias": True, -} - - -def test_event_inputs_are_rows_embeddings_or_both(): - events = pack_events([_objects([0, 0, 3, 3, 4]), _objects([1, 3, 4])], [0, 1]) - batch = collate_sequence_batch(events, torch.tensor([0, 1])) - width = RowLayout((6, 2, 2, 3, 1)).width - for inputs, expected in (("rows", width), ("embeddings", 16), ("both", width + 16)): - config = ModelConfig(name="tabicl", tabular_inputs=inputs, **ENCODER) - model = EventInputs.of(config, num_features=7, num_node_types=5) - assert model(batch).shape == (2, expected) - with pytest.raises(ValueError, match="tabular_inputs"): - EventInputs.of( - ModelConfig(tabular_inputs="objects"), num_features=7, num_node_types=5 - ) - - -@pytest.fixture -def jepa_encoder(tmp_path): - small = [f"model.{key}={value}" for key, value in ENCODER.items()] - config = fixture_config( - "delphes_jepa", - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path / "jepa", - "train.max_epochs=1", - *small, - ) - path = str(save_config(config, tmp_path / "jepa.yaml")) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - return f"{config.path}/last.pt" - - -def _hybrid(name, tmp_path, jepa_encoder, *overrides): - small = [f"model.{key}={value}" for key, value in ENCODER.items()] - config = fixture_config( - name, - GAMGAM_FIXTURES if name.startswith("gamgam") else [TTH_FIXTURE, TTH_FIXTURE], - list(range(5)) if name.startswith("gamgam") else [0, 1], - tmp_path / name, - "model.tabular_inputs=both", - f"model.pretrained={jepa_encoder}", - *small, - *overrides, - ) - path = str(save_config(config, tmp_path / f"{name}.yaml")) - assert main(["prepare", path]) == 0 - return config, path - - -def test_tabicl_classifies_gamgam_on_rows_and_jepa_embeddings(tmp_path, jepa_encoder): - pytest.importorskip("tabicl") - config, path = _hybrid( - "gamgam_tabicl", - tmp_path, - jepa_encoder, - "data.train_events=100", - "model.tabicl_estimators=1", - "split=validation", - ) - assert main(["evaluate", path]) == 0 - metrics = json.loads((config.path / "validation_metrics.json").read_text()) - assert math.isfinite(metrics["metrics"]["auc"]) - - -@pytest.mark.parametrize("freeze", [False, True]) -def test_meta_train_tabicl_with_its_encoder(tmp_path, jepa_encoder, freeze): - pytest.importorskip("tabicl") - meta, path = _hybrid( - "delphes_tabicl", - tmp_path, - jepa_encoder, - "train.max_epochs=1", - "task.context_events=[4,16]", - "task.query_events=4", - "task.steps_per_epoch=2", - "task.validation_steps=1", - f"model.freeze_backbone={str(freeze).lower()}", - ) - assert main(["train", path]) == 0 - state = load_checkpoint(meta.path / "best.pt")["model_state_dict"] - encoder = [key for key in state if key.startswith("inputs.encoder.")] - # A trained encoder is saved; a frozen one comes from model.pretrained. - assert bool(encoder) != freeze - gamgam, path = _hybrid( - "gamgam_tabicl", - tmp_path, - jepa_encoder, - f"checkpoint={meta.path}/best.pt", - "data.train_events=40", - "split=validation", - f"model.freeze_backbone={str(freeze).lower()}", - ) - assert main(["evaluate", path]) == 0 - metrics = json.loads((gamgam.path / "validation_metrics.json").read_text()) - assert math.isfinite(metrics["metrics"]["auc"]) diff --git a/tests/test_tasks.py b/tests/test_tasks.py deleted file mode 100644 index 3e0e799c282986bb2452b8e5f2e7286ce97829ec..0000000000000000000000000000000000000000 --- a/tests/test_tasks.py +++ /dev/null @@ -1,73 +0,0 @@ -import itertools -import math - -import pytest -import torch - -from gnn4colliders.tasks import ( - BinaryClassificationTask, - MulticlassClassificationTask, - macro_scores, - per_class_weighted_mean, - weighted_binary_auc, -) - - -def _pairwise_auc(labels, scores, weights): - total = correct = 0.0 - for i, j in itertools.product(range(len(labels)), repeat=2): - if labels[i] and not labels[j]: - pair = weights[i] * weights[j] - total += pair - if scores[i] > scores[j]: - correct += pair - elif scores[i] == scores[j]: - correct += 0.5 * pair - return correct / total - - -def test_sorted_auc_matches_pairwise_definition_with_ties(): - generator = torch.Generator().manual_seed(0) - labels = torch.rand(60, generator=generator) > 0.4 - scores = torch.randint(0, 8, (60,), generator=generator).float() - weights = torch.rand(60, generator=generator) - assert weighted_binary_auc(labels, scores, weights) == pytest.approx( - _pairwise_auc(labels.tolist(), scores.tolist(), weights.tolist()) - ) - assert math.isnan(weighted_binary_auc(labels & False, scores, weights)) - - -def test_loss_is_the_mean_of_per_class_weighted_means(): - loss = torch.tensor([1.0, 2.0, 3.0, 4.0]) - labels = torch.tensor([0, 0, 2, 2]) - weights = torch.tensor([1.0, 3.0, 1.0, 1.0]) - expected = ((1 * 1 + 3 * 2) / 4 + (3 + 4) / 2) / 2 - assert float(per_class_weighted_mean(loss, labels, weights, 3)) == pytest.approx( - expected - ) - - -def test_macro_scores(): - scores = macro_scores(torch.tensor([0, 0, 1, 1]), torch.tensor([0, 1, 1, 1]), 2) - assert scores["precision"] == pytest.approx((1 + 2 / 3) / 2) - assert scores["recall"] == pytest.approx((0.5 + 1) / 2) - - -def test_multiclass_metrics_exclude_non_positive_weights(): - task = MulticlassClassificationTask() - logits = torch.tensor([[2.0, 0.0], [0.0, 2.0], [2.0, 0.0]]) - labels = torch.tensor([0, 1, 1]) - metrics = task.metrics(logits, labels, torch.tensor([1.0, 1.0, -1.0])) - assert metrics["auc"] == 1.0 - assert metrics["accuracy"] == pytest.approx(2 / 3) - - -def test_binary_task_balances_classes_and_uses_absolute_weights(): - task = BinaryClassificationTask(absolute_weights=True) - logits = torch.tensor([[0.0], [0.0], [0.0]]) - labels = torch.tensor([1, 0, 0]) - loss = task.loss(logits, labels, torch.tensor([-1.0, 1.0, 2.0])) - assert float(loss) == pytest.approx(float(torch.log(torch.tensor(2.0)))) - assert task.metrics(torch.tensor([3.0, -1.0, 0.0]), labels, torch.ones(3))[ - "auc" - ] == pytest.approx(1.0) diff --git a/tests/test_train.py b/tests/test_train.py deleted file mode 100644 index 7413f44aecc326c5e998af66bfc26babac46a0fa..0000000000000000000000000000000000000000 --- a/tests/test_train.py +++ /dev/null @@ -1,153 +0,0 @@ -import copy - -import pytest -import torch -from conftest import random_events - -from gnn4colliders.graphs import collate_graph_batch -from gnn4colliders.loader import EventLoader -from gnn4colliders.models import EdgeNetwork -from gnn4colliders.tasks import MulticlassClassificationTask -from gnn4colliders.train import ( - EarlyStopping, - Trainer, - load_checkpoint, - restore, - save_checkpoint, -) - -cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") - - -def _trainer(seed=0, **options): - torch.manual_seed(seed) - model = EdgeNetwork(7, 3, 0, 8, 3, 2, 1) - optimizer = torch.optim.Adam(model.parameters(), lr=1e-2) - scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.5) - return Trainer( - model, MulticlassClassificationTask(), optimizer, scheduler, **options - ) - - -def _loader(events, **options): - return EventLoader(events, collate_graph_batch, 8, **options) - - -def test_training_reduces_the_loss_and_follows_the_schedule(): - events = random_events() - trainer = _trainer() - seen = [] - history = trainer.fit( - _loader(events, shuffle=True), - _loader(events), - max_epochs=4, - on_epoch_end=lambda train, validation: seen.append( - (trainer.epoch, trainer.optimizer.param_groups[0]["lr"]) - ), - ) - assert seen == [(0, 5e-3), (1, 2.5e-3), (2, 1.25e-3), (3, 6.25e-4)] - assert history[-1][0].loss < history[0][0].loss - assert history[0][1].num_samples == 40 - assert {"accuracy", "auc", "f1"} <= set(history[0][1].metrics) - - -def test_training_loss_early_stopping_needs_no_validation(): - trainer = _trainer(early_stopping=EarlyStopping("train_loss", 1, min_delta=10.0)) - history = trainer.fit(_loader(random_events()), max_epochs=5) - # The first epoch sets the best; the second cannot beat it by min_delta. - assert len(history) == 2 and trainer.epoch == 1 - - -def test_non_finite_loss_is_reported_at_the_end_of_the_epoch(): - events = random_events() - events.object_features[0, 0] = float("nan") - with pytest.raises(FloatingPointError, match="batch 0"): - _trainer().train_epoch(_loader(events)) - - -def test_resumed_training_continues_identically(tmp_path): - events = random_events() - straight = _trainer() - straight.fit(_loader(events, shuffle=True), max_epochs=3) - first = _trainer() - first.fit(_loader(events, shuffle=True), max_epochs=1) - save_checkpoint(tmp_path / "last.pt", first) - resumed = _trainer(seed=1) - restore(resumed, load_checkpoint(tmp_path / "last.pt")) - resumed.fit(_loader(events, shuffle=True), max_epochs=3) - assert resumed.epoch == 2 - for left, right in zip(straight.model.parameters(), resumed.model.parameters()): - torch.testing.assert_close(left, right) - - -def test_bfloat16_autocast_keeps_float32_logits_and_weights(): - events = random_events(seed=6) - batch = next(iter(_loader(events))) - fast = _trainer(precision="bfloat16") - logits = fast.forward(batch) - assert logits.dtype == torch.float32 - assert all(p.dtype == torch.float32 for p in fast.model.parameters()) - torch.testing.assert_close(logits, _trainer().forward(batch), rtol=5e-2, atol=5e-2) - with pytest.raises(ValueError, match="precision"): - _trainer(precision="float16") - - -@pytest.mark.gpu -@cuda -def test_cuda_graph_steps_match_eager_gradients(): - from gnn4colliders.cuda_graphs import CudaGraphTrainStep - - batches = [b.to("cuda") for b in _loader(random_events(64, seed=5), shuffle=False)] - task = MulticlassClassificationTask() - torch.manual_seed(0) - eager = EdgeNetwork(7, 3, 0, 8, 3, 2, 1).cuda() - graphed = copy.deepcopy(eager) - # A zero learning rate keeps parameters fixed so gradients are comparable. - optimizer = torch.optim.Adam(graphed.parameters(), lr=0.0, capturable=True) - for batch in batches[:3]: - optimizer.zero_grad(set_to_none=True) - task.loss(graphed(batch), batch.labels, batch.weight).backward() - optimizer.step() - step = CudaGraphTrainStep(graphed, task, optimizer) - for batch in batches[3:]: - eager.zero_grad(set_to_none=True) - expected = task.loss(eager(batch), batch.labels, batch.weight) - expected.backward() - loss, logits = step(batch) - torch.testing.assert_close(loss, expected.detach(), rtol=1e-5, atol=1e-6) - assert logits.shape == (8, 3) - for left, right in zip(eager.parameters(), graphed.parameters()): - torch.testing.assert_close(right.grad, left.grad, rtol=1e-4, atol=1e-6) - - -@pytest.mark.gpu -@cuda -def test_compiled_graph_steps_with_bfloat16_track_eager_training(): - loader = EventLoader(random_events(96, seed=7), collate_graph_batch, 16) - - def run(**options): - trainer = _trainer(device="cuda", **options) - return [trainer.train_epoch(loader).loss for _ in range(3)] - - eager = run() - fast = run(compile_model=True, cuda_graphs=True, precision="bfloat16") - assert fast == pytest.approx(eager, rel=5e-2) - - -@pytest.mark.gpu -@cuda -def test_cuda_graph_steps_follow_the_learning_rate_schedule(tmp_path): - loader = EventLoader(random_events(96, seed=8), collate_graph_batch, 16) - trainer = _trainer(device="cuda", cuda_graphs=True) - trainer.train_epoch(loader) # warmup steps, then captured replays - # A zero rate must reach the already captured steps. - torch.optim.lr_scheduler.ExponentialLR(trainer.optimizer, gamma=0.0).step() - before = [p.detach().clone() for p in trainer.model.parameters()] - trainer.train_epoch(loader) - for old, new in zip(before, trainer.model.parameters()): - assert torch.equal(old, new) - # Checkpoints stay loadable by eager training on any device. - groups = load_checkpoint(save_checkpoint(tmp_path / "c.pt", trainer))[ - "optimizer_state_dict" - ]["param_groups"] - assert all(type(g["lr"]) is float and not g["capturable"] for g in groups) diff --git a/tests/test_workflows.py b/tests/test_workflows.py deleted file mode 100644 index a0527f257c0c5ed11f9672cecacda2e307e3d972..0000000000000000000000000000000000000000 --- a/tests/test_workflows.py +++ /dev/null @@ -1,205 +0,0 @@ -"""End-to-end command-line runs on the 64-event ROOT fixtures.""" - -import json -import math -import multiprocessing as mp -import os -import socket -import sys - -import numpy as np -import pytest -import torch -import uproot -from conftest import GAMGAM_FIXTURES, PRETRAINED, TTH_FIXTURE, fixture_config - -from gnn4colliders.cli import main -from gnn4colliders.config import save_config -from gnn4colliders.data import load_events -from gnn4colliders.inference import write_root_scores -from gnn4colliders.train import load_checkpoint - - -def _config_file(config, tmp_path): - return str(save_config(config, tmp_path / "experiment.yaml")) - - -def _history(run): - return [ - json.loads(line) for line in (run / "history.jsonl").read_text().splitlines() - ] - - -def test_gamgam_train_evaluate_predict_export_and_resume(gamgam, tmp_path): - path, run = _config_file(gamgam, tmp_path), gamgam.path - assert main(["train", path, "train.max_epochs=3"]) == 0 - history = _history(run) - assert [row["epoch"] for row in history] == [0, 1, 2] - for row in history: - assert row["learning_rate"][0] == pytest.approx( - 1e-4 * 0.99 ** (row["epoch"] + 1) - ) - assert {"accuracy", "auc", "f1", "precision", "recall"} <= set( - row["validation"] - ) - best = load_checkpoint(run / "best.pt") - best_auc = max(row["validation"]["auc"] for row in history) - assert best["best_value"] == pytest.approx(best_auc) - assert history[best["epoch"]]["validation"]["auc"] == pytest.approx(best_auc) - assert load_checkpoint(run / "last.pt")["epoch"] == 2 - - assert main(["evaluate", path]) == 0 - summary = json.loads((run / "test_metrics.json").read_text()) - events = load_events(gamgam.data.cache) - assert summary["num_samples"] == int((events.fold == 9).sum()) - assert all(math.isfinite(value) for value in summary["metrics"].values()) - assert set(summary["seconds"]) == {"setup", "inference"} - - subset = tmp_path / "subset.json" - assert main(["evaluate", path, "data.eval_events=10", f"output={subset}"]) == 0 - assert json.loads(subset.read_text())["num_samples"] == 10 - - assert main(["predict", path, "split=validation"]) == 0 - predictions = np.load(run / "validation_predictions.npz") - assert predictions["scores"].shape == (int((events.fold == 8).sum()), 5) - np.testing.assert_allclose(predictions["scores"].sum(axis=1), 1, rtol=1e-5) - assert (predictions["fold"] == 8).all() - first = str(predictions["source_file"][0]) - entries = predictions["entry"][predictions["source_file"] == first] - tree = uproot.open(first)["mini"].arrays(["eventNumber"], library="np") - assert (tree["eventNumber"][entries] % 10 == 8).all() - - pytest.importorskip("onnxruntime") - assert main(["export", path]) == 0 - assert (run / "model.onnx").is_file() and (run / "model.json").is_file() - - assert ( - main(["train", path, "train.max_epochs=4", f"train.resume={run}/last.pt"]) == 0 - ) - assert [row["epoch"] for row in _history(run)] == [0, 1, 2, 3] - - -def test_root_scores_are_aligned_by_entry(tmp_path): - entries = np.array([3, 0, 5]) - scores = np.array([[0.1, 0.9], [0.8, 0.2], [0.5, 0.5]]) - output = write_root_scores( - scores, - entries, - input_path=GAMGAM_FIXTURES[0], - output_path=tmp_path / "scores.root", - tree_name="mini", - ) - tree = uproot.open(output)["mini"].arrays(library="np") - assert tree["selection_pass"].tolist()[:6] == [1, 0, 0, 1, 0, 1] - assert tree["score_class_1"][3] == pytest.approx(0.9) - assert np.isnan(tree["score_class_0"][1]) - # Jagged branches are copied unchanged. - original = uproot.open(GAMGAM_FIXTURES[0])["mini"]["jet_pt"].array() - copied = uproot.open(output)["mini"]["jet_pt"].array() - assert copied.tolist() == original.tolist() - - -def test_fine_tuning_from_the_pretrained_model(tmp_path): - config = fixture_config( - "gamgam_finetune", - GAMGAM_FIXTURES, - list(range(5)), - tmp_path, - "train.max_epochs=1", - f"model.pretrained={PRETRAINED}/multiclass_pretrained_model_12/model_epoch_71.pt", - ) - path = _config_file(config, tmp_path) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - assert main(["evaluate", path]) == 0 - state = load_checkpoint(config.path / "best.pt")["model_state_dict"] - assert state["classifier.weight"].shape == (5, 64) - - -@pytest.mark.parametrize("name", ["gamgam_transformer", "gamgam_deep_sets"]) -def test_sequence_models_train_on_the_same_cache(name, gamgam, tmp_path): - config = fixture_config( - name, - GAMGAM_FIXTURES, - list(range(5)), - tmp_path, - f"data.cache={gamgam.data.cache}", - "train.max_epochs=1", - ) - path = _config_file(config, tmp_path) - assert main(["train", path]) == 0 - assert main(["evaluate", path]) == 0 - pytest.importorskip("onnxruntime") - assert main(["export", path]) == 0 - assert json.loads((config.path / "model.json").read_text())["inputs"] == [ - "tokens", - "token_mask", - ] - - -def test_binary_delphes_task(tmp_path): - # The same file twice with labels 0 and 1 exercises the binary metrics. - config = fixture_config( - "tth_cp_even_odd", - [TTH_FIXTURE, TTH_FIXTURE], - [0, 1], - tmp_path, - "train.max_epochs=1", - ) - path = _config_file(config, tmp_path) - assert main(["prepare", path]) == 0 - assert main(["train", path]) == 0 - assert math.isfinite(_history(config.path)[0]["validation"]["auc"]) - - -def _port() -> int: - with socket.socket() as sock: - sock.bind(("127.0.0.1", 0)) - return sock.getsockname()[1] - - -def _rank(rank: int, port: int, argv: list[str], queue) -> None: - os.environ.update( - RANK=str(rank), - LOCAL_RANK=str(rank), - WORLD_SIZE="2", - MASTER_ADDR="127.0.0.1", - MASTER_PORT=str(port), - ) - torch.set_num_threads(1) - sys.argv = ["gnn4colliders", *argv] - queue.put((rank, main(argv))) - - -def test_training_runs_on_two_ddp_ranks(gamgam, tmp_path): - config = _config_file(gamgam, tmp_path) - # Gloo ranks on the CPU: MPS has no collectives. - argv = ["train", config, "train.max_epochs=2", "train.device=cpu"] - context = mp.get_context("spawn") - queue = context.Queue() - port = _port() - processes = [ - context.Process(target=_rank, args=(rank, port, argv, queue)) - for rank in range(2) - ] - try: - for process in processes: - process.start() - for process in processes: - process.join(300) - assert [process.exitcode for process in processes] == [0, 0] - assert sorted(queue.get(timeout=5) for _ in processes) == [(0, 0), (1, 0)] - finally: - # A regression can deadlock a rank in a collective; never leave it. - for process in processes: - if process.is_alive(): - process.kill() - process.join(5) - history = _history(gamgam.path) - events = load_events(gamgam.data.cache) - assert len(history) == 2 - assert load_checkpoint(gamgam.path / "last.pt")["epoch"] == 1 - # Validation covers every event once; training pads ranks to equal size. - assert history[0]["validation"]["events"] == int((events.fold == 8).sum()) - train_events = int((events.fold < 8).sum()) - assert history[0]["train"]["events"] == train_events + train_events % 2 diff --git a/uv.lock b/uv.lock deleted file mode 100644 index 2476f38dd074dfe77352c2e2886c592397991ded..0000000000000000000000000000000000000000 --- a/uv.lock +++ /dev/null @@ -1,1703 +0,0 @@ -version = 1 -revision = 2 -requires-python = "==3.14.*" -resolution-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", - "sys_platform == 'win32'", - "sys_platform == 'emscripten'", - "(platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", -] - -[[package]] -name = "annotated-types" -version = "0.8.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" }, -] - -[[package]] -name = "awkward" -version = "2.13.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "awkward-cpp" }, - { name = "fsspec" }, - { name = "numpy" }, - { name = "packaging" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/43/97/2728ab879ed3edaecfbd8e9daa7d88a2f89a7ea584e413c4d8971fc0414d/awkward-2.13.0.tar.gz", hash = "sha256:36f127573295e1ddf65b551bf071b803ad4af21bd14519f4c4dd34953cb3efc2", size = 6479149, upload-time = "2026-08-14T16:52:43.873Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/25/34/b8eece9a8d024defc08418ea883361b3d7e3300392719ee27535e93885a5/awkward-2.13.0-py3-none-any.whl", hash = "sha256:ff40879e7179a6f14f4c4ee5f5297e3dcd027b99b4c64188bdb7dccdf625e982", size = 990112, upload-time = "2026-08-14T16:52:42.018Z" }, -] - -[[package]] -name = "awkward-cpp" -version = "56" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/43/92/6325292ccc2be3ede227fd0b4bd7b23bc7597790f31227171bec4009d2df/awkward_cpp-56.tar.gz", hash = "sha256:cd3635fab926c6630c0a85d92b59a15851c4f5a881ea749984069027634f8f52", size = 1501723, upload-time = "2026-08-14T15:30:10.748Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/93/21/43db481a3a182b0389483558dbbf0e181de1231bc123579b7f054ae3a343/awkward_cpp-56-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:cb0e45e17432143d3de05846c4059ddfdb4d373670674e45849f0f57f454d1e5", size = 628964, upload-time = "2026-08-14T15:29:39.698Z" }, - { url = "https://files.pythonhosted.org/packages/07/f8/e4d048a16a692b5eb728f7b6d6377bc4ddf7649504259263be24fc425c7f/awkward_cpp-56-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:82729789ce90205a7364bf25b7bd27bdfaf4df7a88dc3682a0b6ffd437751d53", size = 581954, upload-time = "2026-08-14T15:29:41.363Z" }, - { url = "https://files.pythonhosted.org/packages/6e/a9/b0359161706b20b98e5a39b68f29abe6f4921a83cb50cc62446fde6a4168/awkward_cpp-56-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a813ec55efcf3080e04c0e83e40ef4509152e9d413854eac1fd65a8350e2364f", size = 626876, upload-time = "2026-08-14T15:29:42.913Z" }, - { url = "https://files.pythonhosted.org/packages/58/b3/2232b297322554f0711e6f0c6b006ab2a5d066c445c28df6029bcee732ac/awkward_cpp-56-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ae238b66a2e59d43af82dabbfc3e779f871f82a3cf96f54172e5861a849aeb3", size = 699174, upload-time = "2026-08-14T15:29:44.448Z" }, - { url = "https://files.pythonhosted.org/packages/c2/35/cab44d2b01db3986492184bb48f2f670b09a1dc44e901448cac1fc3bd998/awkward_cpp-56-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ce618698397ef177e3c642f5197f6b79c1bafdbce4a49e274d044bd3363dd254", size = 1627728, upload-time = "2026-08-14T15:29:46.085Z" }, - { url = "https://files.pythonhosted.org/packages/9c/f3/f2a194edda3f76def8b5dd6f59009a19df39ca89dd4b17a52279053ed153/awkward_cpp-56-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ca637ee30a127d92f86268bfef16074557b3c21b6ad18b89dad811a53dd9f3bf", size = 1745659, upload-time = "2026-08-14T15:29:47.765Z" }, - { url = "https://files.pythonhosted.org/packages/a8/05/086e5fdd4cb75fcab4f9c5a7f161687c8415cd64c0ba5b548cd3dd9c3739/awkward_cpp-56-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:33b24ea7b2f54e8d1cf514ef8246513ab8c9f3d335407b7c5b3ea5174c5c66b2", size = 235559, upload-time = "2026-08-14T15:29:49.253Z" }, - { url = "https://files.pythonhosted.org/packages/02/a2/35bea035807d1bfc25fe38f7a317410812ee0b94b97ee0da96e54839f245/awkward_cpp-56-cp314-cp314-win32.whl", hash = "sha256:9abb36eaac6f0f6a9419284f18fa1b0031800982cb82b3779e2aa4eb96967cb9", size = 726000, upload-time = "2026-08-14T15:29:50.888Z" }, - { url = "https://files.pythonhosted.org/packages/19/29/1361c17bbfd544a4872b6724ba840d12ff5a76665e6e16977a52dc434a57/awkward_cpp-56-cp314-cp314-win_amd64.whl", hash = "sha256:6a7f5f0ff66aaf731c63b1ca4cd0a67026140396ffa83531d3802ee5909cc88d", size = 761510, upload-time = "2026-08-14T15:29:52.351Z" }, - { url = "https://files.pythonhosted.org/packages/4f/4d/f8beec1c1b17b8e733b7eda7b279083938f620a7c4cf720c2f4efb2909a8/awkward_cpp-56-cp314-cp314-win_arm64.whl", hash = "sha256:517c9eaafb0cb877ea9860f13c4cf737d1f98b017c5e1ca2fd74eb78b424f45c", size = 903318, upload-time = "2026-08-14T15:29:54.506Z" }, - { url = "https://files.pythonhosted.org/packages/65/c1/60bdb584d22dc664cef3bef929d349da2a7243ba35c76eb2f2560cb6d4db/awkward_cpp-56-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:212ed06dc00591e51e63452cc3fb9eb5a3ae31340ab677ecad9dc25b763b857e", size = 640282, upload-time = "2026-08-14T15:29:56.143Z" }, - { url = "https://files.pythonhosted.org/packages/76/b5/faebdd67804aca811b188dd8525725aeadaade6da6166e017d91a3628ec2/awkward_cpp-56-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b76a58c754d259b463d1dd6837a51e8b1c368621e3a3d5939f5fa9c49d405ea", size = 594824, upload-time = "2026-08-14T15:29:57.96Z" }, - { url = "https://files.pythonhosted.org/packages/d8/35/19dea873740c369b5802faf823e7266ae624b01745eb8981a94ca3ba1f4a/awkward_cpp-56-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ef6da944624eb3aa46a4eba10b395de99c2f7bdeb130d66e01122108f6747258", size = 631000, upload-time = "2026-08-14T15:29:59.514Z" }, - { url = "https://files.pythonhosted.org/packages/ab/e0/e256df80030dcef60151fb952a334f5a1cb582da81e2c9700c20ceafd699/awkward_cpp-56-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e7e5ad0a07d6211fd97c33a67887d1a81392fb53bc351f23e0414d2ea2c90ad5", size = 702667, upload-time = "2026-08-14T15:30:01.031Z" }, - { url = "https://files.pythonhosted.org/packages/d7/3c/3ed9f61fd7e2561c84c81120acbe7578148b7ebb10ed4509e46f51899ade/awkward_cpp-56-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ac0e58f2b8eb7b54a167b0a2885a8eb05733c60019b0145a464a8b526d9e924a", size = 1633280, upload-time = "2026-08-14T15:30:02.776Z" }, - { url = "https://files.pythonhosted.org/packages/85/1a/c722b659c20cad01c334e897a592573ce0d715cefbbcbeff1951dd9aaad2/awkward_cpp-56-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:8ffd50a0c5b7d0afb780334e7898b17e6044bbc03eb292201471e34652d79a98", size = 1745426, upload-time = "2026-08-14T15:30:04.608Z" }, - { url = "https://files.pythonhosted.org/packages/e7/cf/66264c97a7db4c122190f25f2ce06bb85e172728e5ac1b68e68f9506a7d6/awkward_cpp-56-cp314-cp314t-win32.whl", hash = "sha256:4f53df218cdf67e290235c6b75bf9bec142b73159c261dc3997865d5b1e016e6", size = 733789, upload-time = "2026-08-14T15:30:06.07Z" }, - { url = "https://files.pythonhosted.org/packages/3f/67/e89d209d147ffa6c7c5ef636829aa97b939a14e83a002ed012936067e6d9/awkward_cpp-56-cp314-cp314t-win_amd64.whl", hash = "sha256:6bc4ebdbd1df44bf8e19de77cef39d67b023b3b7b753e30ace2d75d1f8ecc0e0", size = 768029, upload-time = "2026-08-14T15:30:07.476Z" }, - { url = "https://files.pythonhosted.org/packages/04/fc/2db6fbe8c3e94e33c8b2bd4fbb989bbf11c2b070e68bba0bff5a2bb0756a/awkward_cpp-56-cp314-cp314t-win_arm64.whl", hash = "sha256:f3693744ef6d149518a1e46bd5c2dfbad4260b96a52721911a5a56edd5836870", size = 913341, upload-time = "2026-08-14T15:30:09.03Z" }, -] - -[[package]] -name = "certifi" -version = "2026.7.22" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, -] - -[[package]] -name = "charset-normalizer" -version = "3.5.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e5/3f/143b048436775b0f76ac3eec145c019e8173ccc2885c8f20319b996d5e83/charset_normalizer-3.5.1.tar.gz", hash = "sha256:6117b84ea48435e5356dc737f5121485c30920ba43375fa7b434fd753df0eac3", size = 171764, upload-time = "2026-08-15T08:20:44.807Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/29/cd/2b812ce5e888f1ce69a5350281e58aab07ae64a958ecae8912f30865718e/charset_normalizer-3.5.1-cp314-cp314-android_24_arm64_v8a.whl", hash = "sha256:774d157f112367ff4abd29019f38f023c24e00e56edc7829c20e358a5a913ad8", size = 212318, upload-time = "2026-08-15T08:18:04.403Z" }, - { url = "https://files.pythonhosted.org/packages/9e/4a/a6ee107430768a5334e6d63f31f148a04a1a491ef161a1ac9415a73f2fa8/charset_normalizer-3.5.1-cp314-cp314-android_24_x86_64.whl", hash = "sha256:26422d45fd13551cf564c58932f7d72b4f58b93b0fcf18c35ba6be12b46bb102", size = 224897, upload-time = "2026-08-15T08:18:05.997Z" }, - { url = "https://files.pythonhosted.org/packages/c3/d9/35ae3f64f29d0179c35c3baefe575904df2913dde519129c7f75995a2b1d/charset_normalizer-3.5.1-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:09a7bba9f739468c8e78c36a75c33768e53cb1959fc638f510454c14683f00d5", size = 194848, upload-time = "2026-08-15T08:18:07.397Z" }, - { url = "https://files.pythonhosted.org/packages/74/76/f2fc7380f056cc273a53af37f50d08ad54b2c59f61078f31432edcf1c2bd/charset_normalizer-3.5.1-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:4c9548dc78002099910abaebc0a72ac58b7d30931869e0351c09b507dff4ece3", size = 198163, upload-time = "2026-08-15T08:18:08.989Z" }, - { url = "https://files.pythonhosted.org/packages/e9/40/095ce62fa078483cccc1fa2b36e6bc9580b85422a20ee9f925341c50e44f/charset_normalizer-3.5.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:c428c6c31eb5f4277d7f8eccaf767fbd548ddd5ce3c8b4f4cbbfab3d96b5904c", size = 341823, upload-time = "2026-08-15T08:18:10.458Z" }, - { url = "https://files.pythonhosted.org/packages/f1/5a/0e58b1c04a1596e0256f407274a92d5fb2ee21324409d1fab1da48a65b5b/charset_normalizer-3.5.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2f06b7eae9dbe77fe1d644ca244dad508de8d302870a43f3c559b521270938a0", size = 242458, upload-time = "2026-08-15T08:18:11.989Z" }, - { url = "https://files.pythonhosted.org/packages/22/95/b4618ce912e6db0b1aae89ba788e38e8a7eba0f3025cc66e8c0699f977b2/charset_normalizer-3.5.1-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:6b7430cf5728e68f6c462254009a6ef4086e1bea43cf2f57aa9c55fb4f50ff96", size = 226717, upload-time = "2026-08-15T08:18:13.401Z" }, - { url = "https://files.pythonhosted.org/packages/8a/76/c681192bbda3d55356db5dadd64381d5202b37c6b598fcda5282e88b5d3d/charset_normalizer-3.5.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ab743e9bc90c1f73552ec33e10e3331315acd2c397b36065b591b0181de533cc", size = 266111, upload-time = "2026-08-15T08:18:14.961Z" }, - { url = "https://files.pythonhosted.org/packages/88/be/55127bfca72c0cff6c022488d140d7c5b04c771e3b72e9bdb4836d54979d/charset_normalizer-3.5.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f6f7deae3feb4edfa2efaf7c574fe88cbf055038a6abdb40188e4fff66d5699f", size = 263128, upload-time = "2026-08-15T08:18:16.515Z" }, - { url = "https://files.pythonhosted.org/packages/e0/91/39c3af510b0aa32bbda03374259200f28430febfd1bf5e511fe765282ce5/charset_normalizer-3.5.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:15f024313246a4ed976c60f440bb8d257815513a681d212ff74fd46f7d715a90", size = 251240, upload-time = "2026-08-15T08:18:18.127Z" }, - { url = "https://files.pythonhosted.org/packages/1c/a5/cbe418bbc6ecdfc3e05a0116002897c4b403a5e838d697e64c78e9f0190d/charset_normalizer-3.5.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:823f82903d189af463d7df250ef1f7f696f3cee08cc8d91deb565e8d425f6506", size = 245282, upload-time = "2026-08-15T08:18:19.625Z" }, - { url = "https://files.pythonhosted.org/packages/cc/a4/689bb42e8e7cd492f3cb64907c6bc00ad247ec9a3628cd3f8eed126e8ae1/charset_normalizer-3.5.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:01e93745f7f219b703b60ba7afead36cfc4242782be5af484673fc500df12da5", size = 244597, upload-time = "2026-08-15T08:18:21.121Z" }, - { url = "https://files.pythonhosted.org/packages/c1/ce/9962938e179cf9f699d3f1e7b3114b5d7642dee6a893745229f9dd04f274/charset_normalizer-3.5.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:329fc3ccb63ad22d867d84c2adea759a64079a37ba4a343433b02c7a2816871e", size = 231376, upload-time = "2026-08-15T08:18:22.57Z" }, - { url = "https://files.pythonhosted.org/packages/85/54/46000450ada53bd9eac5429a2c8c54cd2d9b39c0c255f229aea9af0948a5/charset_normalizer-3.5.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:bb57753e36e4855b8ca375069482250a6246372331a3e4f3407eaebb007443f5", size = 266715, upload-time = "2026-08-15T08:18:24.235Z" }, - { url = "https://files.pythonhosted.org/packages/3d/bb/618749d70f792b44252a777bf89bfb86823b9bbc1ea13fe8ce759b07f38a/charset_normalizer-3.5.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:fce8cbd4997efeb450bd298b54f755dcdff18d496f7a5ddbb4867c6d7c88fdc3", size = 245848, upload-time = "2026-08-15T08:18:25.726Z" }, - { url = "https://files.pythonhosted.org/packages/7e/3f/ffb64458527c7668031d5eb095d978de561958dc9f5b53f8e488a533e603/charset_normalizer-3.5.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:6c9cdde8becb25a7fde49924511aa2644d6f8081cc8df8e9452724303348d8e3", size = 264521, upload-time = "2026-08-15T08:18:27.193Z" }, - { url = "https://files.pythonhosted.org/packages/4f/ab/74a55fd803916a35ac461daf002708191aac19b546b80dc8cabfedc63d98/charset_normalizer-3.5.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:9ac4444d8d4fd4c4bd08bf451ed3167aa9e7ec6cdb41b648794f1d1103652e36", size = 253054, upload-time = "2026-08-15T08:18:28.568Z" }, - { url = "https://files.pythonhosted.org/packages/a0/2a/6a9034b7d3c60b17499afb482df5878bf9fa20b50cc3887d5ef017a833db/charset_normalizer-3.5.1-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:f03ac127268b43ef4fe9e6ab6794a6794b49485a0cc0c1db79876d2f33f75bc7", size = 140580, upload-time = "2026-08-15T08:18:30.214Z" }, - { url = "https://files.pythonhosted.org/packages/f3/46/1d362e1a00d035d66b9869e1281eee115907f7e390a16a07824ab5737360/charset_normalizer-3.5.1-cp314-cp314-win32.whl", hash = "sha256:1f5883d77fd409a261abb5dc8ccbe335720d798b1de4abb3b1d47ccbbc76b53b", size = 180325, upload-time = "2026-08-15T08:18:31.877Z" }, - { url = "https://files.pythonhosted.org/packages/7a/7c/4938c329b6a9d446f6a59aa2092ff7118f274209b5ed0e26893d1d30a63c/charset_normalizer-3.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:c658c50ac0c98cd755a2dd50b7977d3bca7df401dcc47fbdfa87db53ef7d4e8b", size = 204175, upload-time = "2026-08-15T08:18:33.466Z" }, - { url = "https://files.pythonhosted.org/packages/ac/33/eeb384dbd8dec570661354592f4f2e1b2fcc92585624d146a000caf53841/charset_normalizer-3.5.1-cp314-cp314-win_arm64.whl", hash = "sha256:4bea7f8ebe90bbd7f0e4a2de42ca6924ba23e3e76418c408ff82f1d46fabd687", size = 184123, upload-time = "2026-08-15T08:18:34.913Z" }, - { url = "https://files.pythonhosted.org/packages/1c/6c/c73fa9d5a85f6ab05395de61c5f6984e0a9ff40bb5ff888d46dff02526c6/charset_normalizer-3.5.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:fbc597639158fd7c14d55e808718848319540f51b0e6746e3eefa59723a4a348", size = 381682, upload-time = "2026-08-15T08:18:36.349Z" }, - { url = "https://files.pythonhosted.org/packages/30/c7/63565f860921457feba93bae6c86fb7746deb4cffeed2f375cb845318146/charset_normalizer-3.5.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e71c909f353863b2b89c83de2ebed71ea6d0df8a6ef65a128193c5e650766bef", size = 240826, upload-time = "2026-08-15T08:18:37.887Z" }, - { url = "https://files.pythonhosted.org/packages/06/ae/7ae8807410dfa33f8e6f1715740adeaafa8a816cc4cb33508f54b1f7c896/charset_normalizer-3.5.1-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:7ac76cf9afd34929d76eb7fcb63be476a4853d8a96f0dcf2d0db68a0cbdf9885", size = 227861, upload-time = "2026-08-15T08:18:39.315Z" }, - { url = "https://files.pythonhosted.org/packages/e9/a3/887c1642f0da26000b0e0652d91071113c0e72cea33952e225cf589f49a9/charset_normalizer-3.5.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a3a370082ce34d0612f421e15fe011c53bb1feff21a26d06ad4fb244dab5a375", size = 260758, upload-time = "2026-08-15T08:18:40.88Z" }, - { url = "https://files.pythonhosted.org/packages/3e/11/e6f5b9a3d0e55b0ef7505cd3765cdd48f22db89994c947b316f52f801fd8/charset_normalizer-3.5.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:256dd4d85d9e4dc595e2bc983c980e73f62ddeb3165c58b4c3dfe78c5c8548c1", size = 259950, upload-time = "2026-08-15T08:18:42.351Z" }, - { url = "https://files.pythonhosted.org/packages/1b/ee/e4e10a94d51cd1ee638aa7e00b65399e6b2a4e8376ab6d2eac9f95586671/charset_normalizer-3.5.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58d4aa13a59c969dbfdf9e6a9560e242cbfd9e8a8f50c2747714df1a423adf65", size = 249329, upload-time = "2026-08-15T08:18:43.914Z" }, - { url = "https://files.pythonhosted.org/packages/c4/25/d5f4198819e6059735a84e8d0bfb72dc33976da67b97adcd3fb5a5e07ec6/charset_normalizer-3.5.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:0c6dfb5ca6723eeed15aa8e564a014d69fcb8812f94eef11fe3631e0508199f5", size = 243137, upload-time = "2026-08-15T08:18:45.368Z" }, - { url = "https://files.pythonhosted.org/packages/a5/e9/e925ca7569cf9fb9701fd82503fee73eea5268fdb856bdd64947092d3daa/charset_normalizer-3.5.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:c010f5581d9c612804cc59fcf7b524b707fbcb72828551237ab545bb5c7034af", size = 242820, upload-time = "2026-08-15T08:18:46.842Z" }, - { url = "https://files.pythonhosted.org/packages/34/17/672c251a888ed2aebcdd2fe830ad0104e25ff83c43f5c4f9c15e9fc6853c/charset_normalizer-3.5.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:52ec005752a56ae79547a05c0139ca2501a0c866390b6115008456b9f0e7cde1", size = 230504, upload-time = "2026-08-15T08:18:48.353Z" }, - { url = "https://files.pythonhosted.org/packages/3f/fc/f6a85abebd42ce4da2f1db0aa56cc6a0df1995e318b3875d14401b8381d1/charset_normalizer-3.5.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:2bced4061f000f7187254a02ad3433ae17eaf991747ceea2f478422590a5bba9", size = 263087, upload-time = "2026-08-15T08:18:49.859Z" }, - { url = "https://files.pythonhosted.org/packages/98/66/7c42677e739ba66746b297e2046918d793078094dc239e1e72768cffccc6/charset_normalizer-3.5.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:9eea3ab2597a5e65fe65296e2d6a84570845a6b55532d90333d740d48bbc850a", size = 243269, upload-time = "2026-08-15T08:18:51.601Z" }, - { url = "https://files.pythonhosted.org/packages/de/d8/a50b79237f417af10f8c2a501ce8d1ca87829a22e69117891ca4ba20a69e/charset_normalizer-3.5.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:496846868fea80e479324862fa877f02411f2fd0f83b79ccee2607aa68b2a032", size = 258766, upload-time = "2026-08-15T08:18:53.23Z" }, - { url = "https://files.pythonhosted.org/packages/2e/1d/0fc91aeaeb3c83b748f532399ce67cf84604b48297405d740000f7a9e786/charset_normalizer-3.5.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:85d5855daafc240cc045c026d7a15fd198a09b0fc8ff6f5ecbb5297b509cb11e", size = 250814, upload-time = "2026-08-15T08:18:54.768Z" }, - { url = "https://files.pythonhosted.org/packages/ae/10/3d8c777cf9024615295aa1b808324ad5b4a77855869c00824bad74ffaf8a/charset_normalizer-3.5.1-cp314-cp314t-win32.whl", hash = "sha256:58d3e12c88e0950bca850ae1f7c256055c097639c2edb9eb123af9807d8b15e4", size = 191074, upload-time = "2026-08-15T08:18:56.305Z" }, - { url = "https://files.pythonhosted.org/packages/4d/81/ae557d3c44d1a1d688696d60563413a0866a91b7ebc50f20df838be3d8c8/charset_normalizer-3.5.1-cp314-cp314t-win_amd64.whl", hash = "sha256:acaf604462bf330b0d07e7a07c1d6e4adac79e5fb13e9c5140590542cafacc00", size = 216476, upload-time = "2026-08-15T08:18:57.889Z" }, - { url = "https://files.pythonhosted.org/packages/27/e9/61c01fb8b804692569c036b3fc50495814502dcf13a60649c6055390b02c/charset_normalizer-3.5.1-cp314-cp314t-win_arm64.whl", hash = "sha256:fdb8a068947befafba9952162645dc2fecaeb400e64584829ed5e9b2fbe21a7f", size = 194115, upload-time = "2026-08-15T08:18:59.418Z" }, - { url = "https://files.pythonhosted.org/packages/5b/97/fb4e82231aba271ffd775a1b4993b0defc4e3059f286ae41d9433409fe85/charset_normalizer-3.5.1-cp37-abi3-macosx_10_9_universal2.whl", hash = "sha256:41876ee62a3dddf48ff1121ad8f0798032aa03f2fd35f21f34a4cab14f18d8d2", size = 331467, upload-time = "2026-08-15T08:19:50.959Z" }, - { url = "https://files.pythonhosted.org/packages/9f/2f/fe3f187327aac18e2d54e9d2b08e15d27bf9b642d9e51c219f130fc34d1a/charset_normalizer-3.5.1-cp37-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:a6dac12ff6b846103483683f60c5f8fee205121adc58ffd87e90a90a3af69e99", size = 253057, upload-time = "2026-08-15T08:19:52.654Z" }, - { url = "https://files.pythonhosted.org/packages/d7/c7/9e48cee5c161fe24da823b61bf381921d77cb994a0a4de148e95018c1984/charset_normalizer-3.5.1-cp37-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cee5dd7c6fb5dd52a0fe2a740f9bc6e3593f5f8b1788bde49de02086f30182b2", size = 240930, upload-time = "2026-08-15T08:19:54.163Z" }, - { url = "https://files.pythonhosted.org/packages/49/e0/716601f3cc69be7b198951150c75ead1ece33c3c8036ff6ffa46029659a0/charset_normalizer-3.5.1-cp37-abi3-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:343fb4f2821043bd87095f7b08a1a181febc8e36ac64212143bbfd0a0e1bc235", size = 230822, upload-time = "2026-08-15T08:19:55.807Z" }, - { url = "https://files.pythonhosted.org/packages/d3/05/71bfc5caa0abcc45aea1f6a4d50ac68e59605ddc7666fe8494f4cd229665/charset_normalizer-3.5.1-cp37-abi3-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:ae4a097991662cd4fff0ddc74e0fe7874f82e00042fa0ea00855645ed0c79598", size = 260037, upload-time = "2026-08-15T08:19:57.312Z" }, - { url = "https://files.pythonhosted.org/packages/c3/92/de7e32ed05341e7a9c4c877c318418197b7f2d66a3b68d561bf2ac57ca3e/charset_normalizer-3.5.1-cp37-abi3-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:4b599739b93b2cbeded49645ae3c8d1405c29ddfbceac1545c87a3f9580a9e96", size = 255097, upload-time = "2026-08-15T08:19:59.056Z" }, - { url = "https://files.pythonhosted.org/packages/f5/7b/ade0a122600319dfa0b1000ab0f9731c94a817904cf3c5de408c73a4ede7/charset_normalizer-3.5.1-cp37-abi3-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b39b69b347e5e47a3b5b8cfc005c68c1ba347474e3960236c4944a8ecd174962", size = 250166, upload-time = "2026-08-15T08:20:00.612Z" }, - { url = "https://files.pythonhosted.org/packages/75/9c/019fbb9f4834491a160951349b1a3714439376f66e5f7cf18b4f18f0c7aa/charset_normalizer-3.5.1-cp37-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:a2028475ba855475b8b4d3cfeb4994269c967aea8b9892dfba907f4263a863a3", size = 241821, upload-time = "2026-08-15T08:20:02.321Z" }, - { url = "https://files.pythonhosted.org/packages/2b/b8/11d4840bfc99330cc7fbcc2681ee5a044553a6e77655508d8f9b2bff7b34/charset_normalizer-3.5.1-cp37-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:36047af20e17097c3bb9476c2b7655f2f7aa51322c0ba58c07695bedf755a950", size = 232529, upload-time = "2026-08-15T08:20:04.008Z" }, - { url = "https://files.pythonhosted.org/packages/18/96/2b3a21492d9f65171ac75d872f5018260013d00bfa0ff70ec9f179148cbd/charset_normalizer-3.5.1-cp37-abi3-musllinux_1_2_ppc64le.whl", hash = "sha256:4c4fb141a727957c93edfe5c32a26ceb6b5f6461d67146e2d39f51e16170bea8", size = 260348, upload-time = "2026-08-15T08:20:05.877Z" }, - { url = "https://files.pythonhosted.org/packages/d6/aa/a69a2028e8bd052476c245460ab19d7de595de084dd968f2d75cd50c3e25/charset_normalizer-3.5.1-cp37-abi3-musllinux_1_2_riscv64.whl", hash = "sha256:2f293479cce755c75f1697e87c409b7ae4c555c7dfecb6e988ad13abba943031", size = 247234, upload-time = "2026-08-15T08:20:07.487Z" }, - { url = "https://files.pythonhosted.org/packages/35/8a/3d130aeabcaf3d2466af76b7b141c08d9e89c9016ab4b7cdd0f7dc2d1c62/charset_normalizer-3.5.1-cp37-abi3-musllinux_1_2_s390x.whl", hash = "sha256:3588e376b3ea2eea84976f67273d679f229e24c66dce7b82ae45aef04ff6e072", size = 256917, upload-time = "2026-08-15T08:20:09.142Z" }, - { url = "https://files.pythonhosted.org/packages/80/c2/a7379b840292d0c1ab9fbd17d1f3967aa81794dc95bc74be8999d7fedcf7/charset_normalizer-3.5.1-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e199fb99720074809a7720f1c0b4d919eea8b87e88713e0f8f602f7bef543d9d", size = 254846, upload-time = "2026-08-15T08:20:10.727Z" }, - { url = "https://files.pythonhosted.org/packages/01/65/d43b714731bb2f40d4053dfa00ecfc1c5a301f8e3316c5db3a09af59fe94/charset_normalizer-3.5.1-cp37-abi3-win32.whl", hash = "sha256:dd732602a7009217f658d5863d12d79d373a4de0eebc111094bcdd3bb8e0a6cc", size = 174216, upload-time = "2026-08-15T08:20:12.334Z" }, - { url = "https://files.pythonhosted.org/packages/35/4f/b911ed898b26a09789eba9c9200c999aff6c61b4bafaf4838e56d1a1e1a3/charset_normalizer-3.5.1-cp37-abi3-win_amd64.whl", hash = "sha256:70055ff39b97c99e7ae40ea3e393fb62aa2e44dbd9b29f8d14f42fb0025c3959", size = 199764, upload-time = "2026-08-15T08:20:13.908Z" }, - { url = "https://files.pythonhosted.org/packages/f0/a7/920baf467bfd9bf689f3b318340f37aee4572a71f162bd8db51da55ba4fa/charset_normalizer-3.5.1-cp37-abi3-win_arm64.whl", hash = "sha256:87e4f41d375c0b9be2fb5251aee4b8a689169e134535aed81bf085c3b647451e", size = 287318, upload-time = "2026-08-15T08:20:15.551Z" }, - { url = "https://files.pythonhosted.org/packages/cc/61/d01fc49b8dea277640b55a9e15960dbca9fdc8c9fde18e572d39c59f4019/charset_normalizer-3.5.1-py3-none-any.whl", hash = "sha256:6df0ec430f9a831772c23ca5a224cba36517a58a84bb32c32bb59a9fa67c47f6", size = 68658, upload-time = "2026-08-15T08:20:43.306Z" }, -] - -[[package]] -name = "cloudpickle" -version = "3.1.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/27/fb/576f067976d320f5f0114a8d9fa1215425441bb35627b1993e5afd8111e5/cloudpickle-3.1.2.tar.gz", hash = "sha256:7fda9eb655c9c230dab534f1983763de5835249750e85fbcef43aaa30a9a2414", size = 22330, upload-time = "2025-11-03T09:25:26.604Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/88/39/799be3f2f0f38cc727ee3b4f1445fe6d5e4133064ec2e4115069418a5bb6/cloudpickle-3.1.2-py3-none-any.whl", hash = "sha256:9acb47f6afd73f60dc1df93bb801b472f05ff42fa6c84167d25cb206be1fbf4a", size = 22228, upload-time = "2025-11-03T09:25:25.534Z" }, -] - -[[package]] -name = "colorama" -version = "0.4.6" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, -] - -[[package]] -name = "contourpy" -version = "1.4.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/83/5a/a55177dd22553a277388e8a1b3220e92de91bacb28356cdc73caa240121d/contourpy-1.4.0.tar.gz", hash = "sha256:20156f5a1ac4f8ce02656e39a61e82164a3d359796dc8026f75b062783d500e1", size = 13323726, upload-time = "2026-09-11T19:05:05.808Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/75/b0/b548628c8dddf6e5e0a981b7fa8e4c008024df09c7d17c6a9e271edfa0b3/contourpy-1.4.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:a3e67bc1a6a4618a1dac7c3053a9ffece5ddfa2046b2670b342cf430bb0b87d1", size = 297354, upload-time = "2026-09-11T19:03:35.017Z" }, - { url = "https://files.pythonhosted.org/packages/04/5a/513484208742f65af2648c519f0b7ab034616e183c4403538024c136a1ab/contourpy-1.4.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:072a702e2e178f4fcf96f04775d0c059e4c917925abf0d3208ebb6865df3c23d", size = 283729, upload-time = "2026-09-11T19:03:36.818Z" }, - { url = "https://files.pythonhosted.org/packages/35/da/5a6562febc994b2c4bf9d01e57a50458ef8a2051bb3e7b35a00333e0fc26/contourpy-1.4.0-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:404dbcd9233513dfd1323b66ea593fef90e899f658f6b8a9ed8e93bd0ca669db", size = 357819, upload-time = "2026-09-11T19:03:39.058Z" }, - { url = "https://files.pythonhosted.org/packages/6d/99/7b358ce888aaa426f755daa096c99c59d17c352037344e836f20e5d3b957/contourpy-1.4.0-cp314-cp314-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:61624f6722480aa7e168fd746164e7bfdf48f8ccb6542f1613741200dc37e2b1", size = 407795, upload-time = "2026-09-11T19:03:41.293Z" }, - { url = "https://files.pythonhosted.org/packages/be/02/34d7548adf08c60967435d79d7a90882eacc84f5093435287974c7e99943/contourpy-1.4.0-cp314-cp314-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:2e741a39dfc96babe1722561e81351aeec104526c57dbe91c167c1db606c2d55", size = 408552, upload-time = "2026-09-11T19:03:42.728Z" }, - { url = "https://files.pythonhosted.org/packages/cf/8c/9677dad226e5aaac5b823d5b09d881ebc32ff1648af417597649841e516c/contourpy-1.4.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e439ab450c93455feb0218532ded3e946ec7a9b5f459069f122cfa44b89229f6", size = 384850, upload-time = "2026-09-11T19:03:44.402Z" }, - { url = "https://files.pythonhosted.org/packages/08/f5/e3fbeaa489c223629fbb434737cb3de4ecaca9cae4825e77b92490f640f8/contourpy-1.4.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:417be048f7e122cefbe2a34c51a1f2b0410f8eb35796f794d45475ddb7872d9d", size = 1357601, upload-time = "2026-09-11T19:03:46.746Z" }, - { url = "https://files.pythonhosted.org/packages/88/be/55fec293d342c222ed58a5d4056df4362e192af2e59527a3d5d822de8828/contourpy-1.4.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:43d80072a32299bf945de0a6dd4e034ea0162e832261de2e07a274a5adbba9c9", size = 1427340, upload-time = "2026-09-11T19:03:48.745Z" }, - { url = "https://files.pythonhosted.org/packages/86/40/0faabf453edf59b0e3633381d7ca72c3414fb9b72b5cbb08b06ed792d3e5/contourpy-1.4.0-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:fa1b787362a3856e63dd2b89449f6c384d55ee1beb8f94f4bcc871b40f02462c", size = 118355, upload-time = "2026-09-11T19:03:50.047Z" }, - { url = "https://files.pythonhosted.org/packages/38/0c/7ea75aa3559ae8a6805015e15a19c92b6480752ecc92526361e5776e88c5/contourpy-1.4.0-cp314-cp314-win32.whl", hash = "sha256:738c44fa71735a617f36da58e32810512407d283d5d1bb5b1cecb00d7eeba7cf", size = 356977, upload-time = "2026-09-11T19:04:11.401Z" }, - { url = "https://files.pythonhosted.org/packages/15/41/3df8cc14bb572c8b447f9dfbc876703f0f2f2896058285b5c87a1d4a638d/contourpy-1.4.0-cp314-cp314-win_amd64.whl", hash = "sha256:181bea01bc742734ae672fa00c717d855dd35e1f029536406538c52e7b0cd73d", size = 240101, upload-time = "2026-09-11T19:04:12.772Z" }, - { url = "https://files.pythonhosted.org/packages/66/01/ee6830a5aa4565662a345172f6f35eeeed625d20f1c9d4ecd35318fe4ade/contourpy-1.4.0-cp314-cp314-win_arm64.whl", hash = "sha256:64039341e2d8804f1a13bda8c69083eab935167bbd7e856cf096241f17d5bd45", size = 576359, upload-time = "2026-09-11T19:04:14.983Z" }, - { url = "https://files.pythonhosted.org/packages/51/af/49ee9c9cf012699e505c1b63c531cc99f19df50f205173092c4d541e3ec7/contourpy-1.4.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:ef47fef9a912c77e3d84b014f701e5a9340f25703e9ed3bcebe37f91fb69dc49", size = 313601, upload-time = "2026-09-11T19:03:51.452Z" }, - { url = "https://files.pythonhosted.org/packages/ad/59/ce411ab2be0f805038626d8f4ca9ce563487ff51ff6c363e7512271564ed/contourpy-1.4.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:96019f059bcb3774acc0104be3360a57b36d33eaaa2e1d7f77c800ca8e07618b", size = 302812, upload-time = "2026-09-11T19:03:53.022Z" }, - { url = "https://files.pythonhosted.org/packages/95/b0/390915f9af14c1e2d3b2a9ad99d85b98df4352346a51e368cbb825f813c3/contourpy-1.4.0-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ddf5a2596d716fd3793434844caf31abc7a40b1e8718431420c89858268fb909", size = 348996, upload-time = "2026-09-11T19:03:54.594Z" }, - { url = "https://files.pythonhosted.org/packages/c7/85/e0952576d54322f3e9983b1e049a7a31f42ce82c5c29d01821e6dd3780d7/contourpy-1.4.0-cp314-cp314t-manylinux_2_26_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:dd59df9e2fb0aff8bd7fd1adb9f349b7c835245a9d0f18c2f2deeb537190f2b1", size = 400272, upload-time = "2026-09-11T19:03:56.822Z" }, - { url = "https://files.pythonhosted.org/packages/d6/7d/f13543a5e1598e4a4ad01d9d5e2d65f9c09c37660616429913d606c5977d/contourpy-1.4.0-cp314-cp314t-manylinux_2_26_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b7794ea07cab575633daad8d6e963b662053391022294d3aacc1d9ba9ef54114", size = 412895, upload-time = "2026-09-11T19:03:58.487Z" }, - { url = "https://files.pythonhosted.org/packages/c5/a8/6e5eaff53507dfd2dfc428240eecda32453067277d4b3e7ca71210b2d611/contourpy-1.4.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a1c8a74744fa746eeaa12f0f9ed7f8b4add9d8f1b14d95e3b5e133675eac888f", size = 377060, upload-time = "2026-09-11T19:04:00.122Z" }, - { url = "https://files.pythonhosted.org/packages/16/a3/d00e44d35511b3cead5bd794ee672d4eecf30d9de4ceac57a49d6204064f/contourpy-1.4.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:107ec46f7aa1664266d69b1181aadc953d58d39f7d8d648fa5b795d4a060b0de", size = 1345688, upload-time = "2026-09-11T19:04:02.705Z" }, - { url = "https://files.pythonhosted.org/packages/1e/01/1e533b0cae32c8edd63fe40435ec3b29345dab87321afcd28722d7f957d1/contourpy-1.4.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5a4c89c7f38a0d7a94356d74e3391073c4325c4d47ad2fd065c3fe7dbae16c0c", size = 1416140, upload-time = "2026-09-11T19:04:04.77Z" }, - { url = "https://files.pythonhosted.org/packages/6d/0c/4ee66759f00d6a5d6909782f14f441d639fa18a4a43142e334351efc47f8/contourpy-1.4.0-cp314-cp314t-win32.whl", hash = "sha256:912c6afaa106e2f74ba22b30416b77ef7eb94e3bdc5a4df51ab147845dba9b6d", size = 371982, upload-time = "2026-09-11T19:04:06.233Z" }, - { url = "https://files.pythonhosted.org/packages/e6/ac/b3e5324138c3e3741f526663f4f265f4f69dfdf26415a39428940864a5f6/contourpy-1.4.0-cp314-cp314t-win_amd64.whl", hash = "sha256:ef9440f6f8506246269a82734f5ff9e2e4c5e775b3996fc883cc491c5257eca6", size = 262907, upload-time = "2026-09-11T19:04:07.882Z" }, - { url = "https://files.pythonhosted.org/packages/19/37/c9aa45e47819dc15a38fc5c81a2fb987fde55e9d3b991fbde514e3b6b5f5/contourpy-1.4.0-cp314-cp314t-win_arm64.whl", hash = "sha256:fc9feef8f1f001c5b87decadc67c4a5d1eebb62ca39c4763d1237ff62cf2b707", size = 587071, upload-time = "2026-09-11T19:04:09.898Z" }, -] - -[[package]] -name = "cramjam" -version = "2.11.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/14/12/34bf6e840a79130dfd0da7badfb6f7810b8fcfd60e75b0539372667b41b6/cramjam-2.11.0.tar.gz", hash = "sha256:5c82500ed91605c2d9781380b378397012e25127e89d64f460fea6aeac4389b4", size = 99100, upload-time = "2025-07-27T21:25:07.559Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/de/07/a1051cdbbe6d723df16d756b97f09da7c1adb69e29695c58f0392bc12515/cramjam-2.11.0-cp314-cp314-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:7ba5e38c9fbd06f086f4a5a64a1a5b7b417cd3f8fc07a20e5c03651f72f36100", size = 3554141, upload-time = "2025-07-27T21:23:17.938Z" }, - { url = "https://files.pythonhosted.org/packages/74/66/58487d2e16ef3d04f51a7c7f0e69823e806744b4c21101e89da4873074bc/cramjam-2.11.0-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:b8adeee57b41fe08e4520698a4b0bd3cc76dbd81f99424b806d70a5256a391d3", size = 1860353, upload-time = "2025-07-27T21:23:19.593Z" }, - { url = "https://files.pythonhosted.org/packages/67/b4/67f6254d166ffbcc9d5fa1b56876eaa920c32ebc8e9d3d525b27296b693b/cramjam-2.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:b96a74fa03a636c8a7d76f700d50e9a8bc17a516d6a72d28711225d641e30968", size = 1693832, upload-time = "2025-07-27T21:23:21.185Z" }, - { url = "https://files.pythonhosted.org/packages/55/a3/4e0b31c0d454ae70c04684ed7c13d3c67b4c31790c278c1e788cb804fa4a/cramjam-2.11.0-cp314-cp314-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:c3811a56fa32e00b377ef79121c0193311fd7501f0fb378f254c7f083cc1fbe0", size = 2027080, upload-time = "2025-07-27T21:23:23.303Z" }, - { url = "https://files.pythonhosted.org/packages/d9/c7/5e8eed361d1d3b8be14f38a54852c5370cc0ceb2c2d543b8ba590c34f080/cramjam-2.11.0-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c5d927e87461f8a0d448e4ab5eb2bca9f31ca5d8ea86d70c6f470bb5bc666d7e", size = 1761543, upload-time = "2025-07-27T21:23:24.991Z" }, - { url = "https://files.pythonhosted.org/packages/09/0c/06b7f8b0ce9fde89470505116a01fc0b6cb92d406c4fb1e46f168b5d3fa5/cramjam-2.11.0-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:f1f5c450121430fd89cb5767e0a9728ecc65997768fd4027d069cb0368af62f9", size = 1854636, upload-time = "2025-07-27T21:23:26.987Z" }, - { url = "https://files.pythonhosted.org/packages/6f/c6/6ebc02c9d5acdf4e5f2b1ec6e1252bd5feee25762246798ae823b3347457/cramjam-2.11.0-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:724aa7490be50235d97f07e2ca10067927c5d7f336b786ddbc868470e822aa25", size = 2032715, upload-time = "2025-07-27T21:23:28.603Z" }, - { url = "https://files.pythonhosted.org/packages/a2/77/a122971c23f5ca4b53e4322c647ac7554626c95978f92d19419315dddd05/cramjam-2.11.0-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:54c4637122e7cfd7aac5c1d3d4c02364f446d6923ea34cf9d0e8816d6e7a4936", size = 2069039, upload-time = "2025-07-27T21:23:30.319Z" }, - { url = "https://files.pythonhosted.org/packages/19/0f/f6121b90b86b9093c066889274d26a1de3f29969d45c2ed1ecbe2033cb78/cramjam-2.11.0-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:17eb39b1696179fb471eea2de958fa21f40a2cd8bf6b40d428312d5541e19dc4", size = 1979566, upload-time = "2025-07-27T21:23:32.002Z" }, - { url = "https://files.pythonhosted.org/packages/e0/a3/f95bc57fd7f4166ce6da816cfa917fb7df4bb80e669eb459d85586498414/cramjam-2.11.0-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:36aa5a798aa34e11813a80425a30d8e052d8de4a28f27bfc0368cfc454d1b403", size = 2030905, upload-time = "2025-07-27T21:23:33.696Z" }, - { url = "https://files.pythonhosted.org/packages/fc/52/e429de4e8bc86ee65e090dae0f87f45abd271742c63fb2d03c522ffde28a/cramjam-2.11.0-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:449fca52774dc0199545fbf11f5128933e5a6833946707885cf7be8018017839", size = 2155592, upload-time = "2025-07-27T21:23:35.375Z" }, - { url = "https://files.pythonhosted.org/packages/6c/6c/65a7a0207787ad39ad804af4da7f06a60149de19481d73d270b540657234/cramjam-2.11.0-cp314-cp314-musllinux_1_1_i686.whl", hash = "sha256:d87d37b3d476f4f7623c56a232045d25bd9b988314702ea01bd9b4a94948a778", size = 2170839, upload-time = "2025-07-27T21:23:37.197Z" }, - { url = "https://files.pythonhosted.org/packages/b2/c5/5c5db505ba692bc844246b066e23901d5905a32baf2f33719c620e65887f/cramjam-2.11.0-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:26cb45c47d71982d76282e303931c6dd4baee1753e5d48f9a89b3a63e690b3a3", size = 2157236, upload-time = "2025-07-27T21:23:38.854Z" }, - { url = "https://files.pythonhosted.org/packages/b0/22/88e6693e60afe98901e5bbe91b8dea193e3aa7f42e2770f9c3339f5c1065/cramjam-2.11.0-cp314-cp314-win32.whl", hash = "sha256:4efe919d443c2fd112fe25fe636a52f9628250c9a50d9bddb0488d8a6c09acc6", size = 1604136, upload-time = "2025-07-27T21:23:40.56Z" }, - { url = "https://files.pythonhosted.org/packages/cc/f8/01618801cd59ccedcc99f0f96d20be67d8cfc3497da9ccaaad6b481781dd/cramjam-2.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:ccec3524ea41b9abd5600e3e27001fd774199dbb4f7b9cb248fcee37d4bda84c", size = 1710272, upload-time = "2025-07-27T21:23:42.236Z" }, - { url = "https://files.pythonhosted.org/packages/40/81/6cdb3ed222d13ae86bda77aafe8d50566e81a1169d49ed195b6263610704/cramjam-2.11.0-cp314-cp314t-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:966ac9358b23d21ecd895c418c048e806fd254e46d09b1ff0cdad2eba195ea3e", size = 3559671, upload-time = "2025-07-27T21:23:44.504Z" }, - { url = "https://files.pythonhosted.org/packages/cb/43/52b7e54fe5ba1ef0270d9fdc43dabd7971f70ea2d7179be918c997820247/cramjam-2.11.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:387f09d647a0d38dcb4539f8a14281f8eb6bb1d3e023471eb18a5974b2121c86", size = 1867876, upload-time = "2025-07-27T21:23:46.987Z" }, - { url = "https://files.pythonhosted.org/packages/9d/28/30d5b8d10acd30db3193bc562a313bff722888eaa45cfe32aa09389f2b24/cramjam-2.11.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:665b0d8fbbb1a7f300265b43926457ec78385200133e41fef19d85790fc1e800", size = 1695562, upload-time = "2025-07-27T21:23:48.644Z" }, - { url = "https://files.pythonhosted.org/packages/d9/86/ec806f986e01b896a650655024ea52a13e25c3ac8a3a382f493089483cdc/cramjam-2.11.0-cp314-cp314t-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:ca905387c7a371531b9622d93471be4d745ef715f2890c3702479cd4fc85aa51", size = 2025056, upload-time = "2025-07-27T21:23:50.404Z" }, - { url = "https://files.pythonhosted.org/packages/09/43/c2c17586b90848d29d63181f7d14b8bd3a7d00975ad46e3edf2af8af7e1f/cramjam-2.11.0-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3c1aa56aef2c8af55a21ed39040a94a12b53fb23beea290f94d19a76027e2ffb", size = 1764084, upload-time = "2025-07-27T21:23:52.265Z" }, - { url = "https://files.pythonhosted.org/packages/2b/a9/68bc334fadb434a61df10071dc8606702aa4f5b6cdb2df62474fc21d2845/cramjam-2.11.0-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e5db59c1cdfaa2ab85cc988e602d6919495f735ca8a5fd7603608eb1e23c26d5", size = 1854859, upload-time = "2025-07-27T21:23:54.085Z" }, - { url = "https://files.pythonhosted.org/packages/5b/4e/b48e67835b5811ec5e9cb2e2bcba9c3fd76dab3e732569fe801b542c6ca9/cramjam-2.11.0-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b1f893014f00fe5e89a660a032e813bf9f6d91de74cd1490cdb13b2b59d0c9a3", size = 2035970, upload-time = "2025-07-27T21:23:55.758Z" }, - { url = "https://files.pythonhosted.org/packages/c4/70/d2ac33d572b4d90f7f0f2c8a1d60fb48f06b128fdc2c05f9b49891bb0279/cramjam-2.11.0-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:c26a1eb487947010f5de24943bd7c422dad955b2b0f8650762539778c380ca89", size = 2069320, upload-time = "2025-07-27T21:23:57.494Z" }, - { url = "https://files.pythonhosted.org/packages/1d/4c/85cec77af4a74308ba5fca8e296c4e2f80ec465c537afc7ab1e0ca2f9a00/cramjam-2.11.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7d5c8bfb438d94e7b892d1426da5fc4b4a5370cc360df9b8d9d77c33b896c37e", size = 1982668, upload-time = "2025-07-27T21:23:59.126Z" }, - { url = "https://files.pythonhosted.org/packages/55/45/938546d1629e008cc3138df7c424ef892719b1796ff408a2ab8550032e5e/cramjam-2.11.0-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:cb1fb8c9337ab0da25a01c05d69a0463209c347f16512ac43be5986f3d1ebaf4", size = 2034028, upload-time = "2025-07-27T21:24:00.865Z" }, - { url = "https://files.pythonhosted.org/packages/01/76/b5a53e20505555f1640e66dcf70394bcf51a1a3a072aa18ea35135a0f9ed/cramjam-2.11.0-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:1f6449f6de52dde3e2f1038284910c8765a397a25e2d05083870f3f5e7fc682c", size = 2155513, upload-time = "2025-07-27T21:24:02.92Z" }, - { url = "https://files.pythonhosted.org/packages/84/12/8d3f6ceefae81bbe45a347fdfa2219d9f3ac75ebc304f92cd5fcb4fbddc5/cramjam-2.11.0-cp314-cp314t-musllinux_1_1_i686.whl", hash = "sha256:382dec4f996be48ed9c6958d4e30c2b89435d7c2c4dbf32480b3b8886293dd65", size = 2170035, upload-time = "2025-07-27T21:24:04.558Z" }, - { url = "https://files.pythonhosted.org/packages/4b/85/3be6f0a1398f976070672be64f61895f8839857618a2d8cc0d3ab529d3dc/cramjam-2.11.0-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:d388bd5723732c3afe1dd1d181e4213cc4e1be210b080572e7d5749f6e955656", size = 2160229, upload-time = "2025-07-27T21:24:06.729Z" }, - { url = "https://files.pythonhosted.org/packages/57/5e/66cfc3635511b20014bbb3f2ecf0095efb3049e9e96a4a9e478e4f3d7b78/cramjam-2.11.0-cp314-cp314t-win32.whl", hash = "sha256:0a70ff17f8e1d13f322df616505550f0f4c39eda62290acb56f069d4857037c8", size = 1610267, upload-time = "2025-07-27T21:24:08.428Z" }, - { url = "https://files.pythonhosted.org/packages/ce/c6/c71e82e041c95ffe6a92ac707785500aa2a515a4339c2c7dd67e3c449249/cramjam-2.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:028400d699442d40dbda02f74158c73d05cb76587a12490d0bfedd958fd49188", size = 1713108, upload-time = "2025-07-27T21:24:10.147Z" }, -] - -[[package]] -name = "cuda-bindings" -version = "13.4.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/e6/f0/d71d1dd45caa358a39c6898dcc9f1bd22a9f22b92376b967835c047b50bc/cuda_bindings-13.4.4-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a429132f06d0c3a853aa419aabe4c0a7bfc3f3886fb4c336d81729f63754d31", size = 6407946, upload-time = "2026-10-14T12:35:39.826Z" }, - { url = "https://files.pythonhosted.org/packages/15/65/0b8f2f2509e5f1a8e556a2e083ac0287fe4af79afdba3153454229306f39/cuda_bindings-13.4.4-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:680e70dcb8b61fdae16d3bb4876c41d55aadc9fe8b3825ea567f2f2c85606b0a", size = 6978954, upload-time = "2026-10-14T12:35:42.892Z" }, - { url = "https://files.pythonhosted.org/packages/f1/dd/81903887199ffd90f9ee29eee8ecd497b734007e78b80689e0dd2f904bd8/cuda_bindings-13.4.4-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:03a8fd2a25530dbae26ecd4462d418156028e35b3f3b081440ff6368506e10eb", size = 6314269, upload-time = "2026-10-14T12:35:50.607Z" }, - { url = "https://files.pythonhosted.org/packages/90/44/f6325a6061f1fc619e216dd6c7cbd9a0b475d3937bb32296c84cc27e6eeb/cuda_bindings-13.4.4-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0ce9208c7cf03a947ac11553b08ee2fca594ac63d34360639a9060980f546a62", size = 6865978, upload-time = "2026-10-14T12:35:52.874Z" }, -] - -[[package]] -name = "cuda-pathfinder" -version = "1.8.3" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b9/fb/f8e1890428f9f590b4beebd63b068aac1ce32a3331510c847b9f9a78f261/cuda_pathfinder-1.8.3-py3-none-any.whl", hash = "sha256:e29e59829c297a7a5233bd9cc71094fc5bddbd076951482670178f9eade39b1f", size = 62561, upload-time = "2026-10-02T03:20:23.712Z" }, -] - -[[package]] -name = "cuda-toolkit" -version = "13.0.3.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/c7/a79086a62c98befcdb8349656c6f114e2db3b8b2422f6e25c97a7f2a9a3c/cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f", size = 2512, upload-time = "2026-04-14T00:50:08.173Z" }, -] - -[package.optional-dependencies] -cublas = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cudart = [ - { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cufft = [ - { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cufile = [ - { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cupti = [ - { name = "nvidia-cuda-cupti", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -curand = [ - { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cusolver = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cusparse = [ - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -nvtx = [ - { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - -[[package]] -name = "cycler" -version = "0.12.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, -] - -[[package]] -name = "einops" -version = "0.8.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/2c/77/850bef8d72ffb9219f0b1aac23fbc1bf7d038ee6ea666f331fa273031aa2/einops-0.8.2.tar.gz", hash = "sha256:609da665570e5e265e27283aab09e7f279ade90c4f01bcfca111f3d3e13f2827", size = 56261, upload-time = "2026-01-26T04:13:17.638Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/09/f8d8f8f31e4483c10a906437b4ce31bdf3d6d417b73fe33f1a8b59e34228/einops-0.8.2-py3-none-any.whl", hash = "sha256:54058201ac7087911181bfec4af6091bb59380360f069276601256a76af08193", size = 65638, upload-time = "2026-01-26T04:13:18.546Z" }, -] - -[[package]] -name = "faiss-cpu" -version = "1.12.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, - { name = "packaging" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/7d/80/bb75a7ed6e824dea452a24d3434a72ed799324a688b10b047d441d270185/faiss_cpu-1.12.0.tar.gz", hash = "sha256:2f87cbcd603f3ed464ebceb857971fdebc318de938566c9ae2b82beda8e953c0", size = 69292, upload-time = "2025-08-13T06:07:26.553Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/12/35/01a4a7c179d67bee0d8a027b95c3eae19cb354ae69ef2bc50ac3b93bc853/faiss_cpu-1.12.0-cp314-cp314-macosx_13_0_x86_64.whl", hash = "sha256:ff7db774968210d08cd0331287f3f66a6ffef955a7aa9a7fcd3eb4432a4ce5f5", size = 8036142, upload-time = "2025-08-13T06:06:38.894Z" }, - { url = "https://files.pythonhosted.org/packages/08/23/bac2859490096608c9d527f3041b44c2e43f8df0d4aadd53a4cc5ce678ac/faiss_cpu-1.12.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:220b5bb5439c64e417b35f9ade4c7dc3bf7df683d6123901ba84d6d764ecd486", size = 3363747, upload-time = "2025-08-13T06:06:40.73Z" }, - { url = "https://files.pythonhosted.org/packages/7b/1d/e18023e1f43a18ec593adcd69d356f1fa94bde20344e38334d5985e5c5cc/faiss_cpu-1.12.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:693d0bf16f79e8d16a1baaeda459f3375f37da0354e97dc032806b48a2a54151", size = 3835232, upload-time = "2025-08-13T06:06:42.172Z" }, - { url = "https://files.pythonhosted.org/packages/cd/2b/1c1fea423d3f550f44c5ec3f14d8400919b49c285c3bd146687c63e40186/faiss_cpu-1.12.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bcc6587dee21e17430fb49ddc5200625d6f5e1de2bdf436f14827bad4ca78d19", size = 31432677, upload-time = "2025-08-13T06:06:44.348Z" }, - { url = "https://files.pythonhosted.org/packages/de/d2/3483e92a02f30e2d8491a256f470f54b7f5483266dfe09126d28741d31ec/faiss_cpu-1.12.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b80e5965f001822cc99ec65c715169af1b70bdae72eccd573520a2dec485b3ee", size = 9765504, upload-time = "2025-08-13T06:06:46.567Z" }, - { url = "https://files.pythonhosted.org/packages/ce/2f/d97792211a9bd84b8d6b1dcaa1dcd69ac11e026c6ef19c641b6a87e31025/faiss_cpu-1.12.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:98279f1b4876ef9902695a329b81a99002782ab6e26def472022009df6f1ac68", size = 24169930, upload-time = "2025-08-13T06:06:48.916Z" }, - { url = "https://files.pythonhosted.org/packages/ee/b8/b707ca4d88af472509a053c39d3cced53efd19d096b8dff2fadc18c4b82d/faiss_cpu-1.12.0-cp314-cp314-win_amd64.whl", hash = "sha256:11670337f9f5ee9ff3490e30683eea80add060c300cf6f6cb0e8faf3155fd20e", size = 18475400, upload-time = "2025-08-13T06:06:51.233Z" }, - { url = "https://files.pythonhosted.org/packages/77/11/42e41ddebde4dfe77e36e92d0110b4f733c8640883abffde54f802482deb/faiss_cpu-1.12.0-cp314-cp314-win_arm64.whl", hash = "sha256:7ac1c8b53609b5c722ab60f1749260a7cb3c72fdfb720a0e3033067e73591da5", size = 8281229, upload-time = "2025-08-13T06:06:53.735Z" }, - { url = "https://files.pythonhosted.org/packages/1c/9a/8ae5bbeabe70eb673c37fc7c77e2e476746331afb6654b2df97d8b6d380d/faiss_cpu-1.12.0-cp314-cp314t-macosx_13_0_x86_64.whl", hash = "sha256:110b21b7bb4c93c4f1a5eb2ffb8ef99dcdb4725f8ab2e5cd161324e4d981f204", size = 8087247, upload-time = "2025-08-13T06:06:55.407Z" }, - { url = "https://files.pythonhosted.org/packages/f4/df/b3d79098860b67b126da351788c04ac243c29718dadc4a678a6f5e7209c0/faiss_cpu-1.12.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:82eb5515ce72be9a43f4cf74447a0d090e014231981df91aff7251204b506fbf", size = 3411043, upload-time = "2025-08-13T06:06:56.983Z" }, - { url = "https://files.pythonhosted.org/packages/bc/2f/b1a2a03dd3cce22ff9fc434aa3c7390125087260c1d1349311da36eaa432/faiss_cpu-1.12.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:754eef89cdf2b35643df6b0923a5a098bdfecf63b5f4bd86c385042ee511b287", size = 3801789, upload-time = "2025-08-13T06:06:58.688Z" }, - { url = "https://files.pythonhosted.org/packages/a3/a8/16ad0c6a966e93d04bfd5248d2be1d8b5849842b0e2611c5ecd26fcaf036/faiss_cpu-1.12.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7285c71c8f5e9c58b55175f5f74c78c518c52c421a88a430263f34e3e31f719c", size = 31231388, upload-time = "2025-08-13T06:07:00.55Z" }, - { url = "https://files.pythonhosted.org/packages/62/a1/9c16eca0b8f8b13c32c47a5e4ff7a4bc0ca3e7d263140312088811230871/faiss_cpu-1.12.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:84a50d7a2f711f79cc8b65aa28956dba6435e47b71a38b2daea44c94c9b8e458", size = 9737605, upload-time = "2025-08-13T06:07:03.018Z" }, - { url = "https://files.pythonhosted.org/packages/a8/4a/2c2d615078c9d816a836fb893aaef551ad152f2eb00bc258698273c240c0/faiss_cpu-1.12.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:7f3e0a14e4edec6a3959a9f51afccb89e863138f184ff2cc24c13f9ad788740b", size = 23922880, upload-time = "2025-08-13T06:07:05.099Z" }, - { url = "https://files.pythonhosted.org/packages/30/aa/99b8402a4dac678794f13f8f4f29d666c2ef0a91594418147f47034ebc81/faiss_cpu-1.12.0-cp314-cp314t-win_amd64.whl", hash = "sha256:8b3239cc371df6826ac43c62ac04eec7cc497bedb43f681fcd8ea494f520ddbb", size = 18750661, upload-time = "2025-08-13T06:07:07.551Z" }, - { url = "https://files.pythonhosted.org/packages/a3/a2/b546e9a20ba157eb2fbe141289f1752f157ee6d932899f4853df4ded6d4b/faiss_cpu-1.12.0-cp314-cp314t-win_arm64.whl", hash = "sha256:58b23456db725ee1bd605a6135d2ef55b2ac3e0b6fe873fd99a909e8ef4bd0ff", size = 8302032, upload-time = "2025-08-13T06:07:09.602Z" }, -] - -[[package]] -name = "filelock" -version = "3.32.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7d/64/a02e6765de08964ed371eca577870593245afc9dfac16d037de7c10d18e6/filelock-3.32.3.tar.gz", hash = "sha256:0ffa185a3540854c95caa7fa76b76cb219d907415e2c5dc9af25fd970563487f", size = 218135, upload-time = "2026-08-13T16:00:05.577Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a7/8e/50f46a9c0ce8d2861a394c1347caae037ea0431d2f67d7feb151cbc4649a/filelock-3.32.3-py3-none-any.whl", hash = "sha256:7f0ca4bcc0e181c60dbbd8aa9ab5b120ebb99e4e064e83636340056f833a1f09", size = 98901, upload-time = "2026-08-13T16:00:03.974Z" }, -] - -[[package]] -name = "flatbuffers" -version = "25.12.19" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, -] - -[[package]] -name = "fonttools" -version = "4.67.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/94/36/102e180f8f5dbaee88b26595b01ca8aa80bf4e62128d9aa94265b3996c96/fonttools-4.67.0.tar.gz", hash = "sha256:3cb57e6600ca77c0b1729cf8adc23bc0652633a37f18cfa934d9c7bc3de25519", size = 3750028, upload-time = "2026-10-14T13:20:28.294Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/25/5b/c418f48918e40ef8c3f0f555567fe013c0c8058a8afa8040d6baeec80683/fonttools-4.67.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:846982e89b1861d6c9d7fcd6567aec3fa5a10ad313e7f2076045fcd339cfbd8e", size = 3110155, upload-time = "2026-10-14T13:19:14.877Z" }, - { url = "https://files.pythonhosted.org/packages/30/18/49013c643c3d56fce1b7e909ef7c01c36a5bd906dfb58571c9dcdaa4dc38/fonttools-4.67.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:952eb091689545d86d16e40f719ed7bb086dd810a07dcc9ea2ca0a81004810a3", size = 2599657, upload-time = "2026-10-14T13:19:16.93Z" }, - { url = "https://files.pythonhosted.org/packages/1f/2c/b7f33fa3bd1e4afdf9bf93b760f22486350eda487ce76c47f5931f868957/fonttools-4.67.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e2b5d511ea012dce7bd6df12b279b7d7a5b01b019865717d03ae679f4b944fa5", size = 5385200, upload-time = "2026-10-14T13:19:18.868Z" }, - { url = "https://files.pythonhosted.org/packages/79/fe/fef04b2cc2930edba11095f9e9b5c2797f8594fc54316195cc39d3c3bc63/fonttools-4.67.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:916836845e4b1c1447bb61390ffb3cb5f2940fd9f5d6de4685539a81806c7764", size = 5328754, upload-time = "2026-10-14T13:19:21.179Z" }, - { url = "https://files.pythonhosted.org/packages/2e/c6/41cd4f6137f61dd059cc0609b73d9556091ecfcc8cb4d3cc543129c8ec24/fonttools-4.67.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:775364ac079e2ea7a2eedb5f9172c57b059d638ff79e2bf8d4257e5805713f32", size = 5327159, upload-time = "2026-10-14T13:19:23.153Z" }, - { url = "https://files.pythonhosted.org/packages/53/5c/08abd0a6d5c36624411e1b934745b4689d4309b03e98d8cf49f9469c63b6/fonttools-4.67.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:b3ddf350e74508102b33dc6b32984b6dd751359a7c57732bcd39f9d7cb37d71e", size = 5460332, upload-time = "2026-10-14T13:19:25.454Z" }, - { url = "https://files.pythonhosted.org/packages/b5/0f/59e835023817fe3932653067fde74960a0800fb95535375d8206aa9ecd68/fonttools-4.67.0-cp314-cp314-win32.whl", hash = "sha256:72d6d316dffc92eadb771f697f289ea7b60f689580931328905a267bd170f93b", size = 2448714, upload-time = "2026-10-14T13:19:27.73Z" }, - { url = "https://files.pythonhosted.org/packages/b3/d3/5230265a5ff16aead01ce1a432a6b5bbdabe086f433988f41a1395e6dff8/fonttools-4.67.0-cp314-cp314-win_amd64.whl", hash = "sha256:4e2c1586b5b6588a47d02e2588170eefdc996b708f2659c44dbe169bd6fcacb5", size = 2500790, upload-time = "2026-10-14T13:19:29.906Z" }, - { url = "https://files.pythonhosted.org/packages/b3/38/d899d7bbbe04d27dd509ac6b8f58f73fc240bb1dfe0ada9a9d33ad3bf9f2/fonttools-4.67.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:84a3aed005de106fb1794372dace82eca50859d52ae26da4bb6c602480a41250", size = 3183808, upload-time = "2026-10-14T13:19:32.015Z" }, - { url = "https://files.pythonhosted.org/packages/c3/f6/4f465a62972e383b3d82205841b93f625a4e5ece6e5693c5be2a691ffe6d/fonttools-4.67.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:64e56d0d6a39780fee86955c758674538387b18f911ea904a4aae8f8e30fa26f", size = 2632826, upload-time = "2026-10-14T13:19:33.854Z" }, - { url = "https://files.pythonhosted.org/packages/d7/91/ce1ae8f8baa75feb2320caf6f74d2c228eba210a13b3e0895c0403e5e987/fonttools-4.67.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8c21073cfe7129aaa070d94f575c1e2a880ae4aae1dcffd5352f174b96d27d16", size = 5557857, upload-time = "2026-10-14T13:19:36.086Z" }, - { url = "https://files.pythonhosted.org/packages/fe/1c/495fe0a6bb8625e693c1417e178aeac42a11aa47e79efd7611c7bc5fb81e/fonttools-4.67.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:720bcf27727193b0fe1883c2e036dc88e37047916e977f5c3daf6ee4316e9656", size = 5364450, upload-time = "2026-10-14T13:19:38.5Z" }, - { url = "https://files.pythonhosted.org/packages/19/9c/d9730d3dd32e39583d6db929d0867df02042539bb0ebc3ad3d92a52a6aaf/fonttools-4.67.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:6c19a770a8d273371a37969003c143eaa629ab893c3db028af8b91d04c6f9a6d", size = 5425231, upload-time = "2026-10-14T13:19:40.659Z" }, - { url = "https://files.pythonhosted.org/packages/f0/c6/d41c1163431828b0fa2172e867798e0c4517ac6606e774b9175e048fb666/fonttools-4.67.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:13d7507252c5a5d7941a5fa1be27d335c378ef07983ea2bb24988bf600eadd5e", size = 5460574, upload-time = "2026-10-14T13:19:43.22Z" }, - { url = "https://files.pythonhosted.org/packages/95/af/14885b78b1c1ff7219f890b79a5a6f76608d907c171b40e43839de995f54/fonttools-4.67.0-cp314-cp314t-win32.whl", hash = "sha256:07a2f36b3263faadf5b7b548f62fd3cac401e490189c82b16f7139ac0df91cd4", size = 2482468, upload-time = "2026-10-14T13:19:45.91Z" }, - { url = "https://files.pythonhosted.org/packages/cf/33/3d660eb850d24a81b4097ed46a1352c4ac0e4c10025526fa115e1871fc64/fonttools-4.67.0-cp314-cp314t-win_amd64.whl", hash = "sha256:fd79e36c2968e9fc3e1b082f2ba7dc63ae88a161a3d8ceaa0746b906455f3617", size = 2532326, upload-time = "2026-10-14T13:19:48.023Z" }, - { url = "https://files.pythonhosted.org/packages/3d/61/4161946319472aaa9b897bd18ad5108a5b10f5ebaa503d921a001ac4fff9/fonttools-4.67.0-py3-none-any.whl", hash = "sha256:4304f03ed7f4ba000a8dcc941ad854bfa52e2f3b6112b8f099b6f431cf98e701", size = 1213142, upload-time = "2026-10-14T13:20:26.258Z" }, -] - -[[package]] -name = "fsspec" -version = "2026.7.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88", size = 317040, upload-time = "2026-07-28T16:34:51.052Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279", size = 206583, upload-time = "2026-07-28T16:34:49.538Z" }, -] - -[[package]] -name = "gnn4colliders" -source = { editable = "." } -dependencies = [ - { name = "awkward" }, - { name = "huggingface-hub" }, - { name = "numpy" }, - { name = "pyyaml" }, - { name = "torch", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "uproot" }, -] - -[package.optional-dependencies] -onnx = [ - { name = "onnx" }, - { name = "onnxruntime" }, -] -plots = [ - { name = "matplotlib" }, -] -tabdpt = [ - { name = "tabdpt" }, -] -tabicl = [ - { name = "tabicl" }, -] -tabpfn = [ - { name = "tabpfn" }, -] - -[package.dev-dependencies] -dev = [ - { name = "onnx" }, - { name = "onnxruntime" }, - { name = "pytest" }, - { name = "ruff" }, -] - -[package.metadata] -requires-dist = [ - { name = "awkward", specifier = ">=2.0" }, - { name = "huggingface-hub", specifier = ">=0.36,<1" }, - { name = "matplotlib", marker = "extra == 'plots'", specifier = ">=3.9,<4" }, - { name = "numpy", specifier = ">=2.1" }, - { name = "onnx", marker = "extra == 'onnx'", specifier = ">=1.16,<2" }, - { name = "onnxruntime", marker = "extra == 'onnx'", specifier = ">=1.18,<2" }, - { name = "pyyaml", specifier = ">=6" }, - { name = "tabdpt", marker = "extra == 'tabdpt'", specifier = ">=1.3.1,<2" }, - { name = "tabicl", marker = "extra == 'tabicl'", specifier = ">=2.2,<3" }, - { name = "tabpfn", marker = "extra == 'tabpfn'", specifier = ">=9.2,<10" }, - { name = "torch", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'", specifier = "==2.14.1" }, - { name = "torch", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'", specifier = "==2.14.1", index = "https://download.pytorch.org/whl/cu130" }, - { name = "uproot", specifier = ">=5.0" }, -] -provides-extras = ["onnx", "tabpfn", "tabicl", "tabdpt", "plots"] - -[package.metadata.requires-dev] -dev = [ - { name = "onnx", specifier = ">=1.16,<2" }, - { name = "onnxruntime", specifier = ">=1.18,<2" }, - { name = "pytest", specifier = ">=8" }, - { name = "ruff", specifier = ">=0.6" }, -] - -[[package]] -name = "hf-xet" -version = "1.6.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/1b/ab/522a2ab67f27971a9d48ca666d4fca85ef7d5282d142e31fd087e27b1bbe/hf_xet-1.6.0.tar.gz", hash = "sha256:2e58454a340b3556dfa4972d5451aff4fba8dd42a236600ba1a1d2b1514f0fef", size = 920527, upload-time = "2026-08-03T22:33:13.243Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/41/62/3c062f593bd92ef4e77a0ef39541e3d82a0a1d3947c8a777a02a13a27828/hf_xet-1.6.0-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:70cbb9c896901600128cb9b6f06e132954fbede1db30f31f7c6c63f84cb7c31d", size = 4074584, upload-time = "2026-08-03T22:32:47.364Z" }, - { url = "https://files.pythonhosted.org/packages/bb/1e/c0ad437dd267a8e435bef594acf781bbc3874ff0b6435b4962d03ecf7cc4/hf_xet-1.6.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:23379c2f9ec8696d952b16414a2bae72cad86a52df869b050698ba60f538c675", size = 3867381, upload-time = "2026-08-03T22:32:49.049Z" }, - { url = "https://files.pythonhosted.org/packages/d5/ee/7c0d7b6ab336167531b1c30af2af003f054af4c749becbd7209ae33a77c3/hf_xet-1.6.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f2f7278c05c22fd60cb436cda1269649b3e81db65ecdc8496e5e164aa4143e7b", size = 4453982, upload-time = "2026-08-03T22:32:50.568Z" }, - { url = "https://files.pythonhosted.org/packages/63/06/ad8eab1c9525246650cbaa821caa3cdbaca734ab1a5b8c91bea09cbd8d69/hf_xet-1.6.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:948f15d3a9545cfe5932f6bd8b440f6ae630aee108f14b7bd6c561f7c2dcc522", size = 4249445, upload-time = "2026-08-03T22:32:52.391Z" }, - { url = "https://files.pythonhosted.org/packages/d8/26/1eee8aedb0dafc1ab9717dc9ac602cde33361b232dc06803f1f6ed18b58c/hf_xet-1.6.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5153e6bb103ad49d6ea9f1b2e230db5a2ea32551ad09a706d2f61d7c7c80d80e", size = 4451099, upload-time = "2026-08-03T22:32:54.114Z" }, - { url = "https://files.pythonhosted.org/packages/67/57/0b88af1f194ab6c9c650547d9cc06bfeaab836ae4dcdb331676bfb8be95a/hf_xet-1.6.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:35cec30d75c6f9eb9c16a77cef68e85a103b72e24d4b473714ec9ff06428bab9", size = 4664712, upload-time = "2026-08-03T22:32:55.547Z" }, - { url = "https://files.pythonhosted.org/packages/53/a0/26b717a9d1840e8abf48dcec64b5ed8fbe472671d38ad28d30e147132b33/hf_xet-1.6.0-cp314-cp314t-win_amd64.whl", hash = "sha256:5789835d7c6bc9436962853192082374297fb72d7eff7e7762ec25ceb7e25338", size = 4025906, upload-time = "2026-08-03T22:32:57.391Z" }, - { url = "https://files.pythonhosted.org/packages/49/f6/4a9966633c6fef83af997e2cff68ec1963676d412bdfd096df2a93b8e185/hf_xet-1.6.0-cp314-cp314t-win_arm64.whl", hash = "sha256:75765820ce4700db3750c94acc8fe27c5fae4c9ec000a0dbac3ca082acf97765", size = 3849221, upload-time = "2026-08-03T22:32:59.123Z" }, - { url = "https://files.pythonhosted.org/packages/a2/50/7afa2c9c787405864fc47a0d1bbc02c62e9101947ed43c1f43899fc7d91d/hf_xet-1.6.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:633dc0cd71d32da58ab8c03ad38e2fac452c15c2b0a2866ebf6ededfe0a5061d", size = 4071729, upload-time = "2026-08-03T22:33:00.721Z" }, - { url = "https://files.pythonhosted.org/packages/4b/69/55b8dcf636142ae660fec1869fcac14c4da2e8412e14d6eee1523be77e9f/hf_xet-1.6.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:f0906082d9932ae0c0057fa194041c22b4e2cdb46b2592ef3b91f020d62a081a", size = 3876287, upload-time = "2026-08-03T22:33:02.251Z" }, - { url = "https://files.pythonhosted.org/packages/67/4e/a28359bf1c1ecf11eba22123168c138698f7cb576ac678f5a2e16cd5da08/hf_xet-1.6.0-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d62671bb130879cef0ee4c9ebe47a14af6c66ec53e6d84dc15936e5ffdfac82f", size = 4464663, upload-time = "2026-08-03T22:33:03.802Z" }, - { url = "https://files.pythonhosted.org/packages/9a/69/1f0cbc2fb22ae6082d094f743d1b8945a3f36f6089cb95f42b7ee348cda7/hf_xet-1.6.0-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:0e6e21fa3cdfcdcd76748564bf593870a5e013f47d97cf10aed63aa222cff5b7", size = 4262538, upload-time = "2026-08-03T22:33:05.287Z" }, - { url = "https://files.pythonhosted.org/packages/d1/3a/4f4f2301ade26e404462d3336fa11f7958d914cabbabdd6e03c3c5d5658c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:4fc74352a17015bd0ee90038bc9efe38db894cde45f268b6712b04fce8cd0acb", size = 4460520, upload-time = "2026-08-03T22:33:06.81Z" }, - { url = "https://files.pythonhosted.org/packages/ab/5f/311725e2a905534dfee2dcb5b08414f249147f1f12252bfc2bd24caa075c/hf_xet-1.6.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8fb4f71cba6129110c3374a33f919001ff130488fc23553698e34cc1c2a1198c", size = 4675937, upload-time = "2026-08-03T22:33:08.616Z" }, - { url = "https://files.pythonhosted.org/packages/98/b7/8c59a66d15205024662f1d66968136f13893f96df1ddc5087e2e281fc95f/hf_xet-1.6.0-cp38-abi3-win_amd64.whl", hash = "sha256:fb4fadde1b2b70bf4c0c14a6dccbe7194b1c28947fefd5bbe3fed9d940676c3b", size = 4033128, upload-time = "2026-08-03T22:33:10.171Z" }, - { url = "https://files.pythonhosted.org/packages/73/63/ca511b6f802f28cf3489b280fe77475bcca8de85e81a6299d7916b5b5555/hf_xet-1.6.0-cp38-abi3-win_arm64.whl", hash = "sha256:3dc3e35441ba395006af5aaacc40ef2e603c51ef46c3530b9156185f00935ea3", size = 3859359, upload-time = "2026-08-03T22:33:11.725Z" }, -] - -[[package]] -name = "huggingface-hub" -version = "0.36.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "filelock" }, - { name = "fsspec" }, - { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, - { name = "packaging" }, - { name = "pyyaml" }, - { name = "requests" }, - { name = "tqdm" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/7c/b7/8cb61d2eece5fb05a83271da168186721c450eb74e3c31f7ef3169fa475b/huggingface_hub-0.36.2.tar.gz", hash = "sha256:1934304d2fb224f8afa3b87007d58501acfda9215b334eed53072dd5e815ff7a", size = 649782, upload-time = "2026-02-06T09:24:13.098Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl", hash = "sha256:48f0c8eac16145dfce371e9d2d7772854a4f591bcb56c9cf548accf531d54270", size = 566395, upload-time = "2026-02-06T09:24:11.133Z" }, -] - -[[package]] -name = "idna" -version = "3.19" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/5f/f7/abb373e5757eaec4b922b92f97ec8d6d7e057cf06778247604fbc4e7c3f3/idna-3.19.tar.gz", hash = "sha256:5e0811a4383b21dc5838069f801c4fb62113b7447663d2530d2bd6e77b49bf15", size = 215237, upload-time = "2026-08-18T05:14:24.27Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/57/b0/0e52c878c53f245edd3a11020f20979b3f490f245af532c7cae3027754b5/idna-3.19-py3-none-any.whl", hash = "sha256:815e7be7a7806d54abb586dc943addc79e8b2ee16915059658cbeff4b1b43bf4", size = 68550, upload-time = "2026-08-18T05:14:22.343Z" }, -] - -[[package]] -name = "iniconfig" -version = "2.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, -] - -[[package]] -name = "jinja2" -version = "3.1.6" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "markupsafe" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, -] - -[[package]] -name = "joblib" -version = "1.6.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cloudpickle" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d5/1d/537ab090f302b838943a1b56497dd53059b9a9b46a074936470173a2e207/joblib-1.6.0.tar.gz", hash = "sha256:2ccc96785b12046c08fd6d55839c12857831b54a3c1673ffadd2f04bfc4eda03", size = 327903, upload-time = "2026-08-31T09:39:04.122Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/18/53/84099323c2ec4be98d935f63c033ac4151ee83836ca1050ede3b3aadf155/joblib-1.6.0-py3-none-any.whl", hash = "sha256:3dbbf9f6e4b592a2357b854608e980fe6390d131d7a82f011a377ef2ebef7aba", size = 306115, upload-time = "2026-08-31T09:39:02.298Z" }, -] - -[[package]] -name = "kiwisolver" -version = "1.5.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ba/07/bd78e6a8fae171ea041ef5bba3ed21a003522fa088834b069b1909981f30/kiwisolver-1.5.1.tar.gz", hash = "sha256:f1303ef2eec81262a4b708c3e858afe58d7c75ad91c1c05266eda7673369859a", size = 104395, upload-time = "2026-08-28T10:28:27.153Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/dd/87/2d5dfad0daf17dcc18d98c48ed2332fc3f051cf599e60be6182a30dd4cf1/kiwisolver-1.5.1-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:0324cd2567259b7a095f6cf18a52b0ffc6f3de9e69528ff1bc0e7a37bd43ff1a", size = 62337, upload-time = "2026-08-28T10:26:18.778Z" }, - { url = "https://files.pythonhosted.org/packages/08/c8/83e1624f15d6262b470dbcc80b09979fd4d5b2ea3ddfc6b6e3327e235726/kiwisolver-1.5.1-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:74ea337e0ec3f6f342a36a4f1b5cd94dd9affddcd28ba9aae2905af932ee8c6b", size = 64513, upload-time = "2026-08-28T10:26:19.909Z" }, - { url = "https://files.pythonhosted.org/packages/c9/2d/827ec30eb07f528c08d8459ffb318ae91a56d793ee8acbea8b491f0ff906/kiwisolver-1.5.1-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:ee9df1f0d77b9c6e94f4ac0fec533fbddd5ea3a327807f18d7b069ae019ded80", size = 66287, upload-time = "2026-08-28T10:26:21.078Z" }, - { url = "https://files.pythonhosted.org/packages/53/11/5c43a562529dad8def4b81e5e1877c612a7e0298105a5939b3b409d2079c/kiwisolver-1.5.1-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:fc271a6f0a2126958f4090e5507b9da5848927dae331f8f763bd4aa642b3d2cd", size = 123940, upload-time = "2026-08-28T10:26:22.475Z" }, - { url = "https://files.pythonhosted.org/packages/64/db/9bd6c505c95128c258a55236bfbb3a7a3fb6023f863316b6d7d9f3c69052/kiwisolver-1.5.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:9b3092d8992a1d69b7a59c3e39f35e1b9be327a17f68a7c35fc17329e337d6f2", size = 66493, upload-time = "2026-08-28T10:26:23.743Z" }, - { url = "https://files.pythonhosted.org/packages/e1/f4/b3007a3ed5c9be73f81161140684cf7d9bdb9c4b632f5f484d2a1c713fb9/kiwisolver-1.5.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c2306e8bb53601979fcb3fa09cc65e031876d9ae01eff2fcbcd7a84ef94d5bc1", size = 64720, upload-time = "2026-08-28T10:26:24.95Z" }, - { url = "https://files.pythonhosted.org/packages/a1/13/08188f0cafa3a800403e4ff62b9aad4e7a17f9c4c7e080dc8f18c64794cf/kiwisolver-1.5.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:18a0cfb124546a4c2e6087c5f3029c7f44b37c85b142e0ced71f73a7599ac208", size = 1475867, upload-time = "2026-08-28T10:26:26.393Z" }, - { url = "https://files.pythonhosted.org/packages/8a/3e/053bdc3c9abdb8f2606225eda398adca25c0c91ab90add8222a69db65ee0/kiwisolver-1.5.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:34ec467940442c9943016fb2d4c81d1ba84351eeca2f1a78f8bc87f1ba0d414c", size = 1282865, upload-time = "2026-08-28T10:26:28.118Z" }, - { url = "https://files.pythonhosted.org/packages/5e/64/a44c341b36b610588cc2f1e89b3cae072a3119aa8be578e90987cd640751/kiwisolver-1.5.1-cp314-cp314-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a83ee7107df13abe42a54a6654670eef9bb39425cf2e27f65e0007465e1286ab", size = 1300865, upload-time = "2026-08-28T10:26:30.125Z" }, - { url = "https://files.pythonhosted.org/packages/60/5e/7e7d716dca38c714478b741257a5b4a321d9932b8d851551a136dcaf3984/kiwisolver-1.5.1-cp314-cp314-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bebb89489b279b2f5661bbbb2abcc87bcd4a46607bb4a5c966f04f1db6b8df9a", size = 1348071, upload-time = "2026-08-28T10:26:31.829Z" }, - { url = "https://files.pythonhosted.org/packages/10/1a/2b98fdda8bf45b7be317e48ed12393d44334394d315c74b81f4a14c0e31b/kiwisolver-1.5.1-cp314-cp314-manylinux_2_39_riscv64.whl", hash = "sha256:509735237ae0d849e8a843551d423d2500d2e0a9ac1611a145658b29c0fb9f85", size = 992191, upload-time = "2026-08-28T10:26:33.544Z" }, - { url = "https://files.pythonhosted.org/packages/b6/55/d893f5ede0e50f9e3fcf01f6015f42ec7d9cd221e26772701fe4a98745f9/kiwisolver-1.5.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:824c3d763a05ea9e9003610145186b0e9848c7584a5575c79bac5a8e7cd80bad", size = 2233854, upload-time = "2026-08-28T10:26:35.282Z" }, - { url = "https://files.pythonhosted.org/packages/b1/82/f85f6279555a6ee1639fef7bfe83adb037a03e11a6fc9eaa54b8d0380339/kiwisolver-1.5.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:1fff05e239575b1481b6ed1a782f6fad616efbf1f0b1f44e6e85c4dfe426e483", size = 2330621, upload-time = "2026-08-28T10:26:36.9Z" }, - { url = "https://files.pythonhosted.org/packages/56/31/e11aea078f66fc2fffcc179d38ca90d9da97652a241b64519169742ba46a/kiwisolver-1.5.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0627b9bceb9c3cdcf12b8a18655eedfed2692b038df27423383c120d0b7dc2d6", size = 1982848, upload-time = "2026-08-28T10:26:39.01Z" }, - { url = "https://files.pythonhosted.org/packages/af/ea/2956b63bf5140ca46aa2c2818e6aa03e2d5754dd2fa41db1c6b28922940c/kiwisolver-1.5.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:8a708a47ade1fe19e8371d5da076bac0dd4b0a5a7985ad6c637f7f7e361b6baa", size = 2494850, upload-time = "2026-08-28T10:26:40.837Z" }, - { url = "https://files.pythonhosted.org/packages/11/d1/3829542258d8b3fc0898d221e7ef0e2c83eca0d348709bb8dbe54f3d4005/kiwisolver-1.5.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:007a5553dfc4f4e8d184f588a0200e2cd4b63a59cc8796df3c39909e679dc7a0", size = 2298067, upload-time = "2026-08-28T10:26:42.803Z" }, - { url = "https://files.pythonhosted.org/packages/4e/0e/49522e1ab5788cbaf63a26fbd3b851f9028616828c961b8a31b35cb96df8/kiwisolver-1.5.1-cp314-cp314-win_amd64.whl", hash = "sha256:f4167e87b397f273dc2356fcf1eaf50a6bac51e6105f45103ef7129c8efb0255", size = 72282, upload-time = "2026-08-28T10:26:44.268Z" }, - { url = "https://files.pythonhosted.org/packages/f5/b6/22e7ca5315d363e6f81c9f37c9472e12e7b298731e77c0428e6a911a2c39/kiwisolver-1.5.1-cp314-cp314-win_arm64.whl", hash = "sha256:5c490db2168a508088f59140dd392556a54b8bd1048fc6383c8baff13c359673", size = 69855, upload-time = "2026-08-28T10:26:45.725Z" }, - { url = "https://files.pythonhosted.org/packages/30/8c/03a9cfbe871964c8758a816eb03ac96c806da2795a9a7cd9bf9648bfb594/kiwisolver-1.5.1-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:4d4ca09bf13cff792b1884f64b98ee6c2467930d632233be25c56b442d99f10e", size = 126289, upload-time = "2026-08-28T10:26:47.023Z" }, - { url = "https://files.pythonhosted.org/packages/f2/e3/14ce3041ca79dff9c9d884ca00c7bf32374e76028a865a9ecd99b4f5a517/kiwisolver-1.5.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:44b8faef94f1857e77fa0238f3390ff1ac51d2ea20a487e2e452a59fd2b5f5ca", size = 67709, upload-time = "2026-08-28T10:26:48.268Z" }, - { url = "https://files.pythonhosted.org/packages/af/c4/45030471a66ec8ef042e9f96ffe1d522c9ab12da180186a0898966fc1385/kiwisolver-1.5.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2ae70bc59790d2af72a3f76f24b272403e135070340281108b447cb77ea70819", size = 65909, upload-time = "2026-08-28T10:26:49.523Z" }, - { url = "https://files.pythonhosted.org/packages/42/73/17dce073a6ae259bb32cf9d686c4079d2e538868bc45967462bf33df914a/kiwisolver-1.5.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:43844c1a7ad6d723d5b5b4c4fc7f5bd399c40e288120d16257c7c9e8765c6e85", size = 1584907, upload-time = "2026-08-28T10:26:50.933Z" }, - { url = "https://files.pythonhosted.org/packages/8c/84/ae3c75909f507283cbfcc7e916c7e822579ef962020b97e6882b27b4478f/kiwisolver-1.5.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:22d5e5aaad6be121f2515765e3b1c444352cb8eb4c86510801db8f2e50757316", size = 1392474, upload-time = "2026-08-28T10:26:52.638Z" }, - { url = "https://files.pythonhosted.org/packages/34/31/8bcc83caad5bce8fa4577152389848bf6bc110e51e573a2b4e7c2aa34c89/kiwisolver-1.5.1-cp314-cp314t-manylinux_2_24_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:3fa5855898f6d3d01b72ccd48a2d65cbdee301251603fefe34e2025bddba219c", size = 1405246, upload-time = "2026-08-28T10:26:54.248Z" }, - { url = "https://files.pythonhosted.org/packages/55/72/220345537d790cf4ae54f8acfff4b5cc2468e0702a384d651cf7a771c63e/kiwisolver-1.5.1-cp314-cp314t-manylinux_2_24_s390x.manylinux_2_28_s390x.whl", hash = "sha256:d66a64dd5dec136040ec2ae94aa026a912ee60fdd45bc28d3db30037fd809e88", size = 1456099, upload-time = "2026-08-28T10:26:56.042Z" }, - { url = "https://files.pythonhosted.org/packages/cd/10/3725fd2398f66d18c34b4e0f81a8d03764cd4f4f089f58a527f0b4428086/kiwisolver-1.5.1-cp314-cp314t-manylinux_2_39_riscv64.whl", hash = "sha256:9e51c119992ea8820706871c30a4642ec76de20ae82f9b50b9a45517d8e9f810", size = 1073695, upload-time = "2026-08-28T10:26:57.658Z" }, - { url = "https://files.pythonhosted.org/packages/86/cb/28d6e09e66b93e4588b2e6b7d84d020ccefea09e2f4de788510a07efeab7/kiwisolver-1.5.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:70ed9a45c7484d2b30cdacf60d220f494a1763b9fec1ad03285c6553fa0889f2", size = 2335355, upload-time = "2026-08-28T10:26:59.202Z" }, - { url = "https://files.pythonhosted.org/packages/6a/b3/a0f31d5e4e40af7dc97c36b8a74fdd3a36cf3c8bbd098da9a23466ff6a94/kiwisolver-1.5.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:98b208a7cc42c803445ef551d6753cc42a5ea13e9cab1ee66cd8b9cb70195330", size = 2426524, upload-time = "2026-08-28T10:27:01.181Z" }, - { url = "https://files.pythonhosted.org/packages/2f/c9/728f63bd58c72cafdc79fc306abeeac7391bec03b757a48dadeb30906521/kiwisolver-1.5.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:c6834b92dd2428e2dd85ef3d85f723d3c12f20aaf43a2ddd4f944ca25d833408", size = 2063430, upload-time = "2026-08-28T10:27:03.06Z" }, - { url = "https://files.pythonhosted.org/packages/84/df/ce188b96f92f9a2c958231da140768918cba53c9713dc887b82f85462118/kiwisolver-1.5.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:5d142e352eb13facc7dd047489aebdff6ba78576c239f1ea04931979caaf0567", size = 2597513, upload-time = "2026-08-28T10:27:05.072Z" }, - { url = "https://files.pythonhosted.org/packages/69/d6/76947c8203768968382e5bd74d9cc95654746703a61ea53015f2c74a2e06/kiwisolver-1.5.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f9b1c4900736e489a812c529100de4b8fb617d4db075e931e213c57424b83d9b", size = 2394488, upload-time = "2026-08-28T10:27:07.423Z" }, - { url = "https://files.pythonhosted.org/packages/e4/d4/14b21e4eb203c4d15425e8b6a2c625a320b4a1f2f7557eead63ffc30ffb7/kiwisolver-1.5.1-cp314-cp314t-win_amd64.whl", hash = "sha256:5978c3340f16a35c30f8ab2fa7bcf559973c55f1a5ef6970e1f621acf3c4db13", size = 75404, upload-time = "2026-08-28T10:27:08.892Z" }, - { url = "https://files.pythonhosted.org/packages/cb/f5/53157899fc7f45f76421b77b99eb1639dd0f83f26ff9d76300c96bb4a3b0/kiwisolver-1.5.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ca307d6c259e5c98d3cb9ade55342b47a6839762caf2536f3d7b46ee660cc82e", size = 72946, upload-time = "2026-08-28T10:27:10.944Z" }, -] - -[[package]] -name = "lightgbm" -version = "4.7.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "narwhals" }, - { name = "numpy" }, - { name = "scipy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/63/8e/4db5e29290d7e619c307fdb8dab0a0514090af2ce3ec483050e024ec6126/lightgbm-4.7.0.tar.gz", hash = "sha256:f8e20f682c9aabd000bcf4a7ed8aa6f473c1adfecccae34ec24e823d156f4af0", size = 1792896, upload-time = "2026-07-18T21:00:56.139Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/cd/05/7213965863cba1ed0150ad045bceed6276a1afaaaedbaeff4699ec4f0ccb/lightgbm-4.7.0-py3-none-macosx_10_15_x86_64.whl", hash = "sha256:dfc1cfe8e760387be1e7ba7a214688be21fdff96e4ed9749188f83e1877c2477", size = 1877851, upload-time = "2026-07-18T21:00:35.225Z" }, - { url = "https://files.pythonhosted.org/packages/b2/86/f4fe714f2e0bf3941705a20d7f6849dc476276d71236e82ea6b0d6539b86/lightgbm-4.7.0-py3-none-macosx_12_0_arm64.whl", hash = "sha256:129535462686f274df179133643118c5c5c5667167fe6c3a28d955f0b3c8e868", size = 1498914, upload-time = "2026-07-18T21:00:36.549Z" }, - { url = "https://files.pythonhosted.org/packages/c6/a3/b29580948b92e8c2f84dea70118ac702ff067dc52ec4ffb5d73c953536a5/lightgbm-4.7.0-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d4529acec5c6fefe4768302a529707d0ead90f6a6f42df694b856212e09695b8", size = 3349492, upload-time = "2026-07-18T21:00:37.943Z" }, - { url = "https://files.pythonhosted.org/packages/15/eb/837ea3b40cc36e22eeebb9785c01e42b2c255d033eea1d2d9ee8e2540e55/lightgbm-4.7.0-py3-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d23e922acd891e77212e4d0fbcee9ba973c96dee479491341d05ba595357ebb7", size = 3476028, upload-time = "2026-07-18T21:00:39.331Z" }, - { url = "https://files.pythonhosted.org/packages/d5/0b/c5c17d862b12ce292f24cd85d40f2f8f8981668fbdbd43fdc2625eccbc79/lightgbm-4.7.0-py3-none-win_amd64.whl", hash = "sha256:f42d1e5b32b6f170e606d7c689c6165671da98d7bf37f1addec2623efc8740c9", size = 1360833, upload-time = "2026-07-18T21:00:40.865Z" }, -] - -[[package]] -name = "markupsafe" -version = "3.0.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, - { url = "https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, - { url = "https://files.pythonhosted.org/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, - { url = "https://files.pythonhosted.org/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, - { url = "https://files.pythonhosted.org/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, - { url = "https://files.pythonhosted.org/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, - { url = "https://files.pythonhosted.org/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, - { url = "https://files.pythonhosted.org/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, - { url = "https://files.pythonhosted.org/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, - { url = "https://files.pythonhosted.org/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, - { url = "https://files.pythonhosted.org/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, - { url = "https://files.pythonhosted.org/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, - { url = "https://files.pythonhosted.org/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, - { url = "https://files.pythonhosted.org/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, - { url = "https://files.pythonhosted.org/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, - { url = "https://files.pythonhosted.org/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, - { url = "https://files.pythonhosted.org/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, - { url = "https://files.pythonhosted.org/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, - { url = "https://files.pythonhosted.org/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, - { url = "https://files.pythonhosted.org/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, - { url = "https://files.pythonhosted.org/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, - { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, -] - -[[package]] -name = "matplotlib" -version = "3.11.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "contourpy" }, - { name = "cycler" }, - { name = "fonttools" }, - { name = "kiwisolver" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "pillow" }, - { name = "pyparsing" }, - { name = "python-dateutil" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e7/c8/9aa712a0afb882649424dd8de8ad9aa6235e796e84c6052e8f6dc1598d0d/matplotlib-3.11.2.tar.gz", hash = "sha256:cec596316640f2b394b8f0daa0ea61a8eae82d017b620b9f202befb972a59ea4", size = 32660610, upload-time = "2026-09-11T19:05:31.214Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/b4/4facdb700bc236e4c676495ca3798d8134ab8668cbfe0f6b2e9ebf962a38/matplotlib-3.11.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:c9721f81275499da1feeb36a2cf8192ea086283b3bd16b7dc4c9d7aedb7396d6", size = 9478976, upload-time = "2026-09-11T19:04:00.12Z" }, - { url = "https://files.pythonhosted.org/packages/4c/a0/bba05f0a25bfeb0475c2f2a87958b31103aa6b7a5386017afb1a96212825/matplotlib-3.11.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1a3040b209f3968b4e84161df7b174f07a9fad33b0f2d7e48ea3bbd3075e2863", size = 9307932, upload-time = "2026-09-11T19:04:02.976Z" }, - { url = "https://files.pythonhosted.org/packages/d4/42/81d5cba4bb39b41b0998efc880e129f2d7cdb94256ca49678fdcd9c21be5/matplotlib-3.11.2-cp314-cp314-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bf3fe71fbfb8ec0e310e0bc8537c3405a01f38f25f9394ed2135e6202fed542b", size = 10674370, upload-time = "2026-09-11T19:04:06.912Z" }, - { url = "https://files.pythonhosted.org/packages/29/5e/52f56f93d5b20815ee0ae352f7a38eff25580edb180b4fa7df87cfeba68f/matplotlib-3.11.2-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8c8255de28f986d935a64c9ca71c0ec2d2f41d355691f5ea684725dc91413f71", size = 10954597, upload-time = "2026-09-11T19:04:09.948Z" }, - { url = "https://files.pythonhosted.org/packages/6e/ae/f8918a1b564db49c53cd603eb3c8748c5767e4ce0fb2cf1d7c38653e4c5d/matplotlib-3.11.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:a8756cc73d9af9a7fe0deb54ea2e75ef73b01d9e575877e72acad5458e660943", size = 10812583, upload-time = "2026-09-11T19:04:12.887Z" }, - { url = "https://files.pythonhosted.org/packages/ea/61/4315ef0d63937fe56f451887b4a0be3b8e8b22cf1d92a8d157caac54570d/matplotlib-3.11.2-cp314-cp314-win_amd64.whl", hash = "sha256:ecea603dd2fbf8242fd31a305a8b12a4ece2de28096870c65fdd0d1e35b8d9a6", size = 9505972, upload-time = "2026-09-11T19:04:15.788Z" }, - { url = "https://files.pythonhosted.org/packages/d5/00/02398b0a1a62ef73d9af0a957fe1773363eec9a6ca7e433c72ae2790fbce/matplotlib-3.11.2-cp314-cp314-win_arm64.whl", hash = "sha256:01dc8eaaab5a9fce9ff615eca82345728f289e4715b186ee10c6d85272fc26bb", size = 9184534, upload-time = "2026-09-11T19:04:18.691Z" }, - { url = "https://files.pythonhosted.org/packages/ec/c5/1e6ca10baa7e04f113c37856aadb4adc3dbb1c55fc7da6dc830bba1292d7/matplotlib-3.11.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:79a258f58253dfa025af80a9e9bb228d75fced007f0e93ae7423fefbde81a74d", size = 9526415, upload-time = "2026-09-11T19:04:21.39Z" }, - { url = "https://files.pythonhosted.org/packages/e0/93/7561727af07ccc84747953c6a33eee3981cc436d04f465032e707ca0e429/matplotlib-3.11.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c0b83f044ce10a98027b105b3931548719a6e8c7ef986b4362651e0b5367c8dc", size = 9361064, upload-time = "2026-09-11T19:04:24.107Z" }, - { url = "https://files.pythonhosted.org/packages/c1/40/ce270ef2a6794d94a409fe28b8dabedffc07c07f9e5dbb89e0f605bd157d/matplotlib-3.11.2-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e43b188f0a5b75447bcc197728258166aa64365770ed1caa36595a5e1ca4bbba", size = 10684626, upload-time = "2026-09-11T19:04:27.201Z" }, - { url = "https://files.pythonhosted.org/packages/45/46/396e0307dbf1c1f2d0a3654caa6dbc780dd5379f5f11257369cf8d21d9fe/matplotlib-3.11.2-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d3304eb5a59442a8867f6920d484591c0fa09ffc29e9260be2feec3351e25869", size = 10960602, upload-time = "2026-09-11T19:04:30.254Z" }, - { url = "https://files.pythonhosted.org/packages/1c/97/4defb84695477340aa88b27c9482cc1ca4de11f36f66b6c8c0ddd685def1/matplotlib-3.11.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:02329432ae5c6af87cf208ee575b701d698bdf0b1a3bb28cc6d53c36e967e575", size = 10822603, upload-time = "2026-09-11T19:04:32.964Z" }, - { url = "https://files.pythonhosted.org/packages/7d/e7/eab54551c1ec6b72dcd678a078b6c6da8d4c85ec1ff0016e81fe52866304/matplotlib-3.11.2-cp314-cp314t-win_amd64.whl", hash = "sha256:f2ac30cf5eb5dff1b584627ae0b0e1186551a4f69ae3c75073911da497a29170", size = 9551206, upload-time = "2026-09-11T19:04:35.811Z" }, - { url = "https://files.pythonhosted.org/packages/4e/0e/04698032cb4d8fb1e30d8ffec5ded20f6f582b03c8c58b2385c95245f158/matplotlib-3.11.2-cp314-cp314t-win_arm64.whl", hash = "sha256:cf41ecd1b0c0b6f7177ed965a54c2afbe888715c7cf6054dc12d53bc1494002c", size = 9234977, upload-time = "2026-09-11T19:04:38.972Z" }, -] - -[[package]] -name = "ml-dtypes" -version = "0.5.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/72/4e/1339dc6e2557a344f5ba5590872e80346f76f6cb2ac3dd16e4666e88818c/ml_dtypes-0.5.4-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:2b857d3af6ac0d39db1de7c706e69c7f9791627209c3d6dedbfca8c7e5faec22", size = 673781, upload-time = "2025-11-17T22:32:11.364Z" }, - { url = "https://files.pythonhosted.org/packages/04/f9/067b84365c7e83bda15bba2b06c6ca250ce27b20630b1128c435fb7a09aa/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:805cef3a38f4eafae3a5bf9ebdcdb741d0bcfd9e1bd90eb54abd24f928cd2465", size = 5036145, upload-time = "2025-11-17T22:32:12.783Z" }, - { url = "https://files.pythonhosted.org/packages/c6/bb/82c7dcf38070b46172a517e2334e665c5bf374a262f99a283ea454bece7c/ml_dtypes-0.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14a4fd3228af936461db66faccef6e4f41c1d82fcc30e9f8d58a08916b1d811f", size = 5010230, upload-time = "2025-11-17T22:32:14.38Z" }, - { url = "https://files.pythonhosted.org/packages/e9/93/2bfed22d2498c468f6bcd0d9f56b033eaa19f33320389314c19ef6766413/ml_dtypes-0.5.4-cp314-cp314-win_amd64.whl", hash = "sha256:8c6a2dcebd6f3903e05d51960a8058d6e131fe69f952a5397e5dbabc841b6d56", size = 221032, upload-time = "2025-11-17T22:32:15.763Z" }, - { url = "https://files.pythonhosted.org/packages/76/a3/9c912fe6ea747bb10fe2f8f54d027eb265db05dfb0c6335e3e063e74e6e8/ml_dtypes-0.5.4-cp314-cp314-win_arm64.whl", hash = "sha256:5a0f68ca8fd8d16583dfa7793973feb86f2fbb56ce3966daf9c9f748f52a2049", size = 163353, upload-time = "2025-11-17T22:32:16.932Z" }, - { url = "https://files.pythonhosted.org/packages/cd/02/48aa7d84cc30ab4ee37624a2fd98c56c02326785750cd212bc0826c2f15b/ml_dtypes-0.5.4-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:bfc534409c5d4b0bf945af29e5d0ab075eae9eecbb549ff8a29280db822f34f9", size = 702085, upload-time = "2025-11-17T22:32:18.175Z" }, - { url = "https://files.pythonhosted.org/packages/5a/e7/85cb99fe80a7a5513253ec7faa88a65306be071163485e9a626fce1b6e84/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2314892cdc3fcf05e373d76d72aaa15fda9fb98625effa73c1d646f331fcecb7", size = 5355358, upload-time = "2025-11-17T22:32:19.7Z" }, - { url = "https://files.pythonhosted.org/packages/79/2b/a826ba18d2179a56e144aef69e57fb2ab7c464ef0b2111940ee8a3a223a2/ml_dtypes-0.5.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0d2ffd05a2575b1519dc928c0b93c06339eb67173ff53acb00724502cda231cf", size = 5366332, upload-time = "2025-11-17T22:32:21.193Z" }, - { url = "https://files.pythonhosted.org/packages/84/44/f4d18446eacb20ea11e82f133ea8f86e2bf2891785b67d9da8d0ab0ef525/ml_dtypes-0.5.4-cp314-cp314t-win_amd64.whl", hash = "sha256:4381fe2f2452a2d7589689693d3162e876b3ddb0a832cde7a414f8e1adf7eab1", size = 236612, upload-time = "2025-11-17T22:32:22.579Z" }, - { url = "https://files.pythonhosted.org/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" }, -] - -[[package]] -name = "mlx" -version = "0.32.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mlx-metal", marker = "sys_platform == 'darwin'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/99/2f/ed874a2de4ad30b47f27960373880aa24b53d8963ea81125e12e460340c2/mlx-0.32.4-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:f1e6ef441e1e2ce54159022a8fa5597f8ea4f7acf278752049c4f774bbd0501d", size = 668901, upload-time = "2026-10-14T21:16:36.971Z" }, - { url = "https://files.pythonhosted.org/packages/12/c6/1e15326918233e1ee74f55bcc5577f6cce1c4257e5a1bbadcdb30b589dca/mlx-0.32.4-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:83adc889ae89057fd9f292ca1a9ee56da43f181fccc7ae48a060d7aadc179eaa", size = 668899, upload-time = "2026-10-14T21:16:39.057Z" }, - { url = "https://files.pythonhosted.org/packages/95/2f/f84f5f16000adbdb1c53583f53da045165013ebb62f6943c169714be011f/mlx-0.32.4-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:1dd56491b44a2bfa217f60da66a5c9c889add7ad9dc2cbcf526389ea31d83bee", size = 668817, upload-time = "2026-10-14T21:16:40.738Z" }, - { url = "https://files.pythonhosted.org/packages/e3/7c/c90d57df7775421fe43f7a235f6cee85a67c4aecc7325051b23fff58335b/mlx-0.32.4-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:f6cbda9cb8c33c9f8673a2b69fb19bdab75b5032d4e47a42a0e7b36db1c3ebcc", size = 692106, upload-time = "2026-10-14T21:16:50.019Z" }, - { url = "https://files.pythonhosted.org/packages/6e/07/90e629a9667ec35b2aa8c84e0e096ff36fdd58731dc9d182f5ee59a488d0/mlx-0.32.4-cp314-cp314t-macosx_15_0_arm64.whl", hash = "sha256:4f95afac25431585b34846d472ea0ad5db83fc97fb014d62c70082456f6b6b5c", size = 692107, upload-time = "2026-10-14T21:16:51.978Z" }, - { url = "https://files.pythonhosted.org/packages/d0/20/bdb26b961fc8f0d54e6b36dd5360449f0cf8591d14cee10470c14ce1e154/mlx-0.32.4-cp314-cp314t-macosx_26_0_arm64.whl", hash = "sha256:63ac6fde8a55d688448fc25eb32b645276501e0a30ff961a72958dfd0a4065df", size = 692103, upload-time = "2026-10-14T21:16:53.963Z" }, -] - -[[package]] -name = "mlx-metal" -version = "0.32.4" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/77/24/7e5b7ed5d28abacab4577c7a8dd1fca6e9058504f5f95ae4333c09420e2e/mlx_metal-0.32.4-py3-none-macosx_14_0_arm64.whl", hash = "sha256:8019b2db622599267d66cc168b2f0dbf5af929790a6c3b3913e00eabc42a7947", size = 44958619, upload-time = "2026-10-14T21:14:45.727Z" }, - { url = "https://files.pythonhosted.org/packages/54/8c/9754dd6c4db485037f59cc8be6495738e14cdeb8bf19d5bddb5e96da890f/mlx_metal-0.32.4-py3-none-macosx_15_0_arm64.whl", hash = "sha256:9f0a3cc26d89724434d9684932fc6f3e4e83669d2a895d858ed2c8a55eb84e8b", size = 44942974, upload-time = "2026-10-14T21:14:50.012Z" }, - { url = "https://files.pythonhosted.org/packages/2a/d6/4340ed6e184f71693d7f531c9ae55d677451351245760d71f8dd714dcbe5/mlx_metal-0.32.4-py3-none-macosx_26_0_arm64.whl", hash = "sha256:514df92d41a286895d2564b4b18f4da45096bd7a4a0c47e1042f345ad5a7f936", size = 78106861, upload-time = "2026-10-14T21:14:56.184Z" }, -] - -[[package]] -name = "mpmath" -version = "1.3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, -] - -[[package]] -name = "narwhals" -version = "2.27.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/bf/21/f64d6b2dbea7bf3f8c38cdc786dcc6ef012ca3d173ad208c782c9a7bedf6/narwhals-2.27.1.tar.gz", hash = "sha256:aed93076a3ea42d9c32c88e4eb5ea422a21937011cbe1f480f9572a523c82094", size = 735013, upload-time = "2026-10-10T06:52:18.113Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/89/5d4c86da1130d9059681e5b6cd7645df5c10279a6a079c5c37dcb2cc6f3f/narwhals-2.27.1-py3-none-any.whl", hash = "sha256:d057df13f5852b8e157596e82eb5e955fad267425df5e420e0ee9863da483b31", size = 483211, upload-time = "2026-10-10T06:52:16.32Z" }, -] - -[[package]] -name = "networkx" -version = "3.6.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, -] - -[[package]] -name = "numpy" -version = "2.5.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/95/b0/c7453d0b6e2073c3264468b106ee1563750cecc910965e67357e3698c83e/numpy-2.5.4.tar.gz", hash = "sha256:9a94cf751c9ad8ebaa835bcd3d40dacf8534ad086b88c38029b65123c7999d2a", size = 20866315, upload-time = "2026-10-10T20:05:31.422Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/99/ba/005cb5edd580d2f84d7ca3206b92dc17d4388e56e6f87ffe8f2762f83139/numpy-2.5.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:c668b2f0d651605b58892644b0e302c7157f7159544227758c896982ef384b18", size = 17005499, upload-time = "2026-10-10T20:03:37.961Z" }, - { url = "https://files.pythonhosted.org/packages/f3/49/fee7587c33ee35f7977f9051d7f2023d4e7246d62710c80f20c2361ea232/numpy-2.5.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:ffa6ce09a1c6a08e9667dd9c97aa0b14184e8d18f2a14b78b2a2328c9147f076", size = 12019666, upload-time = "2026-10-10T20:03:40.606Z" }, - { url = "https://files.pythonhosted.org/packages/d5/b2/c6ce165acffceb15a82c07b9cc77d391f86b3f379ba62911908ae5d34b91/numpy-2.5.4-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:956555e0603a4d38019ae6925711cb9dc43195c076a928accf7ea5d50bddfe53", size = 5455617, upload-time = "2026-10-10T20:03:43.138Z" }, - { url = "https://files.pythonhosted.org/packages/77/7f/dd85ce260a669a89be06842cf355d7353a33e6cfbc590fb8ebb947d88dc9/numpy-2.5.4-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:2c2c4afffdeb7920e445028dd71eb932cac3e704792e964bc2a232426d4f1255", size = 6791932, upload-time = "2026-10-10T20:03:44.874Z" }, - { url = "https://files.pythonhosted.org/packages/63/d6/34b0a2b0741386a63025a65a2c09caaaaaad6d0ca95b66cd65c30dd7fcb5/numpy-2.5.4-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4054173604cd8658796053f1f3bc0befb68ec1c0762c57fdad61e199256a8617", size = 15710899, upload-time = "2026-10-10T20:03:46.839Z" }, - { url = "https://files.pythonhosted.org/packages/16/d5/928078d2b28f26829b138b4a6c3980045022fb409f570657a224ae60ef4e/numpy-2.5.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d549420b8858885cea8838a727842249218b9c1da24dd517e25c9c7a948310a3", size = 16721710, upload-time = "2026-10-10T20:03:49.489Z" }, - { url = "https://files.pythonhosted.org/packages/f9/cf/673fd1b8f4cd78eb6320e87ec4c90ac19c095644259e3749853a405c70f4/numpy-2.5.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:823874a507a84af050493b622affde94b6f7c3a0dc22cb2801381bc03b871c00", size = 17066182, upload-time = "2026-10-10T20:03:52.25Z" }, - { url = "https://files.pythonhosted.org/packages/f3/92/a77b5061b1b3e2643928c37976d79ee173e1b171ed158b7a3c61056b41bc/numpy-2.5.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:4e263278bfb5ee6409db8aedbc4cc32973b1b82bc1e8d3c668551d04d83a7e37", size = 18480315, upload-time = "2026-10-10T20:03:55.39Z" }, - { url = "https://files.pythonhosted.org/packages/bb/1d/1486ef3d3fb2279fd93c4c43c1bbbf1ca389a19816696684409f71babaab/numpy-2.5.4-cp314-cp314-win32.whl", hash = "sha256:cfd73180400042a7c532d30c5e287bdd03c59ff9ee1b4c0316af0539e29dfe23", size = 6185739, upload-time = "2026-10-10T20:03:58.186Z" }, - { url = "https://files.pythonhosted.org/packages/52/9a/e1e512ebc948d5b9dd33b08736760f0ebbed2848fd4eda1f553088a6dcee/numpy-2.5.4-cp314-cp314-win_amd64.whl", hash = "sha256:2ca144f15135b6212a5c47b1e2aeca6e412f102f95a2d5d88d8aec77eb255de3", size = 12703552, upload-time = "2026-10-10T20:04:00.28Z" }, - { url = "https://files.pythonhosted.org/packages/2c/05/de709a982d7bbcd688a3fad71f002e9ff80c2db39e03ee726609b610f1d1/numpy-2.5.4-cp314-cp314-win_arm64.whl", hash = "sha256:468397ba3c64427474706e5c9123fe266395496714dc684294eac75cd4930d1e", size = 10803901, upload-time = "2026-10-10T20:04:02.659Z" }, - { url = "https://files.pythonhosted.org/packages/13/34/083570ada3bb2a30fbe5d77c8c6fef9141144a15d33e6f793a67e9749ab8/numpy-2.5.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1ef3aa6d7e29bb13677323114280b05acc57607fa2300e66432d665d5418a162", size = 12138695, upload-time = "2026-10-10T20:04:05.012Z" }, - { url = "https://files.pythonhosted.org/packages/94/06/1f9c24db48eef0c2d1207e3b11fffb0478e39dfd8c1e1be7476936885eed/numpy-2.5.4-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:98b053943e5a0474ec0da309d2cb9d3f18ea57f8a2067c2ab7b5f763d1068380", size = 5574615, upload-time = "2026-10-10T20:04:07.316Z" }, - { url = "https://files.pythonhosted.org/packages/da/0f/593fba2e1560e949123bc7d2fc48b5893d56e58cd4bd5a273d2fbf60b220/numpy-2.5.4-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:b64a85f40e154983960a4167d4c1d57a50c7f109b3d3264a3a984154e90a8454", size = 6889383, upload-time = "2026-10-10T20:04:09.918Z" }, - { url = "https://files.pythonhosted.org/packages/eb/9f/b799dfdce4e05e80ed4bc815c71ff343a11533b2c0ffc221cae8538cda63/numpy-2.5.4-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a813ed7719bf45463c51779e6a98d0385fe905e48447526938a4b8337333d551", size = 15753763, upload-time = "2026-10-10T20:04:12.278Z" }, - { url = "https://files.pythonhosted.org/packages/34/88/16c5f12f86f5ad2817c4d103205131fc6c8acb3d1878af05a1a4f23ec859/numpy-2.5.4-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c9b80cdf5cedba0e90d93fa5f9a333c4d65bd545cd669b71bb97ce2b703c9d73", size = 16757212, upload-time = "2026-10-10T20:04:14.799Z" }, - { url = "https://files.pythonhosted.org/packages/ff/4f/a1fe40e18a898e6a5089f4f0d891f0a493eb0574d5b34458f0fbe5aa3e5c/numpy-2.5.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:2199ed071f460487c8db2c0e5c0b564494190edb4772fe80f9aad88b2604def5", size = 17116471, upload-time = "2026-10-10T20:04:17.58Z" }, - { url = "https://files.pythonhosted.org/packages/aa/46/e923a11c78e65c1722e7aaad817c06bd591324174b9d28ce5d31eee4d432/numpy-2.5.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:64f9c9878c1938476365e11ccfb6b770f3b9e5f045ccddc514235041e6959365", size = 18524063, upload-time = "2026-10-10T20:04:20.365Z" }, - { url = "https://files.pythonhosted.org/packages/5a/fa/84ab064514440c1f64a1b21088f2c82756defdd05e07c75ab233899565b2/numpy-2.5.4-cp314-cp314t-win32.whl", hash = "sha256:64d1c8ac28a4077cf987e0a71a7a0ef7e2df70722f07f0baa42dbb7eb6938647", size = 6340926, upload-time = "2026-10-10T20:04:22.865Z" }, - { url = "https://files.pythonhosted.org/packages/7e/7e/6cd886876f435b10685db9b9f7eeb70356f99e052116f4e5f11c5792c714/numpy-2.5.4-cp314-cp314t-win_amd64.whl", hash = "sha256:067374eb538c34c745436365cf7b0112595c1d326f21ce4ff340f61230239fbb", size = 12901584, upload-time = "2026-10-10T20:04:24.99Z" }, - { url = "https://files.pythonhosted.org/packages/38/1b/3c1684f6a06f7307f2335fca6e486cb162847fb97e91d65f8eb5cabad213/numpy-2.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:e94aef2c639da4a960ad0db8e06471208d8589974953d78b61d345b4eb99e394", size = 10891152, upload-time = "2026-10-10T20:04:27.52Z" }, -] - -[[package]] -name = "nvidia-cublas" -version = "13.1.1.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, - { url = "https://files.pythonhosted.org/packages/3b/cd/154ca20c38269e05eff77c1464e6c1da89f50a6390b565e9d82e06bc11e1/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436", size = 423138758, upload-time = "2026-04-08T18:46:58.655Z" }, -] - -[[package]] -name = "nvidia-cuda-cupti" -version = "13.0.85" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, - { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, -] - -[[package]] -name = "nvidia-cuda-nvrtc" -version = "13.0.88" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, - { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, -] - -[[package]] -name = "nvidia-cuda-runtime" -version = "13.0.96" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, - { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, -] - -[[package]] -name = "nvidia-cudnn-cu13" -version = "9.24.0.43" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/ca/30/7c257e3d5cb4fecb147b93895c66e29c93f8e76d74b45bb418ff0587c4ec/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:a6812a554a1ff0413e9c52b84c26c050380649ab9615f9c16bded368ce9f421f", size = 650976863, upload-time = "2026-07-02T16:23:39.248Z" }, - { url = "https://files.pythonhosted.org/packages/5c/ba/791cffd048fe5b044e620df55267e3e95c0e6e07d50b41e377c03dfc910f/nvidia_cudnn_cu13-9.24.0.43-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:71f181cd810e90f9b6023b01186fe82d13d65f0ec098581ee201d39fad769e4b", size = 553099438, upload-time = "2026-07-02T16:27:42.58Z" }, -] - -[[package]] -name = "nvidia-cufft" -version = "12.0.0.61" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, - { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, -] - -[[package]] -name = "nvidia-cufile" -version = "1.15.1.6" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, - { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, -] - -[[package]] -name = "nvidia-curand" -version = "10.4.0.35" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, - { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, -] - -[[package]] -name = "nvidia-cusolver" -version = "12.0.4.66" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-cusparse", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, - { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, -] - -[[package]] -name = "nvidia-cusparse" -version = "12.6.3.3" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, - { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, -] - -[[package]] -name = "nvidia-cusparselt-cu13" -version = "0.8.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" }, - { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" }, -] - -[[package]] -name = "nvidia-nccl-cu13" -version = "2.30.7" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/21/a73174c6157101bdf1ffc22b517f76ff0082613989dd9bc8f43e8034caac/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:ca786ffa5a647c75d4d1f5cc72a6c4f537947e2ba8823d7c8aaf768e7a7b9f77", size = 215983881, upload-time = "2026-06-09T03:23:15.633Z" }, - { url = "https://files.pythonhosted.org/packages/3f/34/c500f90c7ae641b8e0f98965b36b8a7ac79cc8b296e8d251fe3eb592ee54/nvidia_nccl_cu13-2.30.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:cefa7fdb9710efd0f39c5f1be1d61ff6fc9a996c451265bd7fbdcf9455ed4b50", size = 215965170, upload-time = "2026-06-09T03:23:39.73Z" }, -] - -[[package]] -name = "nvidia-nvjitlink" -version = "13.4.92" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1d/6b/eef7a9e32872b8f41e145bf10cddc9af26e153c338852811fe9a9baddf9e/nvidia_nvjitlink-13.4.92-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e0391f24ed94ec879b84e3da4d4ec320c879aff681f2c7a638462f7199284323", size = 42452378, upload-time = "2026-09-16T20:45:29.042Z" }, - { url = "https://files.pythonhosted.org/packages/1f/a8/1cbd4014898af8b419e69b0d7dbc63da2121ee92d92b47d59f4fe9075349/nvidia_nvjitlink-13.4.92-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:25f74fad0d654271c921ac4dca614bd6258bc21791242fc7b2289dad7ae9c099", size = 40420120, upload-time = "2026-09-16T20:45:19.163Z" }, -] - -[[package]] -name = "nvidia-nvshmem-cu13" -version = "3.4.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, - { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, -] - -[[package]] -name = "nvidia-nvtx" -version = "13.0.85" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, - { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, -] - -[[package]] -name = "omegaconf" -version = "2.4.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyyaml" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/cc/10/1cabf3188a4d9684a06568c7fd76f9536eede625733c8399cf1f5337b83e/omegaconf-2.4.0.tar.gz", hash = "sha256:7e219a70bea371da250285b9c01c7db027bc421722799c2cc4ff14bc4a2a753b", size = 3544508, upload-time = "2026-10-09T12:44:02.326Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/4f/bd/0d9f8247b0819385213d09564eda5c52d3673313c2414d51f72883ec24cd/omegaconf-2.4.0-py3-none-any.whl", hash = "sha256:f141a664c777582405e3cc6e0545ecc79429e79f844995b1a6e71c35946a025c", size = 254152, upload-time = "2026-10-09T12:43:59.764Z" }, -] - -[[package]] -name = "onnx" -version = "1.22.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "ml-dtypes" }, - { name = "numpy" }, - { name = "protobuf" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/04/19/8ea73a64b368b75fe339771a20a02bc61ea1f551484c9e3d9d0bfbd0450f/onnx-1.22.0.tar.gz", hash = "sha256:ef40c0aaf0b643857ea9306fc7eddce17eaf9fb0407e4801f1fc5758443a38e0", size = 12024721, upload-time = "2026-06-15T12:50:05.354Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ee/6a/481561f1093834376ed493e4ca42a73e5be0d50031f2969c86593bdc7c96/onnx-1.22.0-cp312-abi3-macosx_12_0_universal2.whl", hash = "sha256:596fbf0490947533c1c1045ba860851dc9fb77471023dac9a71ba5b42ceab103", size = 20167081, upload-time = "2026-06-15T12:49:32.078Z" }, - { url = "https://files.pythonhosted.org/packages/84/55/b34fc2aa30aa54b4a775402d24c4082242c720283a274fe976ac8eb94480/onnx-1.22.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ae5a563f281cd9d2845622cecf6c092a57e4ee1b138f66fdbbdd4200567a5e16", size = 18889249, upload-time = "2026-06-15T12:49:34.7Z" }, - { url = "https://files.pythonhosted.org/packages/09/a6/bd32357e6cc1ecb473afd78193d7231724f284435d2db25696ecfaaa1503/onnx-1.22.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:955e02e1f6d385b53d52f9cd7b9cdf5caf417c300bcfe3c64c6d542be763845b", size = 19106514, upload-time = "2026-06-15T12:49:37.424Z" }, - { url = "https://files.pythonhosted.org/packages/5a/9d/3af461ac6c714b8b369cb71499659932f4f12cfb066250b62f7567c3d530/onnx-1.22.0-cp312-abi3-pyemscripten_2025_0_wasm32.whl", hash = "sha256:82e9f27fc1223cb06d68a56bed6f9d3caf3d0dad1b61bce45006d529b15bd94c", size = 16966387, upload-time = "2026-06-15T12:49:40.918Z" }, - { url = "https://files.pythonhosted.org/packages/d0/f0/68195b5e5a53e333faf2660f5352ee43738d0e42fc5216cc6b1871a9fbfb/onnx-1.22.0-cp312-abi3-win32.whl", hash = "sha256:cc8b66b312f8f03a53e268afb67180a2d97dd12cc79e2b61361c6c0073448016", size = 17081568, upload-time = "2026-06-15T12:49:43.398Z" }, - { url = "https://files.pythonhosted.org/packages/13/a8/734725bb703c5fabb687f79c79e51249475212b3eb37771ac4a4ac9b487f/onnx-1.22.0-cp312-abi3-win_amd64.whl", hash = "sha256:72ccebab3bac07215c204ce8848d42e78eaaa666badbf72d25cd359b9f269e3a", size = 17213290, upload-time = "2026-06-15T12:49:45.933Z" }, - { url = "https://files.pythonhosted.org/packages/bd/2a/8ce48d8ae26a8761ad4e5dc771961b155c5c3c7c8540ec7f2f2d71b69af0/onnx-1.22.0-cp312-abi3-win_arm64.whl", hash = "sha256:f3c120dcdb70ad738f3c061b32798f408ea299eb69f84dd69ab4a6bf3c2ec01f", size = 17207030, upload-time = "2026-06-15T12:49:48.635Z" }, - { url = "https://files.pythonhosted.org/packages/f3/13/47323b97846387848efb1044ded11bb94b83526f3d1fbdb37c6480d4520f/onnx-1.22.0-cp314-cp314t-macosx_12_0_universal2.whl", hash = "sha256:19e45e4af88e3fe3261458d4b8cc461957ae2782a358a3560503569bf3b23b72", size = 20176465, upload-time = "2026-06-15T12:49:51.311Z" }, - { url = "https://files.pythonhosted.org/packages/13/0c/d3b8a7e7eee123938586c608bb9894b5723f2342b9450c0eec59fbec7099/onnx-1.22.0-cp314-cp314t-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c21a0e59fd967a95b358e4a6e756d1f1eec2d304a83480f329f66e30d2bf0223", size = 18894028, upload-time = "2026-06-15T12:49:54.451Z" }, - { url = "https://files.pythonhosted.org/packages/b8/8a/da2a97ab46fe6e0cd9beb3ac14603a22f5be492f9ca347faf8233a07bb33/onnx-1.22.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2632406b8f523ef2e2873c363f90b20a3d88c0fbcfac757d3addffccf8f452c2", size = 19110420, upload-time = "2026-06-15T12:49:57.665Z" }, - { url = "https://files.pythonhosted.org/packages/b9/a3/ce984063017518307ebfaa545782fc400e593dc2d7fdf4f23ce4be1ed197/onnx-1.22.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a3a39fc4643867aecb33417fdddb11e308ee79d2d4a584b9d50cc7aec2091b13", size = 17237547, upload-time = "2026-06-15T12:50:00.382Z" }, - { url = "https://files.pythonhosted.org/packages/00/50/257a880384a1dd502d543b0067945074d63cd17d0840e958355bc8197da8/onnx-1.22.0-cp314-cp314t-win_arm64.whl", hash = "sha256:8e268cdc0547e3949799ffd4a44451dc2b9080b57d0824a2db680b6ec65506f0", size = 17231391, upload-time = "2026-06-15T12:50:03.047Z" }, -] - -[[package]] -name = "onnxruntime" -version = "1.29.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "flatbuffers" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "protobuf" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/ff/a3/718e1b83096a1bc7b0fc8014c23d4cf795559fe666961cfac4fc038a4871/onnxruntime-1.29.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:e74b278af1d949876f5d91d1268fd6c680e79f2bac194967394eaba9fdf69e7e", size = 21431104, upload-time = "2026-08-17T22:54:20.118Z" }, - { url = "https://files.pythonhosted.org/packages/4e/17/c75e78ddc1fe69b6ebaef7fe88ac83f29bfe10955e3a0d2436d93473c91c/onnxruntime-1.29.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:939e5d65f332e6d399774b2bd0d3559fd8fa629c1e77833db29d968d2384f23d", size = 20818488, upload-time = "2026-08-17T22:54:23.147Z" }, - { url = "https://files.pythonhosted.org/packages/65/54/9f197c578d3d3d7bea16971e233e5483981228eec73748585cf7b5933403/onnxruntime-1.29.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:6c0c37b92f67ed68dd36221ce0403e1d9bd4f7efce724439978a2597848530e5", size = 23136994, upload-time = "2026-08-17T22:54:26.321Z" }, - { url = "https://files.pythonhosted.org/packages/24/53/4616a55d2495679cfd0195f968feb3d74fe30e26467d168ee243ac97c089/onnxruntime-1.29.0-cp314-cp314-win_amd64.whl", hash = "sha256:4a3129ae56e70d2618ff773920166916310370a7e3cacb60b9e0e8910092725f", size = 14350643, upload-time = "2026-08-17T22:54:28.794Z" }, - { url = "https://files.pythonhosted.org/packages/0a/0f/c338cb5500a522c7e671a3bb1276f4562404fbecce8a0e274565aa968484/onnxruntime-1.29.0-cp314-cp314-win_arm64.whl", hash = "sha256:e417ef8628dcce310d2d53023e750ea298ec14d4341ae6dc3a572bfd9bc7fa97", size = 14124294, upload-time = "2026-08-17T22:54:31.015Z" }, - { url = "https://files.pythonhosted.org/packages/b8/e7/61064289a9a1301b25c1f0f574fe98aba31c2d388db3c1dbec664f78621f/onnxruntime-1.29.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:11264bb58f7b7cf6af835ab10d36838d73680580820fd6f51d90124a1ca8f449", size = 20826174, upload-time = "2026-08-17T22:54:34.283Z" }, - { url = "https://files.pythonhosted.org/packages/60/21/d0c04b561b46e9bff89b5f500fb7415b8ca0669f7902204f76ab06bb0c7e/onnxruntime-1.29.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:1ea91cef3b971506e51ae9c37c16d027774ec64994a524ec1bdfb027d68a9832", size = 23138547, upload-time = "2026-08-17T22:54:37.491Z" }, -] - -[[package]] -name = "packaging" -version = "26.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, -] - -[[package]] -name = "pandas" -version = "3.0.6" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, - { name = "python-dateutil" }, - { name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e2/17/d7b106e05bfa642e8694451e7d3d759c6a241c5386a5d962e4f66c047e06/pandas-3.0.6.tar.gz", hash = "sha256:66b07ef7315a31bfe1089cd3d71a7de781c9dca986762d0b4fe7c0ef17465d10", size = 4667686, upload-time = "2026-09-17T23:23:18.345Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/75/55/1a8875395b05ccd572cbca0b9255dcd2db6e6508e632a558c1a6884b39ad/pandas-3.0.6-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:ee913a91669056c1de1a6b733fbfeab711de9e54e3bee2dfa5fe79d9457247d1", size = 10492213, upload-time = "2026-09-17T23:21:40.746Z" }, - { url = "https://files.pythonhosted.org/packages/35/61/47ae13476995cc8a40cd609e93e7cf11f273d8692925c2903cb6d38aa0d1/pandas-3.0.6-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:ff51a4459ed036e93d1eb1bb5e6e7b28685d3cb6b7c12b91c05b31024e234729", size = 10156618, upload-time = "2026-09-17T23:21:44.142Z" }, - { url = "https://files.pythonhosted.org/packages/bc/f2/cc5f2adb8d6e86a85d9fb5128f8cf205a61189336f70d1f7faf0d1b53ec9/pandas-3.0.6-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:654aae059295dbba6ecd2328ca12712a2cf1676214c8699f1c29213f7ccf9c34", size = 10375479, upload-time = "2026-09-17T23:21:47.159Z" }, - { url = "https://files.pythonhosted.org/packages/ca/ba/ffdcb19be4ff6bfe7d969e7cef2c567c633df5a3a1cc1053394ad053bca8/pandas-3.0.6-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:62f51d7f651c8054c5e82a69265c98082e795d1442df7ca6edc3a545d61214b1", size = 10783244, upload-time = "2026-09-17T23:21:50.367Z" }, - { url = "https://files.pythonhosted.org/packages/77/5b/e150075b2c6eb69fae896f2d9239bc6ed07db97735971d53d66de6553460/pandas-3.0.6-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:22172a92e7ee678ec0140c7af4fc9366b55413834a1cd86af78b3caa0b0574de", size = 11382223, upload-time = "2026-09-17T23:21:53.355Z" }, - { url = "https://files.pythonhosted.org/packages/d6/8a/b441c587dc7355bf6e1f68a91b4f76a6c29740f0c23be3acc5d4ebbeea6d/pandas-3.0.6-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:583be68728a31d0d750d5b8d9e00f02b153df0d4655f858bde93cb84cfc4227c", size = 11854881, upload-time = "2026-09-17T23:21:56.342Z" }, - { url = "https://files.pythonhosted.org/packages/b7/e9/f43410fada510b43fec09993c08f552086c3d247d3ee801a678f3cb10ea5/pandas-3.0.6-cp314-cp314-win_amd64.whl", hash = "sha256:77ccbe5057aece6fc172b9b77f19c04335af6882bc2e10c8f3ee4e6bfb3da553", size = 9791672, upload-time = "2026-09-17T23:21:59.332Z" }, - { url = "https://files.pythonhosted.org/packages/8b/9e/db14c059c21f9baa1907d436f8bf30e0c76c6288225c5e8b79a08ba8b2c5/pandas-3.0.6-cp314-cp314-win_arm64.whl", hash = "sha256:fb625f426b375bcc96e3a04c5d5d266cd7be6ae5d6866e0e703382ab5164068c", size = 9121246, upload-time = "2026-09-17T23:22:02.123Z" }, - { url = "https://files.pythonhosted.org/packages/67/ba/bad0f8dac020ab38a8637fddab01a57a82da7a496a6e6f19590aad53ab62/pandas-3.0.6-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:9e492cd4bdba6778de4fe0df7f4590c012161ebcf9902dce01b01dc683105514", size = 10920612, upload-time = "2026-09-17T23:22:05.404Z" }, - { url = "https://files.pythonhosted.org/packages/c4/a9/b500982e9aac6d52a58da4ad3f11e14168a315b06906b3f397c427878065/pandas-3.0.6-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:d7dcd21238cbb4828ff148481ba01cac8946dc5121457b5aeba28636f8f99a60", size = 10570213, upload-time = "2026-09-17T23:22:08.44Z" }, - { url = "https://files.pythonhosted.org/packages/4b/fa/e6ecd0073c98be8f840ac3125b955272835d7d9fd69f5944383b164deb5e/pandas-3.0.6-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6ff482fa91fa2bafd92e8fe66ce3645c851824310f295c1f0a2f96e928fc4541", size = 10252356, upload-time = "2026-09-17T23:22:11.302Z" }, - { url = "https://files.pythonhosted.org/packages/04/f5/001e230a7a7803590d9275a1a3f7e1bb605e3a495cfe5e8d3a532090621b/pandas-3.0.6-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:db7ec631f26223beee8e5c9e0b8f23c24d8197bbd1d982421d4e3188bea51965", size = 10655952, upload-time = "2026-09-17T23:22:14.283Z" }, - { url = "https://files.pythonhosted.org/packages/ca/ab/bab587148a3852c96aae26c4b5f9e04ce2221801ad94e166b4fbf969ede0/pandas-3.0.6-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bd75ed0c840f709fc2ae26ddd9534ac77ca1a48ac0cce521a74acaa85f3340a7", size = 11274177, upload-time = "2026-09-17T23:22:17.352Z" }, - { url = "https://files.pythonhosted.org/packages/f3/32/74b48d87df2b80892d713c149abfe36d5db4de41b4eccb042a2bc07dafc1/pandas-3.0.6-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ef738d71d1059245b6bb03e312be06d8b3821326a83486c1ad03b9aba3710e44", size = 11719041, upload-time = "2026-09-17T23:22:20.227Z" }, - { url = "https://files.pythonhosted.org/packages/f6/c6/d64b72d64d7eb0fad9fe424d34138e45dee70ddbea1360dcd0adf30e28f6/pandas-3.0.6-cp314-cp314t-win_amd64.whl", hash = "sha256:429d9df32731ab01383ed98f2baa7a60368090d1a94fc06019a12062510e8630", size = 10191388, upload-time = "2026-09-17T23:22:23.524Z" }, - { url = "https://files.pythonhosted.org/packages/7a/30/5e5b2ccabeca73ae2b03fc82bca3eabb7466cf43737f05ac08d591665d47/pandas-3.0.6-cp314-cp314t-win_arm64.whl", hash = "sha256:a4dbd4dc65cbe645b92b8785d0f96dd7311010dc6606cf620e51b07b8788a12a", size = 9387796, upload-time = "2026-09-17T23:22:26.64Z" }, -] - -[[package]] -name = "pillow" -version = "13.0.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/ce/a4/9cc30690a394a8305e8be4c9036aa16d0dfddadd09079860b582f562496d/pillow-13.0.0.tar.gz", hash = "sha256:cdcd1ba258178f7356d74170c9a81ba4eab6a5d55748bc402a5a2a6fe6be3510", size = 47058631, upload-time = "2026-10-16T22:18:09.718Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8e/45/fbad9c7f5e2f1745c373a434d46072bc84eddfefaa28b852425ac957fda6/pillow-13.0.0-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:44629d4498c9cb6b2051806155f22141df0d8a929c56134b166655232dfb2f8c", size = 4190721, upload-time = "2026-10-16T22:16:27.524Z" }, - { url = "https://files.pythonhosted.org/packages/a2/bf/2074902a259a90692c5dbf132885d21671644cbbd3f690047516243a55db/pillow-13.0.0-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:bf97ba1b5542050e9b0c74e632323b6d8136cc99d64e193abebee22fe6013c76", size = 4290236, upload-time = "2026-10-16T22:16:29.562Z" }, - { url = "https://files.pythonhosted.org/packages/ff/a5/3d1bb93b8ea38d15d21356c1b8e9f50971ac592aaa248f7fb71739094c0e/pillow-13.0.0-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:aa694c49787040943570d61c9b8fc01271bd2a2cfdd55f0d896f8aecdce1a981", size = 3723093, upload-time = "2026-10-16T22:16:31.249Z" }, - { url = "https://files.pythonhosted.org/packages/84/a1/be060135d85187b25fa8b2573786c0e45c62ae27527fff5257675eb97f4a/pillow-13.0.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:49c12af515073e9bf511e4ef736743bf70d531a643668722279314b937742ba4", size = 5572253, upload-time = "2026-10-16T22:16:33.103Z" }, - { url = "https://files.pythonhosted.org/packages/b5/f4/0bf6f1ac9ba7ee2ffe9914aeb410a1035a0c1a3c42f122388d7fc950bc4b/pillow-13.0.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d9b0a2bee3f3e03e9f29057aedb807dbf1e1d96ee1b9d0f422d9db6c8583271d", size = 4852269, upload-time = "2026-10-16T22:16:35.102Z" }, - { url = "https://files.pythonhosted.org/packages/83/16/b603274617b3e6728da3a0133d4fba1ea8d7fbe0c4ea5fef9cbffa398b8b/pillow-13.0.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5933fcea5a5ba64c686f55c886a115cbfd92f2900460784a1e79909fdbafd7b6", size = 6516058, upload-time = "2026-10-16T22:16:37.228Z" }, - { url = "https://files.pythonhosted.org/packages/1d/f2/1371707ce438ba9a43f9d24089828494ec6f9396e28cd7903aa704dd6306/pillow-13.0.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e91c2b1cbeae768c98752e67e717e283d68dae45717b41f2d91eaee648e37b7d", size = 7142574, upload-time = "2026-10-16T22:16:39.172Z" }, - { url = "https://files.pythonhosted.org/packages/b9/a9/62cf68587f910bc87c5951bcb6e9430c7643a65076adf5f18b4bf350e2bd/pillow-13.0.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7f6299e494e46253d40c34653cac369412f44ca53b5f85d848d3f95b1ed0546c", size = 6614295, upload-time = "2026-10-16T22:16:41.117Z" }, - { url = "https://files.pythonhosted.org/packages/6c/c9/9e9b48178649730299be0bafbf5f810a5b6977538f5a92d47e4ba5f5d608/pillow-13.0.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f9bd7d260d986f3553518c5d723544346fa48031f8aef3265d786aeb59845bc0", size = 7224841, upload-time = "2026-10-16T22:16:43.118Z" }, - { url = "https://files.pythonhosted.org/packages/8f/48/c44e91384b1b6c072a714a63d00608acc16a6a968735e478e4be4384e5ac/pillow-13.0.0-cp314-cp314-win32.whl", hash = "sha256:2e31c4e68cffd7b1f116127626813b82fcba19edf6b940a7b0c266d6484d0055", size = 6508694, upload-time = "2026-10-16T22:16:45.883Z" }, - { url = "https://files.pythonhosted.org/packages/50/e2/7ad17747517e4c813dca78257d6bb593466f8231baf1384bc98d92745edb/pillow-13.0.0-cp314-cp314-win_amd64.whl", hash = "sha256:f7cbb7a64e7ec664775195e4f81923bc8a93d9451b963f844e59118f63147a7a", size = 7281830, upload-time = "2026-10-16T22:16:47.725Z" }, - { url = "https://files.pythonhosted.org/packages/41/9f/e5b51f2ce8ef4563215334035403ce61a244f93a2e65c24a940131618f87/pillow-13.0.0-cp314-cp314-win_arm64.whl", hash = "sha256:a6aaa7871aebca081eb5c45f83e845bbd9712ae99fa9286dbfdcee835e63f4b4", size = 2695468, upload-time = "2026-10-16T22:16:49.515Z" }, - { url = "https://files.pythonhosted.org/packages/97/73/a53502f731f6c84805a4ce87a8881e86d07a65aee64728b900d112dc703e/pillow-13.0.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7b999816b5850a35ada08443e6e03945f1f6488144a83afe2d37b240a552aafa", size = 5578196, upload-time = "2026-10-16T22:16:51.291Z" }, - { url = "https://files.pythonhosted.org/packages/21/50/7e949959a56b5559c78ef172b70fc6025706acb289f3f648a15329e4fa6f/pillow-13.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:67871a13cf47275201553b536c025da777f997fa562a19eb798af4ad4b97c18f", size = 4857276, upload-time = "2026-10-16T22:16:53.399Z" }, - { url = "https://files.pythonhosted.org/packages/c9/52/43ded0ac08e77e39d17bd6d894162572201b5e8fdc4a022b14c55a1960a4/pillow-13.0.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e8ca8f69b7c77fa96dad99f8638c1fc4f1c0bed4acaa1baf7a3c31bdb80a0ab1", size = 6607008, upload-time = "2026-10-16T22:16:55.525Z" }, - { url = "https://files.pythonhosted.org/packages/0d/2e/4a57ed216d5bf877387cba4d44c02aa87abe51ad2e06404aa6b597c4da39/pillow-13.0.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6bce6a29c3898191f4c9998e4d4e8f69790caeabde46e54da4bc235ce5bde683", size = 7224281, upload-time = "2026-10-16T22:16:57.738Z" }, - { url = "https://files.pythonhosted.org/packages/5c/3d/2929b248e135752d63aaeabf53f6de797e905b7b6263f9c1e2adade20aba/pillow-13.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:816752a5ef76a80484c8c441c85593b7ea9304f6231fc5776ad1fbff7e773058", size = 6695600, upload-time = "2026-10-16T22:16:59.842Z" }, - { url = "https://files.pythonhosted.org/packages/bc/33/fb0da992cf084ed1ee6ee05e96b6191e35c351eb707ca2df874feb24ce47/pillow-13.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5bdb38ba220b61b4ec31dd1046dea6b016d3c6850a23f43ed8b58332b2fc668e", size = 7300589, upload-time = "2026-10-16T22:17:01.686Z" }, - { url = "https://files.pythonhosted.org/packages/bb/af/c30dd9ef6a3dce84ecb1b2393cac629047a32e66279056c6cb46c3ffd65c/pillow-13.0.0-cp314-cp314t-win32.whl", hash = "sha256:06903aac074d00c988e77a1fc435286dc498506d4ca8d6f69df9babddf4ecc15", size = 6517204, upload-time = "2026-10-16T22:17:03.722Z" }, - { url = "https://files.pythonhosted.org/packages/85/d5/ef3a8ae7609c8ab26570abf18194882534b480d245a9148b895a2262c004/pillow-13.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:9a635068d9793c7f0d929f39f18b58f1bd86c236e8e7b64167b86af9232ebca0", size = 7286695, upload-time = "2026-10-16T22:17:05.543Z" }, - { url = "https://files.pythonhosted.org/packages/84/3c/0861db0c69a03689ebe3e7599d35668a414deb5aaafd8255215e4e62e3d3/pillow-13.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0d1a9cf7bc6917538b20f8d06b9ab2147e542f68849e96fc58eb66bdf6acc692", size = 2698098, upload-time = "2026-10-16T22:17:07.674Z" }, -] - -[[package]] -name = "pluggy" -version = "1.6.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, -] - -[[package]] -name = "protobuf" -version = "7.35.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/10/03/8aeeb7458d22546bf64b5250ca1daeb5ff757d900e8e4a7476c6f0db843e/protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6", size = 433226, upload-time = "2026-06-11T21:55:31.719Z" }, - { url = "https://files.pythonhosted.org/packages/37/4b/dfb89eb0e652a1ff073c39a59fb5e3a83cfe9b57a2c83fa6d78270101767/protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799", size = 328847, upload-time = "2026-06-11T21:55:34.035Z" }, - { url = "https://files.pythonhosted.org/packages/0f/58/dc12f2cd484951524af6e3382c785869b9b3fb5e52ee95ae23add53ee8f9/protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4", size = 344030, upload-time = "2026-06-11T21:55:34.941Z" }, - { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" }, - { url = "https://files.pythonhosted.org/packages/d8/bc/6d6c7ba8709c85f8f2c390b2b118d6fb08a783676a572271851bf45a7d22/protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30", size = 428945, upload-time = "2026-06-11T21:55:37.034Z" }, - { url = "https://files.pythonhosted.org/packages/0a/19/8d0cb6f20a1ef7b18f1c8986ad5783f22f84cce39c6ce9a6e645ea55192e/protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87", size = 439996, upload-time = "2026-06-11T21:55:38.123Z" }, - { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, -] - -[[package]] -name = "psutil" -version = "7.2.2" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/81/69/ef179ab5ca24f32acc1dac0c247fd6a13b501fd5534dbae0e05a1c48b66d/psutil-7.2.2-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:eed63d3b4d62449571547b60578c5b2c4bcccc5387148db46e0c2313dad0ee00", size = 130664, upload-time = "2026-01-28T18:15:09.469Z" }, - { url = "https://files.pythonhosted.org/packages/7b/64/665248b557a236d3fa9efc378d60d95ef56dd0a490c2cd37dafc7660d4a9/psutil-7.2.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:7b6d09433a10592ce39b13d7be5a54fbac1d1228ed29abc880fb23df7cb694c9", size = 131087, upload-time = "2026-01-28T18:15:11.724Z" }, - { url = "https://files.pythonhosted.org/packages/d5/2e/e6782744700d6759ebce3043dcfa661fb61e2fb752b91cdeae9af12c2178/psutil-7.2.2-cp314-cp314t-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1fa4ecf83bcdf6e6c8f4449aff98eefb5d0604bf88cb883d7da3d8d2d909546a", size = 182383, upload-time = "2026-01-28T18:15:13.445Z" }, - { url = "https://files.pythonhosted.org/packages/57/49/0a41cefd10cb7505cdc04dab3eacf24c0c2cb158a998b8c7b1d27ee2c1f5/psutil-7.2.2-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e452c464a02e7dc7822a05d25db4cde564444a67e58539a00f929c51eddda0cf", size = 185210, upload-time = "2026-01-28T18:15:16.002Z" }, - { url = "https://files.pythonhosted.org/packages/dd/2c/ff9bfb544f283ba5f83ba725a3c5fec6d6b10b8f27ac1dc641c473dc390d/psutil-7.2.2-cp314-cp314t-win_amd64.whl", hash = "sha256:c7663d4e37f13e884d13994247449e9f8f574bc4655d509c3b95e9ec9e2b9dc1", size = 141228, upload-time = "2026-01-28T18:15:18.385Z" }, - { url = "https://files.pythonhosted.org/packages/f2/fc/f8d9c31db14fcec13748d373e668bc3bed94d9077dbc17fb0eebc073233c/psutil-7.2.2-cp314-cp314t-win_arm64.whl", hash = "sha256:11fe5a4f613759764e79c65cf11ebdf26e33d6dd34336f8a337aa2996d71c841", size = 136284, upload-time = "2026-01-28T18:15:19.912Z" }, - { url = "https://files.pythonhosted.org/packages/e7/36/5ee6e05c9bd427237b11b3937ad82bb8ad2752d72c6969314590dd0c2f6e/psutil-7.2.2-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:ed0cace939114f62738d808fdcecd4c869222507e266e574799e9c0faa17d486", size = 129090, upload-time = "2026-01-28T18:15:22.168Z" }, - { url = "https://files.pythonhosted.org/packages/80/c4/f5af4c1ca8c1eeb2e92ccca14ce8effdeec651d5ab6053c589b074eda6e1/psutil-7.2.2-cp36-abi3-macosx_11_0_arm64.whl", hash = "sha256:1a7b04c10f32cc88ab39cbf606e117fd74721c831c98a27dc04578deb0c16979", size = 129859, upload-time = "2026-01-28T18:15:23.795Z" }, - { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, - { url = "https://files.pythonhosted.org/packages/63/65/37648c0c158dc222aba51c089eb3bdfa238e621674dc42d48706e639204f/psutil-7.2.2-cp36-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b0726cecd84f9474419d67252add4ac0cd9811b04d61123054b9fb6f57df6e9e", size = 156997, upload-time = "2026-01-28T18:15:27.794Z" }, - { url = "https://files.pythonhosted.org/packages/8e/13/125093eadae863ce03c6ffdbae9929430d116a246ef69866dad94da3bfbc/psutil-7.2.2-cp36-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:fd04ef36b4a6d599bbdb225dd1d3f51e00105f6d48a28f006da7f9822f2606d8", size = 148972, upload-time = "2026-01-28T18:15:29.342Z" }, - { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, - { url = "https://files.pythonhosted.org/packages/b4/90/e2159492b5426be0c1fef7acba807a03511f97c5f86b3caeda6ad92351a7/psutil-7.2.2-cp37-abi3-win_amd64.whl", hash = "sha256:eb7e81434c8d223ec4a219b5fc1c47d0417b12be7ea866e24fb5ad6e84b3d988", size = 137737, upload-time = "2026-01-28T18:15:33.849Z" }, - { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, -] - -[[package]] -name = "pydantic" -version = "2.14.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "annotated-types" }, - { name = "pydantic-core" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/7c/0b/8e10b2e693af8ec54346a14caf36221334775975a747c9ceaa3f8371d96d/pydantic-2.14.1.tar.gz", hash = "sha256:94f478203dd03404682a1ada216965651dd74b1d2d5ffd62e00e0837caab5c26", size = 900362, upload-time = "2026-10-11T18:37:55.396Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/35/ea/a56b9fe5066f3537b7882f77e9c5dfb26d8c2eefdaed9b5fc73d57b4dc22/pydantic-2.14.1-py3-none-any.whl", hash = "sha256:9195d967ec791692a04438115466764fb8b9a27b31f14a760437694f40d6b454", size = 488106, upload-time = "2026-10-11T18:37:53.437Z" }, -] - -[[package]] -name = "pydantic-core" -version = "2.50.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a6/24/af4ec4be49fbc810f35b0bdcacc3d433b56bbfa469fd3bfd4ae116cd9bf1/pydantic_core-2.50.1.tar.gz", hash = "sha256:e50d7b94baac6c7d09927fa5ca5800a0c7ee5015c7fcff65beb3a1931b5a6e09", size = 513515, upload-time = "2026-10-11T18:35:44.82Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/49/33/28b96e81677153715e3eafb9f26663a80841e859fde282a380359b0d3fa1/pydantic_core-2.50.1-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:1cf41f1ae3fa155cf167a72689ad044bcc1e3c97e064123677149bdfb5dafc4a", size = 2224428, upload-time = "2026-10-11T18:33:07.153Z" }, - { url = "https://files.pythonhosted.org/packages/94/40/15c06410c9b7b8da5805d27b64e09bd3f900e986728106db2689dbc51513/pydantic_core-2.50.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4df197990c15b5a37c5a277d131d9f2c67de6133f2e5dafd80d9bba4b99f46f9", size = 2023740, upload-time = "2026-10-11T18:33:08.763Z" }, - { url = "https://files.pythonhosted.org/packages/a1/4e/5eb629f6efc2a27e421d789dd4bfacbb5f6d09c80c28b13d1e87d73163d5/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0036473f5583e6a60e50b8b21651511564277a3f05cc5dab8cf579f552cd5f6c", size = 2049435, upload-time = "2026-10-11T18:33:10.366Z" }, - { url = "https://files.pythonhosted.org/packages/ba/8e/f195aebec49ad12318876ac2368c197a7939f5d52f8e156d2238ef8a4588/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:992c3514ec891fa7858099183e4d64e6bd5a5d4ff452fae29df22faa77a006bb", size = 2172975, upload-time = "2026-10-11T18:33:12.368Z" }, - { url = "https://files.pythonhosted.org/packages/e0/f5/7ad9fb83010cd5ea0948409db105676ed779c4e709e2d3d89b9fe2558794/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:739dc730e6be3bd5ec2f4ab5cfc7eb047cc45fc1497b3bafec74ff2ed07df597", size = 2364197, upload-time = "2026-10-11T18:33:14.244Z" }, - { url = "https://files.pythonhosted.org/packages/ca/fb/bf0aab3e78301d202b82a0322ec968cd703b11fc0623fde9a28708936f62/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:32fad3a91e51b6d2039c572db04a5a873260b399f6bd62c3552671fa7a4a2899", size = 2444259, upload-time = "2026-10-11T18:33:15.79Z" }, - { url = "https://files.pythonhosted.org/packages/45/35/38f6d6564fae57d9b12e5676dfa947e0e7d5c46dbeaa7eb3352261d6d299/pydantic_core-2.50.1-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:42b54c2c90ad348b5e3a85e03e715d572c1fde357ef104cdfe3b03b697a404ea", size = 2159857, upload-time = "2026-10-11T18:33:17.51Z" }, - { url = "https://files.pythonhosted.org/packages/65/a0/fb0a3ca10f139dcf765b2d312d10cf0f65c57c59993229d00ad12b14ceff/pydantic_core-2.50.1-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:2df1ff41884de2bc4b307bafd7c40a691094fad2ff8e767e5b45a319257bcf4e", size = 2233954, upload-time = "2026-10-11T18:33:19.591Z" }, - { url = "https://files.pythonhosted.org/packages/8c/6a/e63842252702aa4ec6e6b7178ca85e541a77459076592c23ebe2d9840b33/pydantic_core-2.50.1-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fe90228920fd8ff2be62622b6bb8a2b11acd65046d50c6b130614b5879605a20", size = 2243876, upload-time = "2026-10-11T18:33:21.393Z" }, - { url = "https://files.pythonhosted.org/packages/39/25/5991cf8318b37e0dfab47b87541619a1df8501a793cba0d978846cba37a7/pydantic_core-2.50.1-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:844b869f118e22a41a091bdcedda8a71bc1b0f62c38d1a0c3211cece47e1d8fc", size = 2277908, upload-time = "2026-10-11T18:33:23.324Z" }, - { url = "https://files.pythonhosted.org/packages/a7/3e/3ee8baaa6cc25a6961c69168cf9ff0f002d56f4e0d541ad6724c18fb61f3/pydantic_core-2.50.1-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:2eb75304506894a281d346220a4f7481a1b8729577c5ed2a05395991966a8396", size = 2449285, upload-time = "2026-10-11T18:33:24.942Z" }, - { url = "https://files.pythonhosted.org/packages/c0/c7/acbec6deac13fe697a80c275a9b6661db62c4d323bac8a47347b1f39c7cd/pydantic_core-2.50.1-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:6b20a4bffabdad0db2927ac034ae3b8a681b1f7a0182f3e60b479ad2fde21ebb", size = 2479013, upload-time = "2026-10-11T18:33:26.925Z" }, - { url = "https://files.pythonhosted.org/packages/10/08/21a3f237b264389f6219d053ee78fc1b5c2fd402c1b1cb2b6cb8b85f9834/pydantic_core-2.50.1-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:99ba9bc2b8062ea0c326a990f7f00e6530c23579de66dd246e72c4cafef950a5", size = 1484400, upload-time = "2026-10-11T18:33:28.693Z" }, - { url = "https://files.pythonhosted.org/packages/09/ce/077a6d262d12ef09108377ac0717f029f420d773ace63cafd6143af75568/pydantic_core-2.50.1-cp314-cp314-win32.whl", hash = "sha256:cf356f70551d40374eaffb1aa63f1eb6d2006681cbd7a9faea173ce0f4dd7cd2", size = 2103222, upload-time = "2026-10-11T18:33:30.326Z" }, - { url = "https://files.pythonhosted.org/packages/14/4c/350a2415209c43d670eb71d3c040f332c04a30583d39e311b05c7ac15762/pydantic_core-2.50.1-cp314-cp314-win_amd64.whl", hash = "sha256:d32f3acc081cc3923386d88f422cde8892335e95f034e0104bb4cf9310d9915f", size = 2171586, upload-time = "2026-10-11T18:33:32.139Z" }, - { url = "https://files.pythonhosted.org/packages/bf/92/9bea6ca96580a0902fed366f064f0829e404f41889b34543279a1162b888/pydantic_core-2.50.1-cp314-cp314-win_arm64.whl", hash = "sha256:bed5163e03b98bc1fa2eb05d74c63d9c5c95d8ed6254985481640fbf5e237dea", size = 2155773, upload-time = "2026-10-11T18:33:33.896Z" }, - { url = "https://files.pythonhosted.org/packages/86/8c/f121f073cf32bdd5cba7e6230b1ce0ac845b0cf43e44dd597d95af272db2/pydantic_core-2.50.1-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:9572c1369e9c9da2d64a7b7992c786d90ff295abc93964cfe3125e4290768070", size = 2234017, upload-time = "2026-10-11T18:33:35.588Z" }, - { url = "https://files.pythonhosted.org/packages/a4/69/2bb2bcd6146cffe98d760a46c42ae71efd5151d9b2f9c9bf6619a3b32083/pydantic_core-2.50.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:2005207aafe1231315718bf6ed5d064a7300fb4772754af35ee72fc68159492e", size = 2001016, upload-time = "2026-10-11T18:33:37.57Z" }, - { url = "https://files.pythonhosted.org/packages/20/b3/fbf854c7d07ec114260c26e9e2071a4381740f9ae09641dbfcbdf2a18c45/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:64f6047f62a6c5ae08d0a6afb035667aa2d97c3d20d69762e034c5ea144d92a5", size = 2023015, upload-time = "2026-10-11T18:33:39.433Z" }, - { url = "https://files.pythonhosted.org/packages/65/20/6de55b2f92cdb614b745c6e9ced639fc4fd7e1e77604825a88bdece6fcbd/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1ef800dd7d85bcdadf4c3076e4c94e43939493558a3b69a1ea830c706d4617bb", size = 2192543, upload-time = "2026-10-11T18:33:41.381Z" }, - { url = "https://files.pythonhosted.org/packages/0c/d9/19e91c94bd5c405945ce2f15526808aea37e162c253160de8ed7bf70b406/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:b0135bcdcaa0f23573f286e4cb5e0fd2962700964ed13df085b85f2b97aeab9e", size = 2388704, upload-time = "2026-10-11T18:33:43.167Z" }, - { url = "https://files.pythonhosted.org/packages/b6/bc/2e24c8415eae1a25ee5a5484946a917123bad2e9d01689a759236928175a/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:0b3a6f334c6a2345ca15318ff894502a90012536404b37c844a976c76c846e0b", size = 2464118, upload-time = "2026-10-11T18:33:44.856Z" }, - { url = "https://files.pythonhosted.org/packages/bd/1f/945b8053cb061c64e102bcaf7bfb9ed740c0bd4349f9bb978a7d4ddff4ab/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:06e01fbbfdb9be777b316a71b6c49efaf4a08b615d0a98d678cda3023f79d019", size = 2150142, upload-time = "2026-10-11T18:33:46.661Z" }, - { url = "https://files.pythonhosted.org/packages/2b/78/96a3e50bf0d64aaae781107eb9335b7529d05a85793ed6e7241c4cd1d931/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:a29a061fec0b4e2d714f277e70a3a18125ecff803f2fea6eade2f2e53711d112", size = 2245577, upload-time = "2026-10-11T18:33:48.574Z" }, - { url = "https://files.pythonhosted.org/packages/7a/5d/a6038a0322232758a6ebfa709f4ca14a60bf5b93940cd0b345056a557da4/pydantic_core-2.50.1-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:f5187624823423e1d1b82b1072ac41dc837389e18d3d0572cc19bbee46cd550a", size = 2263728, upload-time = "2026-10-11T18:33:50.527Z" }, - { url = "https://files.pythonhosted.org/packages/0d/4b/76ded3333a457a9344c4f2d63f2d82d0101654b05178fa4ddfe7d3627674/pydantic_core-2.50.1-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:3e46a9eb0a0901dd6275e6b06ac3a464885ef350ec4121fe486869de8053e4bb", size = 2293630, upload-time = "2026-10-11T18:33:52.362Z" }, - { url = "https://files.pythonhosted.org/packages/a9/00/9eca378335c9c1b72bc779bf6e4c2a82ce784f14ec48a0e87102c9870e05/pydantic_core-2.50.1-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:756d669f04e62ec4148ecfe22be6a4484d9b1181a6ef32e205ebfd200540858b", size = 2468048, upload-time = "2026-10-11T18:33:54.118Z" }, - { url = "https://files.pythonhosted.org/packages/35/ca/e3832e9cf93651251de43c8c5c029ed680a3c1a0a1b17ee26c81bed08cbd/pydantic_core-2.50.1-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:c516cc5367ca3448995d42cb994bf3f4c9002d2a7c22eac9622551269ad1b807", size = 2512159, upload-time = "2026-10-11T18:33:55.99Z" }, - { url = "https://files.pythonhosted.org/packages/4b/f2/773469b5a10a39116a2cb17edaf6d722f017dd8da07161b371465311c542/pydantic_core-2.50.1-cp314-cp314t-win32.whl", hash = "sha256:9d1bed94af6a63835461f3cf7502058eb166c58c4778e11d0f433cfb1bd69e19", size = 2098406, upload-time = "2026-10-11T18:33:58.043Z" }, - { url = "https://files.pythonhosted.org/packages/ee/42/0bb74f8f25204b259b11ab7c12dc7f180893b6bd706118b385147fb6efd5/pydantic_core-2.50.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c8dce1f1e0e5358b682a6ad3fa5e31b31d4560997b8e61417e9217c8d60f8a0c", size = 2148481, upload-time = "2026-10-11T18:33:59.913Z" }, - { url = "https://files.pythonhosted.org/packages/47/0d/d801646c9679a4e630e15cf521d4108b93854b42a6e6e02391bd4c6b1095/pydantic_core-2.50.1-cp314-cp314t-win_arm64.whl", hash = "sha256:ceff0acc940be2715bd6ad17b24c0e5304abf44f6efd0f81ee8499e640f9dc86", size = 2174972, upload-time = "2026-10-11T18:34:01.875Z" }, -] - -[[package]] -name = "pydantic-settings" -version = "2.16.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pydantic" }, - { name = "python-dotenv" }, - { name = "typing-extensions" }, - { name = "typing-inspection" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/05/ca/7817b4bf0c4b15c6a6ab8cdf4cb96c268335a8ca2f525de93b6ebe4c7d28/pydantic_settings-2.16.1.tar.gz", hash = "sha256:aff602a1886338c2df3d1e958e2193a8dc33cec75b91f202e83fd98484b4b0da", size = 303436, upload-time = "2026-10-20T11:37:19.024Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3c/74/fdf926afe145d8629c4b2ff57f4ea1d64ed7abc4d84d97463b2517a2547e/pydantic_settings-2.16.1-py3-none-any.whl", hash = "sha256:30f5313340da4e8f1b440d1910b92a5bc3ccbb6b4ba70ec2b8395bc921f418f5", size = 80771, upload-time = "2026-10-20T11:37:17.298Z" }, -] - -[[package]] -name = "pydot" -version = "4.0.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pyparsing" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/50/35/b17cb89ff865484c6a20ef46bf9d95a5f07328292578de0b295f4a6beec2/pydot-4.0.1.tar.gz", hash = "sha256:c2148f681c4a33e08bf0e26a9e5f8e4099a82e0e2a068098f32ce86577364ad5", size = 162594, upload-time = "2025-06-17T20:09:56.454Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7e/32/a7125fb28c4261a627f999d5fb4afff25b523800faed2c30979949d6facd/pydot-4.0.1-py3-none-any.whl", hash = "sha256:869c0efadd2708c0be1f916eb669f3d664ca684bc57ffb7ecc08e70d5e93fee6", size = 37087, upload-time = "2025-06-17T20:09:55.25Z" }, -] - -[[package]] -name = "pygments" -version = "2.21.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, -] - -[[package]] -name = "pyparsing" -version = "3.3.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e4/11/b213bebff182584360cb8d17c72c1677fec5c5c228de439e63bcf8ab1c8f/pyparsing-3.3.3.tar.gz", hash = "sha256:928ae7e20211f3b6f3915a72f06a0cfd29ab9d24279dd6346b6b1a7146397d36", size = 1050487, upload-time = "2026-09-20T20:59:05.609Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/38/bb/d215ee7c73b61497b28a5503f9f53523f294fcc936762b7caf90e0c1c2b5/pyparsing-3.3.3-py3-none-any.whl", hash = "sha256:ece8c00a69cf01b45d0b1dedabb469c90d8caf996d4fda40f147627a122849a4", size = 126420, upload-time = "2026-09-20T20:59:04.025Z" }, -] - -[[package]] -name = "pytest" -version = "9.1.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, - { name = "iniconfig" }, - { name = "packaging" }, - { name = "pluggy" }, - { name = "pygments" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, -] - -[[package]] -name = "python-dateutil" -version = "2.9.0.post0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "six" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, -] - -[[package]] -name = "python-dotenv" -version = "1.2.4" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/74/26/2fbeedb218a787a5eea551c7532cac4e009f83d689dd2faa0d0353473f86/python_dotenv-1.2.4.tar.gz", hash = "sha256:f0d53e69935a851c0dcc78f3ab7aaccd8cabef0b92382b576b824212902873c0", size = 60824, upload-time = "2026-10-01T05:36:10Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/60/d1/38f3a3405989a89ac18390803e70c6ad7c7760da4f9b83cbeca0c44a0c72/python_dotenv-1.2.4-py3-none-any.whl", hash = "sha256:42269a8a5b3fd54ffa6f3d84b18abed50064717576b4ecf03dc4a55d8aa04fdc", size = 23266, upload-time = "2026-10-01T05:36:08.633Z" }, -] - -[[package]] -name = "pyyaml" -version = "6.0.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/8c/f4bd7f6465179953d3ac9bc44ac1a8a3e6122cf8ada906b4f96c60172d43/pyyaml-6.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:8d1fab6bb153a416f9aeb4b8763bc0f22a5586065f86f7664fc23339fc1c1fac", size = 181814, upload-time = "2025-09-25T21:32:35.712Z" }, - { url = "https://files.pythonhosted.org/packages/bd/9c/4d95bb87eb2063d20db7b60faa3840c1b18025517ae857371c4dd55a6b3a/pyyaml-6.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:34d5fcd24b8445fadc33f9cf348c1047101756fd760b4dacb5c3e99755703310", size = 173809, upload-time = "2025-09-25T21:32:36.789Z" }, - { url = "https://files.pythonhosted.org/packages/92/b5/47e807c2623074914e29dabd16cbbdd4bf5e9b2db9f8090fa64411fc5382/pyyaml-6.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:501a031947e3a9025ed4405a168e6ef5ae3126c59f90ce0cd6f2bfc477be31b7", size = 766454, upload-time = "2025-09-25T21:32:37.966Z" }, - { url = "https://files.pythonhosted.org/packages/02/9e/e5e9b168be58564121efb3de6859c452fccde0ab093d8438905899a3a483/pyyaml-6.0.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:b3bc83488de33889877a0f2543ade9f70c67d66d9ebb4ac959502e12de895788", size = 836355, upload-time = "2025-09-25T21:32:39.178Z" }, - { url = "https://files.pythonhosted.org/packages/88/f9/16491d7ed2a919954993e48aa941b200f38040928474c9e85ea9e64222c3/pyyaml-6.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c458b6d084f9b935061bc36216e8a69a7e293a2f1e68bf956dcd9e6cbcd143f5", size = 794175, upload-time = "2025-09-25T21:32:40.865Z" }, - { url = "https://files.pythonhosted.org/packages/dd/3f/5989debef34dc6397317802b527dbbafb2b4760878a53d4166579111411e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:7c6610def4f163542a622a73fb39f534f8c101d690126992300bf3207eab9764", size = 755228, upload-time = "2025-09-25T21:32:42.084Z" }, - { url = "https://files.pythonhosted.org/packages/d7/ce/af88a49043cd2e265be63d083fc75b27b6ed062f5f9fd6cdc223ad62f03e/pyyaml-6.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5190d403f121660ce8d1d2c1bb2ef1bd05b5f68533fc5c2ea899bd15f4399b35", size = 789194, upload-time = "2025-09-25T21:32:43.362Z" }, - { url = "https://files.pythonhosted.org/packages/23/20/bb6982b26a40bb43951265ba29d4c246ef0ff59c9fdcdf0ed04e0687de4d/pyyaml-6.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:4a2e8cebe2ff6ab7d1050ecd59c25d4c8bd7e6f400f5f82b96557ac0abafd0ac", size = 156429, upload-time = "2025-09-25T21:32:57.844Z" }, - { url = "https://files.pythonhosted.org/packages/f4/f4/a4541072bb9422c8a883ab55255f918fa378ecf083f5b85e87fc2b4eda1b/pyyaml-6.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:93dda82c9c22deb0a405ea4dc5f2d0cda384168e466364dec6255b293923b2f3", size = 143912, upload-time = "2025-09-25T21:32:59.247Z" }, - { url = "https://files.pythonhosted.org/packages/7c/f9/07dd09ae774e4616edf6cda684ee78f97777bdd15847253637a6f052a62f/pyyaml-6.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:02893d100e99e03eda1c8fd5c441d8c60103fd175728e23e431db1b589cf5ab3", size = 189108, upload-time = "2025-09-25T21:32:44.377Z" }, - { url = "https://files.pythonhosted.org/packages/4e/78/8d08c9fb7ce09ad8c38ad533c1191cf27f7ae1effe5bb9400a46d9437fcf/pyyaml-6.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:c1ff362665ae507275af2853520967820d9124984e0f7466736aea23d8611fba", size = 183641, upload-time = "2025-09-25T21:32:45.407Z" }, - { url = "https://files.pythonhosted.org/packages/7b/5b/3babb19104a46945cf816d047db2788bcaf8c94527a805610b0289a01c6b/pyyaml-6.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6adc77889b628398debc7b65c073bcb99c4a0237b248cacaf3fe8a557563ef6c", size = 831901, upload-time = "2025-09-25T21:32:48.83Z" }, - { url = "https://files.pythonhosted.org/packages/8b/cc/dff0684d8dc44da4d22a13f35f073d558c268780ce3c6ba1b87055bb0b87/pyyaml-6.0.3-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:a80cb027f6b349846a3bf6d73b5e95e782175e52f22108cfa17876aaeff93702", size = 861132, upload-time = "2025-09-25T21:32:50.149Z" }, - { url = "https://files.pythonhosted.org/packages/b1/5e/f77dc6b9036943e285ba76b49e118d9ea929885becb0a29ba8a7c75e29fe/pyyaml-6.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:00c4bdeba853cc34e7dd471f16b4114f4162dc03e6b7afcc2128711f0eca823c", size = 839261, upload-time = "2025-09-25T21:32:51.808Z" }, - { url = "https://files.pythonhosted.org/packages/ce/88/a9db1376aa2a228197c58b37302f284b5617f56a5d959fd1763fb1675ce6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:66e1674c3ef6f541c35191caae2d429b967b99e02040f5ba928632d9a7f0f065", size = 805272, upload-time = "2025-09-25T21:32:52.941Z" }, - { url = "https://files.pythonhosted.org/packages/da/92/1446574745d74df0c92e6aa4a7b0b3130706a4142b2d1a5869f2eaa423c6/pyyaml-6.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:16249ee61e95f858e83976573de0f5b2893b3677ba71c9dd36b9cf8be9ac6d65", size = 829923, upload-time = "2025-09-25T21:32:54.537Z" }, - { url = "https://files.pythonhosted.org/packages/f0/7a/1c7270340330e575b92f397352af856a8c06f230aa3e76f86b39d01b416a/pyyaml-6.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4ad1906908f2f5ae4e5a8ddfce73c320c2a1429ec52eafd27138b7f1cbe341c9", size = 174062, upload-time = "2025-09-25T21:32:55.767Z" }, - { url = "https://files.pythonhosted.org/packages/f1/12/de94a39c2ef588c7e6455cfbe7343d3b2dc9d6b6b2f40c4c6565744c873d/pyyaml-6.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:ebc55a14a21cb14062aa4162f906cd962b28e2e9ea38f9b4391244cd8de4ae0b", size = 149341, upload-time = "2025-09-25T21:32:56.828Z" }, -] - -[[package]] -name = "requests" -version = "2.34.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "certifi" }, - { name = "charset-normalizer" }, - { name = "idna" }, - { name = "urllib3" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, -] - -[[package]] -name = "ruff" -version = "0.16.3" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/61/b3/3213589383f8f1b3938781bd1278713f6d18621a14992b3e81fefb8a5ef9/ruff-0.16.3.tar.gz", hash = "sha256:e76d33a347661a84b5be6d043d0347fdc745dfdcf825a8f4fed64b5e26eebdf2", size = 4891904, upload-time = "2026-08-13T15:17:13.381Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bf/96/493770daebd68c0a67f1549fdf519f53be51fc435186c0585bcc272fd76c/ruff-0.16.3-py3-none-linux_armv6l.whl", hash = "sha256:0c5710e247a58a4521e66e124ba9a74655b414f61ba3a2e9e3811e11098f48f7", size = 10902799, upload-time = "2026-08-13T15:16:27.382Z" }, - { url = "https://files.pythonhosted.org/packages/5e/e6/2becf3942fddc29a29b8df47691d456fb1085391a694f74d84513251418c/ruff-0.16.3-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:fe155130631a2471fd2e14a7a664a4dfbd7194b8229c3d7b2a40b21178639081", size = 11135539, upload-time = "2026-08-13T15:16:30.87Z" }, - { url = "https://files.pythonhosted.org/packages/3e/1e/4b8b72f0d006dbf19326aa99f9ca0ee2ff374187c4d301cf529a51aa06fe/ruff-0.16.3-py3-none-macosx_11_0_arm64.whl", hash = "sha256:e2ed719e14aa64d895c2ee922594a90a43c861a93f0575a95ff8c47cdbd13eb9", size = 10475095, upload-time = "2026-08-13T15:16:33.259Z" }, - { url = "https://files.pythonhosted.org/packages/92/32/2201fa49ba1f6c101ee321e83f051ac7a4b8d07b0ef6b4d3f2772b302275/ruff-0.16.3-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:9e0b1da805eb043654645d74d5de1e5ce2edc686e40790d2b86f56d71cc06a84", size = 10668771, upload-time = "2026-08-13T15:16:35.65Z" }, - { url = "https://files.pythonhosted.org/packages/c3/66/4afc5c8363bd04d45effce1b7c8713ca037d7a6740b7451a2403a6e3a972/ruff-0.16.3-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a37bdea0bbe21780f590bf437d6412c8c4e1b6cd010f91a65c2c40c5e5f5f870", size = 10699568, upload-time = "2026-08-13T15:16:38.195Z" }, - { url = "https://files.pythonhosted.org/packages/53/fd/c67d246bf36bf1698551c56de39e95cd07f70e64433e0098e6267d77061b/ruff-0.16.3-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:09571e6d1288ed9be475207a3ac04ada404f1cd898104be0f6ab8d7df438575b", size = 11499365, upload-time = "2026-08-13T15:16:40.623Z" }, - { url = "https://files.pythonhosted.org/packages/67/0b/00ecbceb99a263af7b12f6f05ac3c92bc47b905e91adc3f207a836e3bc01/ruff-0.16.3-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2c18c5a101eb540010638cc1ff3c84944d3adb3df62b8d98ca8f22ba484d3413", size = 12311728, upload-time = "2026-08-13T15:16:43.564Z" }, - { url = "https://files.pythonhosted.org/packages/54/b2/b7b3bb54f4d3f7db504e476ad4ab8de530dceebe2c061384b2757ee419e8/ruff-0.16.3-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:8457c44f15033c85ddbb77b15d451df9e24e4bd03b628396dd3610cedc3b8f82", size = 11699896, upload-time = "2026-08-13T15:16:46.209Z" }, - { url = "https://files.pythonhosted.org/packages/c7/30/4c468429ac195addc5ee1b717b6ab1b66632786737ca3b2ed3443fb0c26a/ruff-0.16.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:294b95c4ae0cda9388525c2047778aa758d6b8d4bb876fd4e9eaa3ebc92343eb", size = 11058736, upload-time = "2026-08-13T15:16:48.823Z" }, - { url = "https://files.pythonhosted.org/packages/43/67/7a113cdaddf24b64d7f75b1242a99d04c82fcef4f6921fdbb832beaffb5f/ruff-0.16.3-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:3d0c7c40c87c2a820509c31ba007968da6e1306468c067b2d82fbfdbcd0e8474", size = 11586911, upload-time = "2026-08-13T15:16:51.913Z" }, - { url = "https://files.pythonhosted.org/packages/f1/c1/2e66f24c0f3ead25a5e660111778685e505e5da353c82802bf49f0cbe7b9/ruff-0.16.3-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:9f738c0fdfa8eed0b2ce7fb27ee7258208a92a68d7949e62aa15164bc7b389da", size = 10954265, upload-time = "2026-08-13T15:16:54.763Z" }, - { url = "https://files.pythonhosted.org/packages/c2/ba/4cee23bf52cba9a058d3726de623624daf50ef9638868edd86f4126157f6/ruff-0.16.3-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:fb785f0be25abe69d320415cd4f833b59e17ba7613d9ba6a958023b6bceb0a50", size = 10709886, upload-time = "2026-08-13T15:16:57.339Z" }, - { url = "https://files.pythonhosted.org/packages/82/df/7da7194fa5d9dc0a285f7e6fa5a4722e7c63faac0b45b614ded9314363a1/ruff-0.16.3-py3-none-musllinux_1_2_i686.whl", hash = "sha256:c5536e3acfbf9563085aa2be7b13c629c3077e902afc5b941ac44024dbb9f506", size = 11210392, upload-time = "2026-08-13T15:17:00.171Z" }, - { url = "https://files.pythonhosted.org/packages/35/85/7795f6e817af050e7517bf3e7aa9b061cce70ef33d280aad902c956c1ecf/ruff-0.16.3-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:a2d85c02f9b8e165d85e6779184d38c4132de12603dab59c51c28e22584f9e4d", size = 11626910, upload-time = "2026-08-13T15:17:03.299Z" }, - { url = "https://files.pythonhosted.org/packages/78/9b/475b927cf27a5cbbda3c7bafb69ed6ff77e1d7923d5d85f17c2749d7ae32/ruff-0.16.3-py3-none-win32.whl", hash = "sha256:388cdf2166642bd9b13d52b5932d3170f34f8abed7e8d9a855f1d84b83645a0a", size = 10931415, upload-time = "2026-08-13T15:17:05.726Z" }, - { url = "https://files.pythonhosted.org/packages/b2/99/e2a2bfc4fbf0a1e8a916bc9ebe6fe6c58cc34c28e0ffc6ce281d572d1c2e/ruff-0.16.3-py3-none-win_amd64.whl", hash = "sha256:e80a7d69ca2a6d1c4d352ec91458cdca6e56c83cdbcabd93e4abe1e53591d948", size = 11445993, upload-time = "2026-08-13T15:17:08.353Z" }, - { url = "https://files.pythonhosted.org/packages/69/3e/4132e539aed78c148854d4997a2685b0ed4dc4e87110b59ce528564e184e/ruff-0.16.3-py3-none-win_arm64.whl", hash = "sha256:b8ca152da82c1acc1fa8d5874b15951935f0eef46f10e6954c83859011b6178a", size = 11399302, upload-time = "2026-08-13T15:17:10.908Z" }, -] - -[[package]] -name = "safetensors" -version = "0.8.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, - { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, - { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, - { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, - { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, - { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, - { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, - { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, - { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, - { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, - { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, - { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, - { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, - { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, - { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, - { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, -] - -[[package]] -name = "scikit-learn" -version = "1.9.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "joblib" }, - { name = "narwhals" }, - { name = "numpy" }, - { name = "scipy" }, - { name = "threadpoolctl" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d2/eb/eaf5e07fcc0da7149b0e084f24e54edd7441b9a89ce7e034032ae97fe3a0/scikit_learn-1.9.1.tar.gz", hash = "sha256:629cada3e33e2b9bf376cdc7614a47a4140b8aedc1d836579e359736fbd82977", size = 7786908, upload-time = "2026-09-10T18:34:04.679Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/64/e3/b58e45082dcf3dcf0eb1192ee03545ec43d8c98441dfe20e88afca8442ce/scikit_learn-1.9.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:d5d117952769b563067656784e03c75a2d8235a7a05cf7fffa78a311e75aac08", size = 8747414, upload-time = "2026-09-10T18:33:10.698Z" }, - { url = "https://files.pythonhosted.org/packages/9b/ed/d68115577c8b42b0442ebd8180945d4008880a33176094640e00b585128e/scikit_learn-1.9.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:8893bc6331f60f18d4ac75e12ed356e2dcf6a564bf767918b5b7ca54c8c8be49", size = 8279175, upload-time = "2026-09-10T18:33:12.846Z" }, - { url = "https://files.pythonhosted.org/packages/c4/6c/06c7eb61a438e389cbf3f7210897069bec5a883dbe03c189ae792781e11a/scikit_learn-1.9.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b5492cf2df5226691c32611de8734bcf42148c6547ae53c7f4e6b847793addc0", size = 8881706, upload-time = "2026-09-10T18:33:14.741Z" }, - { url = "https://files.pythonhosted.org/packages/86/4e/0bab75490ca4b85fad8388739c7ebc71d9db553f8c69e39943ee8db0aaae/scikit_learn-1.9.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:993d332ff80e62efae9e39603b7e872297c418d780f01a01855269a3489c950f", size = 9152030, upload-time = "2026-09-10T18:33:17.436Z" }, - { url = "https://files.pythonhosted.org/packages/9a/13/31c6f8ba1b7eecef9dd9558576c752d2ec5785fd0e456bb9fc59305be23f/scikit_learn-1.9.1-cp314-cp314-win_amd64.whl", hash = "sha256:ca9051447455dae341d4d591eece7deb2d8e3d1020298fc87a81fc51e4da8f53", size = 8377024, upload-time = "2026-09-10T18:33:19.941Z" }, - { url = "https://files.pythonhosted.org/packages/62/e6/6d3cb8a45f5228f915acd66b819dd6b8232ccbe24532f51d278e3991df31/scikit_learn-1.9.1-cp314-cp314-win_arm64.whl", hash = "sha256:90de6573f733a9fb79476ff1371af52a397d41c8b35f9146e20923db010d67b6", size = 8014196, upload-time = "2026-09-10T18:33:22.126Z" }, - { url = "https://files.pythonhosted.org/packages/66/6e/6befb2d5961490d18d9dbc16a5df37aa121d08bc9893316a6a363977a903/scikit_learn-1.9.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7b5cad1624de8b75e5b9ccb7b0ce1ff1d01306340a3efc56d5529c5ba92392eb", size = 9066066, upload-time = "2026-09-10T18:33:24.396Z" }, - { url = "https://files.pythonhosted.org/packages/cb/18/11271f2f7db337db01f598e358721b1e83989407131272e5dd64214c28a8/scikit_learn-1.9.1-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:d137ce8a6142029fb5c35bd82f470c40cd9e760e5e2f7694b362c497c4ab3fa2", size = 8647654, upload-time = "2026-09-10T18:33:26.649Z" }, - { url = "https://files.pythonhosted.org/packages/e7/04/9c15d201e1b6a2e81b8215865df7646c5a360f560831769c8dc92ac1ab9a/scikit_learn-1.9.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:66f852f7325b5070bc28329005aca76055a2def78faac039548ae889aeaa45a6", size = 8889070, upload-time = "2026-09-10T18:33:28.95Z" }, - { url = "https://files.pythonhosted.org/packages/1a/5a/4cb6c85160af4a639e87a3b7bf8b1c25cfc3b504c5af710ca416a6dcfc5f/scikit_learn-1.9.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:748bcb0a4cc04aec470652c9e5ec68450948e867387e7dfade647107ade68d25", size = 9131227, upload-time = "2026-09-10T18:33:31.008Z" }, - { url = "https://files.pythonhosted.org/packages/00/0e/361440972ae3d19b90ea88a84791138a51de0e8432a770ba741b2c8d9ced/scikit_learn-1.9.1-cp314-cp314t-win_amd64.whl", hash = "sha256:38cd925e893e5539be704d5edc64dbe081aacdab6b89d8c2977c1f6a7a453ce5", size = 8683001, upload-time = "2026-09-10T18:33:33.188Z" }, - { url = "https://files.pythonhosted.org/packages/e2/8f/a9f405c5c0e2df6f343a871b40c97fb32969e3ccc38e3033dd118f3c261e/scikit_learn-1.9.1-cp314-cp314t-win_arm64.whl", hash = "sha256:b01e5b01735d38474127ca3f49319b592506225a87793b27559816b5c75cea39", size = 8257996, upload-time = "2026-09-10T18:33:35.343Z" }, -] - -[[package]] -name = "scipy" -version = "1.18.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "numpy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/7e/74/66de6258867beb2ef08f35f9f2ac017a52cacd5081714d239ff1a442d458/scipy-1.18.1.tar.gz", hash = "sha256:52c4b7422442aba924d03ad4019852b08a92e64ea187b933135687bfe2747307", size = 30781235, upload-time = "2026-08-21T23:28:50.599Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/06/d5/d8eb4e280ddb56a4ab2c6f02ee49b56b23f6e977cf0802fd6d68dbef14f5/scipy-1.18.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:83de5453a7799afc9048b4616bd085cef126e36412f0ea2f6370c36a2a3a51e7", size = 31090936, upload-time = "2026-08-21T23:25:28.686Z" }, - { url = "https://files.pythonhosted.org/packages/2a/49/59ea385dc3a62ff498ddf3cfff7c2b41b0f9f9d3c4122b3f1dcb6d6327fe/scipy-1.18.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:9554bcc6d715ee87a633a3cc8e7703c6628b100dd29cb8a2efc4c0533c7ff729", size = 28725221, upload-time = "2026-08-21T23:25:33.244Z" }, - { url = "https://files.pythonhosted.org/packages/70/e8/6b0c288c50942d78193696c9f15f9a0874f5178aa0ddf40f83d9924b3e8d/scipy-1.18.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:011413b7426b75012840e35649e00fe0a2c3bae89fed433876e3a99251572efc", size = 20466839, upload-time = "2026-08-21T23:25:37.516Z" }, - { url = "https://files.pythonhosted.org/packages/4b/e0/54fd3793c729e3b936782f181b59cbb1205bf250ab605a16cb1ba61cdd5e/scipy-1.18.1-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:88f0e784020649f88ea48c9f5ddfa403bf9205820667c0914740b392035afb82", size = 23089121, upload-time = "2026-08-21T23:25:42.019Z" }, - { url = "https://files.pythonhosted.org/packages/0b/56/030af62bea3cf878e0028515dff78c123b01633606a879b63f42d2db99cc/scipy-1.18.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2d3ab0e8c69a17dd3559eab8cbb88f258e285c94d572c2719033f90f83290c89", size = 34053851, upload-time = "2026-08-21T23:25:47.998Z" }, - { url = "https://files.pythonhosted.org/packages/6b/89/2a844506d49651e9aa1af6ef95b6bd8031cb1d5a4375edec6155037e04cf/scipy-1.18.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ac0333bdf38309aa3dcbe7e3fa7ea29e7a2c37c6ea306a757b700ded8e4596ad", size = 35329183, upload-time = "2026-08-21T23:25:53.522Z" }, - { url = "https://files.pythonhosted.org/packages/eb/56/c7370c3640e92ac9613cbf26cb3f729f9b12ddf1727b55b94b53b24d6f48/scipy-1.18.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:911de823097db8b63f034299d12662db93344e6ffa0b881cbb57748974b70168", size = 35672551, upload-time = "2026-08-21T23:25:59.387Z" }, - { url = "https://files.pythonhosted.org/packages/24/16/ec8536f351421f8bf60a1120930638f83790f4710b8230446aca3d6159d4/scipy-1.18.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:95298364e251be3e60249facbeeca03631d3bb7584f85879516ec55ac717b81f", size = 37469416, upload-time = "2026-08-21T23:26:05.432Z" }, - { url = "https://files.pythonhosted.org/packages/52/94/d73da0d28f16c45bb9b0a5691b91610b0275c5ef0eb5e43c87cf2dc1bf31/scipy-1.18.1-cp314-cp314-win_amd64.whl", hash = "sha256:78a0d7c918e74a232394117160e7e3db503377572a45bcef8826e4ab8a35feba", size = 37362755, upload-time = "2026-08-21T23:26:11.366Z" }, - { url = "https://files.pythonhosted.org/packages/89/25/e996e4dc74e10e227b1e14db5eaf6608bb6dd33884a64851c38f18dd4249/scipy-1.18.1-cp314-cp314-win_arm64.whl", hash = "sha256:cbf38d043c1aa4ab306e1ada6ab6eddacc3322a20b7af1b30bc93254b366fe09", size = 25036090, upload-time = "2026-08-21T23:26:15.887Z" }, - { url = "https://files.pythonhosted.org/packages/fa/c9/c00213f92309d753b48903e6a451b87eb52ff5b7a16e789d1568bbf221c4/scipy-1.18.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:0fcb3c93519f27bb4f0c4b0f7802cdcaca7fcf93267b75edda2e9f4e8a55cbd7", size = 31485550, upload-time = "2026-08-21T23:26:20.776Z" }, - { url = "https://files.pythonhosted.org/packages/74/b2/e3067c487982d4eeab2938928529410370c06fea84a4d3f4925e7d96647d/scipy-1.18.1-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:ddef79fb382df40104a19bb7151b3b23e57c1778fcf857c71ceecd9bd264513f", size = 29174642, upload-time = "2026-08-21T23:26:25.395Z" }, - { url = "https://files.pythonhosted.org/packages/d5/ab/374c9fe2d1ec014e576c781a4b5d8e1ba340e8f6b4638c16f711d2b194f0/scipy-1.18.1-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:0e82073ecc7acc6436fac4b31674109c7e1d3e596789767eda01258a8c9e8123", size = 20916357, upload-time = "2026-08-21T23:26:30.112Z" }, - { url = "https://files.pythonhosted.org/packages/90/38/223915c88a17317cafbf8ca2a42b11c265a9fb1e804aa665544132b5fe8a/scipy-1.18.1-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:8bcf3c1ba5d6456e2effd30fcbd3459b044d683fcdac79a2e6830f0bdf7de487", size = 23482611, upload-time = "2026-08-21T23:26:34.846Z" }, - { url = "https://files.pythonhosted.org/packages/c4/d1/db0948da8ca57a80b36520ef0a768b967d99f3af65f4b6f1bf6362ad4dd4/scipy-1.18.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cfbf154f2ba187f2ed6cce2639efff7d105f1140573642c0161615b6d91d6a87", size = 34143202, upload-time = "2026-08-21T23:26:40.4Z" }, - { url = "https://files.pythonhosted.org/packages/87/53/39d046cc7574ed6acacb6bd5723e220107ece80bff12faaf3efc4ddeede4/scipy-1.18.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a1d33a7836f7ddc1993427966a0823468ec41bcbdb1a9f9942d1d7e57f803ba3", size = 35380876, upload-time = "2026-08-21T23:26:46.1Z" }, - { url = "https://files.pythonhosted.org/packages/f9/da/32e0e799d875a85ca57d9bde6c78148afcc0e38276df683d95854eadc8c3/scipy-1.18.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:7f4b8bc363b6d65ee2152bec57568e3c52639bb34c46057b09857a307ed5e21d", size = 35770885, upload-time = "2026-08-21T23:26:51.533Z" }, - { url = "https://files.pythonhosted.org/packages/88/2e/f97a666d362fee68b18f41c9c30ed502ca5c98b549749bfcb52a8b74d1eb/scipy-1.18.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:11c423f1049c5755ad4409af52a9ada1cff96fe9b50795d4af3619f292901239", size = 37525424, upload-time = "2026-08-21T23:26:56.751Z" }, - { url = "https://files.pythonhosted.org/packages/ca/d5/a9e765a84654ebba8479a1fd1b059ced1af72b168a3b2a3a46540ea38d20/scipy-1.18.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c24acac1e18912761c4700239bbc1fd32f615af690f1584d49b35859be51324d", size = 37416961, upload-time = "2026-08-21T23:27:01.546Z" }, - { url = "https://files.pythonhosted.org/packages/ee/16/e79e0d1c63ef698879d85439d37e9fb434e3b804e506a6991038d086ebd9/scipy-1.18.1-cp314-cp314t-win_arm64.whl", hash = "sha256:9f2897bf7737392ad0d5213ea7b6add72a4edf5679b3153106aeb88b6507b3b9", size = 25331848, upload-time = "2026-08-21T23:27:05.884Z" }, -] - -[[package]] -name = "setuptools" -version = "84.2.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/70/b3/c19426f783614d625699a4b247faf628ea9ab1db228674491767b417f9e1/setuptools-84.2.1.tar.gz", hash = "sha256:3ff8b3b5492a0aa66960f4419365a86b84e013070d4313e4d4a6894f5d64d6bc", size = 1230289, upload-time = "2026-10-20T19:01:37.742Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ef/55/6dc6d6bb953bdfc52519a0059592c52b6e3b8ad51cf8256dffceb55b0610/setuptools-84.2.1-py3-none-any.whl", hash = "sha256:ca140b45d5325a9478f880e351484f62fef15f336462e449ca0c03c873d589c2", size = 876263, upload-time = "2026-10-20T19:01:35.678Z" }, -] - -[[package]] -name = "six" -version = "1.17.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, -] - -[[package]] -name = "skrub" -version = "0.10.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cloudpickle" }, - { name = "jinja2" }, - { name = "joblib" }, - { name = "matplotlib" }, - { name = "numpy" }, - { name = "pandas" }, - { name = "pydot" }, - { name = "requests" }, - { name = "scikit-learn" }, - { name = "scipy" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a3/a6/f4d0b7240ebf546d3a77f212958b930282cccb0a53d66b9b0e1639ce8230/skrub-0.10.1.tar.gz", hash = "sha256:39f58492f6889c5563a654478cd523771aaa8dab5a42ad2ce9f430efb3881f53", size = 4205317, upload-time = "2026-09-01T09:33:37.146Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0e/2d/492a8b0d0d8b344918312a6416c537e54081ca5d0aaf19d085a188905e41/skrub-0.10.1-py3-none-any.whl", hash = "sha256:553c6f9b19a2cb97bba776a28255e7e024e77d4f50d7c6c70b4805f4e668a3cc", size = 595919, upload-time = "2026-09-01T09:33:35.208Z" }, -] - -[[package]] -name = "sympy" -version = "1.14.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "mpmath" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, -] - -[[package]] -name = "tabdpt" -version = "1.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "faiss-cpu" }, - { name = "huggingface-hub" }, - { name = "numpy" }, - { name = "omegaconf" }, - { name = "safetensors" }, - { name = "scikit-learn" }, - { name = "scipy" }, - { name = "torch", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "tqdm" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/68/f0/cda7df9f41705c162d075255fec8e9e53d63aa9eee0bb419a01a56fe8f53/tabdpt-1.3.1.tar.gz", hash = "sha256:42e3e5c13ededb3ff7e6ab128d6491de5537642f9bba7a66da54701bdb7b99ff", size = 33832, upload-time = "2026-09-21T18:22:02.147Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/1c/fc711dd8ec3bcf15c3c3c9320a35cc2f798c4f3de3bf7e59aef6e7b7263d/tabdpt-1.3.1-py3-none-any.whl", hash = "sha256:b3ad31443cf20d6ee0d2d9398c61f09d3d2e7b796fb2f9346a716ddf576e933f", size = 27204, upload-time = "2026-09-21T18:22:01.192Z" }, -] - -[[package]] -name = "tabicl" -version = "2.2.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "einops" }, - { name = "huggingface-hub" }, - { name = "numpy" }, - { name = "psutil" }, - { name = "scikit-learn" }, - { name = "scipy" }, - { name = "torch", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "tqdm" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d4/81/ee11ac72b33c2b87f7337c99238fd2700128791edceb558b63bb30556b73/tabicl-2.2.0.tar.gz", hash = "sha256:dfc41f2acd9603e4402cedae9e42f9922cfb8cfc9fcf7901d773e26609e47e21", size = 268084, upload-time = "2026-09-02T15:13:38.59Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/77/22/bbac2df399b8ed5a699bde38496ebb129a045c5d160fe52786727de932f3/tabicl-2.2.0-py3-none-any.whl", hash = "sha256:c80b8ec2719be26bc0bf1dd758a2361f58019d8f7877038b785529a425789ce4", size = 310448, upload-time = "2026-09-02T15:13:36.629Z" }, -] - -[[package]] -name = "tabpfn" -version = "9.2.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "einops" }, - { name = "filelock" }, - { name = "huggingface-hub" }, - { name = "joblib" }, - { name = "lightgbm" }, - { name = "matplotlib" }, - { name = "mlx", marker = "platform_machine == 'arm64' and sys_platform == 'darwin'" }, - { name = "numpy" }, - { name = "pandas" }, - { name = "pydantic" }, - { name = "pydantic-settings" }, - { name = "safetensors" }, - { name = "scikit-learn" }, - { name = "scipy" }, - { name = "skrub" }, - { name = "torch", version = "2.14.1", source = { registry = "https://pypi.org/simple" }, marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "torch", version = "2.14.1+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "tqdm" }, - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/b0/03/1b48e093811877c651f731ac64e36a71b3ac44ec5c5c67a395422a6cbfed/tabpfn-9.2.0.tar.gz", hash = "sha256:e7fb2c6293ffef664c69838eb4cb215c79c85a0c0bd449fdf22489045d32b9eb", size = 937644, upload-time = "2026-10-12T15:12:56.56Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/ac/a7/5c403c9fec9f726c91eaeaebaa8905530813ab58cb134a043c4b17afbd6a/tabpfn-9.2.0-py3-none-any.whl", hash = "sha256:bf8edb860cf490ab021d0ba6e91a9798f1b549bdb220066599661c8fc6d6089f", size = 850035, upload-time = "2026-10-12T15:12:54.671Z" }, -] - -[[package]] -name = "threadpoolctl" -version = "3.7.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/00/dc/6c58154c1c65f758ea979e7139cb76993a9cfc662d14e9be3c4a667cfb77/threadpoolctl-3.7.0.tar.gz", hash = "sha256:61348cfb77d53b9242e0017029244b559b810c142ced65b4e21eeca1843959a7", size = 31961, upload-time = "2026-09-15T15:46:20.263Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/43/3f/f88a53f60a472b46f4023f56d204dd7de33d34c5d2acbfa0d70a674e639e/threadpoolctl-3.7.0-py3-none-any.whl", hash = "sha256:cd8b60b5641b45c67bbf73c64c843235fc2d8a480c87389f52f5dbee893b86be", size = 26362, upload-time = "2026-09-15T15:46:19.168Z" }, -] - -[[package]] -name = "torch" -version = "2.14.1" -source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "sys_platform == 'win32'", - "sys_platform == 'emscripten'", - "(platform_machine != 'x86_64' and sys_platform == 'linux') or (sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32')", -] -dependencies = [ - { name = "cuda-bindings", marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, - { name = "filelock", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "fsspec", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "jinja2", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "networkx", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "nvidia-cudnn-cu13", marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, - { name = "setuptools", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "sympy", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, - { name = "triton", marker = "platform_machine != 'x86_64' and sys_platform == 'linux'" }, - { name = "typing-extensions", marker = "platform_machine != 'x86_64' or sys_platform != 'linux'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/43/19/23a1aed488423a5055727256b25406e4b93bd2bcf1352bef582b9951c10c/torch-2.14.1-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:cee091caf2a6229e248daf41d18ceebf590ba02d9179205610062c64ee5fef03", size = 127325330, upload-time = "2026-09-30T17:52:54.752Z" }, - { url = "https://files.pythonhosted.org/packages/93/f4/94219ada13edd62fda1f976163292b7b8595f1fd2f4ef74af24b3baad375/torch-2.14.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:f68f5476e2bc0e8f60b74f7ca21acda885c4477af6392f8977e4f0d1ea1aa162", size = 454016297, upload-time = "2026-09-30T17:54:03.078Z" }, - { url = "https://files.pythonhosted.org/packages/fd/df/23c69e9b9fd19fe6563422f1bc59f89601bda47234a89a1cf432cfde3aee/torch-2.14.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:1d4df07be7338bbcc4d54085adee21363c91314702b9bd4d78ef72ffac9465ba", size = 554621744, upload-time = "2026-09-30T17:54:22.951Z" }, - { url = "https://files.pythonhosted.org/packages/d9/dc/a36a4431ab5e3ad168a75f341e24677a4bbbf0c5dda59c97a7079262bb94/torch-2.14.1-cp314-cp314-win_amd64.whl", hash = "sha256:d02a4c48a2ca5fb7654e36e71f710f74494d83f1f10aaff8e059dd554adca956", size = 124110781, upload-time = "2026-09-30T17:53:31.141Z" }, - { url = "https://files.pythonhosted.org/packages/2d/bc/1afbd1a22f6023eafcdb9295ac1d97137a7816f54091d693701072da2218/torch-2.14.1-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:06c3ec25f3b497f9a73dc0c36f293f295cfe8d5584fd446238d9a501b30a66d5", size = 127657364, upload-time = "2026-09-30T17:53:53.837Z" }, - { url = "https://files.pythonhosted.org/packages/e4/94/b97e863c9ceef2bc6e082e967cbfb9065b5e232f2a60c5af0f5115a1bc31/torch-2.14.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:9113f94f70429f9f302bf55b090b5411269083e8a72a5d015ffcc5a7f83f2c69", size = 453998592, upload-time = "2026-09-30T17:54:38.268Z" }, - { url = "https://files.pythonhosted.org/packages/03/d8/8272157c438cc26a199a8fa0caf87985477b8a1ed8a09084eb2ae40609d8/torch-2.14.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:e65d5fe136e533b23c2d134377f7f126721c7c78dd75b3dbb735af082a8aeb85", size = 554578730, upload-time = "2026-09-30T17:54:54.838Z" }, - { url = "https://files.pythonhosted.org/packages/45/05/451a69a4287033d8f106f5c65f92c2c0c37229d81ea101d4b047caf758cc/torch-2.14.1-cp314-cp314t-win_amd64.whl", hash = "sha256:e07306caa1de2a4ac1467e11ecfc92fc44f523dd6a521145039aef46d913963c", size = 124400129, upload-time = "2026-09-30T17:54:12.306Z" }, -] - -[[package]] -name = "torch" -version = "2.14.1+cu130" -source = { registry = "https://download.pytorch.org/whl/cu130" } -resolution-markers = [ - "platform_machine == 'x86_64' and sys_platform == 'linux'", -] -dependencies = [ - { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "filelock", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "fsspec", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "jinja2", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "networkx", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cudnn-cu13", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "setuptools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "sympy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, -] -wheels = [ - { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.14.1%2Bcu130-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:5206fa83163aa71dc0186637b2aa2ec3fb7114892039fdde3c91bf570bd23ba8", upload-time = "2026-09-30T20:29:22Z" }, - { url = "https://download-r2.pytorch.org/whl/cu130/torch-2.14.1%2Bcu130-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:f143ff078be3bfd41579a4557a1ebcf01506f8223fad3ac68ce7bb6e88f96a12", upload-time = "2026-09-30T20:32:06Z" }, -] - -[[package]] -name = "tqdm" -version = "4.70.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/21/3b/6c24bec5be5e743ffd99576daa5cc077722fc7d5bbc00bd133fa0c698dc6/tqdm-4.70.0.tar.gz", hash = "sha256:55b0b0dbd97462d06ebee91e4dac24ed4d4702be82b24f07e6c1d27e08cea220", size = 795438, upload-time = "2026-07-27T11:33:15.271Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl", hash = "sha256:7f585706bfddbdebf89daac705b2dfcc16890130727d3197ca62c732b4310953", size = 80184, upload-time = "2026-07-27T11:33:13.167Z" }, -] - -[[package]] -name = "triton" -version = "3.8.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b8/59/bf0e9493118bb353ab59a5d6a65db3618d9b314417cc1459f0121e0ec5c9/triton-3.8.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1f6b48d0591929a3867973acac3dccd4e058585f91bfb41022de496c9ffab304", size = 226488654, upload-time = "2026-08-28T16:08:47.141Z" }, - { url = "https://files.pythonhosted.org/packages/93/d9/08c75f3459f19ad00425b564058e40efa4bcd79b816064cf27499303ea42/triton-3.8.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:387dae4cb0089a7b6ba1a428ae0782b65c4c58f57d94617cb22ca8593d8ccbca", size = 247972313, upload-time = "2026-08-28T15:56:14.007Z" }, - { url = "https://files.pythonhosted.org/packages/7c/34/429c5592181cfb7361a0a8e0bff218e7224b726709d75da2472b3e819f70/triton-3.8.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b84e7d512490ba529111260fa6f7cad8b254a6bb5fbdf41d5ef9a5e57f52d0a", size = 226591133, upload-time = "2026-08-28T16:09:02.271Z" }, - { url = "https://files.pythonhosted.org/packages/fe/d1/aa8a3e935c37efee7945984fdb64d7e0851bf6d920afd97b2d21f9d23360/triton-3.8.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74217bb56ed8692759227758e4c4b3bd2d608a209c1a7a081bf361fb4c2c1bf9", size = 248077577, upload-time = "2026-08-28T15:56:24.94Z" }, -] - -[[package]] -name = "typing-extensions" -version = "4.16.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, -] - -[[package]] -name = "typing-inspection" -version = "0.4.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a3/26/b09b8010994eccc3c09092e6b34058f36a460eea2d4c3e8b910c695975a0/typing_inspection-0.4.4.tar.gz", hash = "sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47", size = 76928, upload-time = "2026-08-12T12:37:25.997Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/67/81/4add07e5172b7ac40d8ed5ff580409a7801a4fe26d529bdd915401dabfbe/typing_inspection-0.4.4-py3-none-any.whl", hash = "sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147", size = 14750, upload-time = "2026-08-12T12:37:24.648Z" }, -] - -[[package]] -name = "tzdata" -version = "2026.5" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/d9/68/f1b440335057bfce71b6e50a9d09445aa2ecbd08359a337976627b8409e7/tzdata-2026.5.tar.gz", hash = "sha256:8cc73c0a0bfca7dbfa59235d60b2eff82231dee33f53d206db1acd9173cfc0a7", size = 200404, upload-time = "2026-10-03T09:23:14.143Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/94/21/1e5995a1c920cce14e4bffae20c665ec10e7ed03ab25e006cd741092b718/tzdata-2026.5-py2.py3-none-any.whl", hash = "sha256:b683bd1b6659ddcd810ff02ad09ba821d4bf1065072805063eb35c49617905ac", size = 347996, upload-time = "2026-10-03T09:23:12.535Z" }, -] - -[[package]] -name = "uproot" -version = "5.7.5" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "awkward" }, - { name = "cramjam" }, - { name = "fsspec" }, - { name = "numpy" }, - { name = "packaging" }, - { name = "xxhash" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/d6/e6/9739727cc9c9ed95c40c23993eeecd9ea8367d5c36cbea20d97605a49de5/uproot-5.7.5.tar.gz", hash = "sha256:c9a30c8b39ffd30d9f87d7a8ec49a82c3d71fe163161b28e193c83d771119cd6", size = 1025717, upload-time = "2026-07-04T10:43:21.14Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/2c/2fb9b2d992bd9a9c4aed623c0ace167988027ccc7421b178a26a9a40ef50/uproot-5.7.5-py3-none-any.whl", hash = "sha256:25aee18f2d11dda6b26c6d76d84a9bc0f0aa0adeeba60548e47391964edae760", size = 401195, upload-time = "2026-07-04T10:43:19.434Z" }, -] - -[[package]] -name = "urllib3" -version = "2.7.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, -] - -[[package]] -name = "xxhash" -version = "4.0.1" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/f6/a5/1386f35da1475fcaeef42581deae73417c6d2a6a0b2d2e8914de18844dcd/xxhash-4.0.1.tar.gz", hash = "sha256:d55bf4ef10eb09b8b6866790e083d26d087d84caa3cc0946ba87c3ca7ecaf7b7", size = 101513, upload-time = "2026-08-17T08:24:08.557Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/81/0e/ea406a02b561d3275232ccfdb3e29df80f7a65414940e3a15721c7bea40f/xxhash-4.0.1-cp314-cp314-android_24_arm64_v8a.whl", hash = "sha256:af05a3f650220a6c59fa0ad2410249f2d2470a05225807c378fb67458693f8df", size = 43747, upload-time = "2026-08-17T08:22:31.37Z" }, - { url = "https://files.pythonhosted.org/packages/f9/f0/b0c94d61ccf6b5d1f8847b58ef8f923125ac4919ed5bd0eb082750ca7cbd/xxhash-4.0.1-cp314-cp314-android_24_x86_64.whl", hash = "sha256:a6e3653df1a70b8ac4191216324242e4be2bca18c9a7c10934e1bd56dc7ca15e", size = 40749, upload-time = "2026-08-17T08:22:29.431Z" }, - { url = "https://files.pythonhosted.org/packages/2f/c5/8085881a538983be0fd1c865d5df236242fea496044e2c8ca32b9f2ba39c/xxhash-4.0.1-cp314-cp314-ios_13_0_arm64_iphoneos.whl", hash = "sha256:4528cf80ebbbf57d40edfb31521ae265daa6dd636d615b1cf0ac86209579e59d", size = 34734, upload-time = "2026-08-17T08:35:33.68Z" }, - { url = "https://files.pythonhosted.org/packages/d3/94/8803d13c968fc75ca434eea991d29ac5fd8a36b4afc9a6a9803c53933db4/xxhash-4.0.1-cp314-cp314-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:90cb2a1c9cc503a054a19612b48ff6e8e47805f618bdb3224a07568aad03a37e", size = 35671, upload-time = "2026-08-17T08:21:48.322Z" }, - { url = "https://files.pythonhosted.org/packages/85/d5/ad91d7f0fd294190d37c08236fe661f5c4e3f83dcd1a121877a2e64681ce/xxhash-4.0.1-cp314-cp314-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:a949b072ea59c6eca0811ccd9e95133cc50d2afda8d464b5b077c78f78efa269", size = 38094, upload-time = "2026-08-17T08:22:39.763Z" }, - { url = "https://files.pythonhosted.org/packages/89/f4/2b7ebdc1869caca5f02c4cba8379b631050d3c3d4adb9187e4dc1a6b8d3c/xxhash-4.0.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:79a3203aadf39637869dfea1185227d8452844d78b837e54fb1117b4d34ba5c3", size = 38244, upload-time = "2026-08-17T08:35:38.081Z" }, - { url = "https://files.pythonhosted.org/packages/90/9d/f66cf6935f528e575f1ae4d6560d376e7587569747186f4fae8777cadc1b/xxhash-4.0.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d9f3848ffaf010bdbabdbf4c25641fa258b6227ff27bc74a4d06edef521a4873", size = 35904, upload-time = "2026-08-17T08:21:37.358Z" }, - { url = "https://files.pythonhosted.org/packages/07/29/34569d7b482f0dc060074faafd163c588f915cbc3e3e218f1ffd8a3ad340/xxhash-4.0.1-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:9283d9dd6b44acad35118e2976fc763a065509e4118debdb61916ec322ed17b9", size = 259595, upload-time = "2026-08-17T08:22:38.153Z" }, - { url = "https://files.pythonhosted.org/packages/ce/d2/a2370acfcd48732cf5c2b87f06cfbf7fa51c0ce0dd736bde42939eb9ebf7/xxhash-4.0.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1c7c642a0f79c3e3cf2965475507574d3d1a50ec71060039d60cb87358667cb2", size = 284279, upload-time = "2026-08-17T08:22:36.396Z" }, - { url = "https://files.pythonhosted.org/packages/08/15/17d33c24e6c4a1c0b9ddc5584f0c25d51d48b34bacde1416a2235a19db4b/xxhash-4.0.1-cp314-cp314-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:96dedccfb09a73a25751053a183159b88f4ee75f388df8166040c152ac0531c6", size = 303973, upload-time = "2026-08-17T08:35:39.22Z" }, - { url = "https://files.pythonhosted.org/packages/ec/e0/4ec0d69ad5738729098a61e631b7ed2df22a922b0e03014b597c72bd863d/xxhash-4.0.1-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:81664268dba92e037b740ecf37fa02f1cab4a391f93f28e35792b3341c60648f", size = 287535, upload-time = "2026-08-17T08:21:52.158Z" }, - { url = "https://files.pythonhosted.org/packages/0f/8b/4f9b17e7a9eb71c65548ecddd9c18b84e3c18ca41c4d436ad2a3000d3f7b/xxhash-4.0.1-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:839f58c5bd9989875be0fd28446dbf32cace2c2cd8bf2f6762acdc38a95cd1aa", size = 519257, upload-time = "2026-08-17T08:22:43.272Z" }, - { url = "https://files.pythonhosted.org/packages/68/35/3276b3e743b8ddbed9c3f71c76d9dd6a75d72aa4e678b1447b635cfd92e0/xxhash-4.0.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ffa44b4c7c5d0ffa31356b4428659516c0e47647825c74079a296b3857b6d99d", size = 268190, upload-time = "2026-08-17T08:35:44.985Z" }, - { url = "https://files.pythonhosted.org/packages/08/d4/f1555de3c96721320930dbb7988c8482d82b85970076aba1a8d40e83ad43/xxhash-4.0.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e681a6fc7e4f715252b9b5acfb30536ec7dd1f75033a32dc617e6fa95af1a3fd", size = 345553, upload-time = "2026-08-17T08:21:41.025Z" }, - { url = "https://files.pythonhosted.org/packages/ac/98/c28908f27007087b61139d290f908dd827ffd40b88af0c43f9e1a1a7ffd5/xxhash-4.0.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:c6301d92545c591ad31c3e050aa40a5f8a4c16413f1f9e6f9322c6f0f9d2b736", size = 280499, upload-time = "2026-08-17T08:22:52.236Z" }, - { url = "https://files.pythonhosted.org/packages/a9/76/3ef57622c65816348f8196273485baab4752aae064959901e85cd867e067/xxhash-4.0.1-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:6efb8f21cc136c79b3e5bb747c8682d37916fb202cdbbc32182de5c4e47f821f", size = 307211, upload-time = "2026-08-17T08:22:40.815Z" }, - { url = "https://files.pythonhosted.org/packages/8a/4c/5804504bbc808968e57d6a50286dd8f8cc06e0ddd6e4ab4b1dc89ae42f35/xxhash-4.0.1-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:760de77279e9cf9c81d012ce0705cba13afccee9b09c480f17d778c8c5cefae8", size = 265865, upload-time = "2026-08-17T08:35:42.727Z" }, - { url = "https://files.pythonhosted.org/packages/aa/ee/8572fdfd70e7aaaf150af899871c2cc0bb88c3295ca82172a31e04ca5168/xxhash-4.0.1-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:a16a3fa6936e36bb1414d16a6bd012c9033e5161b68b426805b61d895392437d", size = 284545, upload-time = "2026-08-17T08:21:56.965Z" }, - { url = "https://files.pythonhosted.org/packages/c1/f8/6eadcca0904660c848b466524e82a233d16c9d2d5258433aaf3546142d86/xxhash-4.0.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:9c3c4b9aa9a27196b921197f7daf9e6c1412739df06a99cfa6e923879362eff6", size = 336022, upload-time = "2026-08-17T08:22:46.346Z" }, - { url = "https://files.pythonhosted.org/packages/27/df/4aa107b81602d6d6d09ab5a607c530d2d3a6b28e2e9a59b01875bd877c54/xxhash-4.0.1-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:863f3d3b44110f7243e86cf994aa5c5d88f2348b6e84ab4402fadadfbf9f7da7", size = 486671, upload-time = "2026-08-17T08:35:49.016Z" }, - { url = "https://files.pythonhosted.org/packages/45/b7/b2bf9b5301e9cd5f2e335fea8da0f5cf209a6594cb1fe77754774ad4a6fd/xxhash-4.0.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:63aa52659bc32bb9bd7cb5caf523b4d14429a477762cfac886132d687c1f80fc", size = 263744, upload-time = "2026-08-17T08:21:56.165Z" }, - { url = "https://files.pythonhosted.org/packages/0b/96/35b1c02177ae26234892c2310fb4822ba62411acccbf425ab8f9fd99354a/xxhash-4.0.1-cp314-cp314-pyemscripten_2026_0_wasm32.whl", hash = "sha256:67e57b834e07ed973cee7b6da1548ff28a56458d77696fd2a5f397f340694848", size = 20563, upload-time = "2026-08-17T08:35:11.924Z" }, - { url = "https://files.pythonhosted.org/packages/51/c2/a06300b165fbd6b0cb4a9742987f2e997a9f447ce3bf7c6ac97b862ce62a/xxhash-4.0.1-cp314-cp314-win32.whl", hash = "sha256:b6c1f9c59bbe593f88a0aad30be4150f15bd57bd64efb95feeabcb8e563f1ecd", size = 35151, upload-time = "2026-08-17T08:22:44.283Z" }, - { url = "https://files.pythonhosted.org/packages/06/96/c5b37296b78f80fc97124c0fee0c7bbd1bdb6f3b18bcd8748bb113b2d8fc/xxhash-4.0.1-cp314-cp314-win_amd64.whl", hash = "sha256:da544672efd9ad76077928a3e6c5d894e52ce82d3bf14002db4a1bf17d1a36a2", size = 37156, upload-time = "2026-08-17T08:35:46.551Z" }, - { url = "https://files.pythonhosted.org/packages/ce/5e/248f9cd169c2fb62236bedfba246d213bce728f74901e99047e3f3c55875/xxhash-4.0.1-cp314-cp314-win_arm64.whl", hash = "sha256:d0d24a4f3fb63852cd09af46ae4b7a4d00cc8b8615a046dca543786e728d1056", size = 34379, upload-time = "2026-08-17T08:21:59.446Z" }, - { url = "https://files.pythonhosted.org/packages/58/c8/db1d37c0da0324d0298f6abd931ca1d4736e049d9f2081230a8421da74d2/xxhash-4.0.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:349775ac30372b344d2338b2a168c0a1312a644194da25b8bec476d55761a128", size = 38656, upload-time = "2026-08-17T08:22:49.119Z" }, - { url = "https://files.pythonhosted.org/packages/c5/8e/e18998ec465fb977bc74272e5bf3c2e886c13b014cbef916cd607802c709/xxhash-4.0.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:43e5f9169e73d0f0db33b5f6b8554bcce69ac278c966daf83d5eb4eb2f13829f", size = 36306, upload-time = "2026-08-17T08:35:52.853Z" }, - { url = "https://files.pythonhosted.org/packages/ef/1a/b83f86f8a987a3cbcb7e005a6824ff64aecae35abc1395a0d44ee16c3319/xxhash-4.0.1-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4a252fb862b0ae2590587e625f47a0e03da05cf0205e8830b67b6596c06038b1", size = 273729, upload-time = "2026-08-17T08:21:58.833Z" }, - { url = "https://files.pythonhosted.org/packages/02/4e/2db15aa8508e0cd5b632927a53b98234f24039ea65377e6cf996c06d2d4f/xxhash-4.0.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:2df3ca8757dc381e75e90a4d7995a6324f58a923c7145220a7b2c0231f66fddc", size = 301083, upload-time = "2026-08-17T08:35:14.113Z" }, - { url = "https://files.pythonhosted.org/packages/26/94/ed759787ffe802bd8e31cfcdad3755cbeca2dcdafd2f790cd6f25d195199/xxhash-4.0.1-cp314-cp314t-manylinux2014_armv7l.manylinux_2_17_armv7l.manylinux_2_31_armv7l.whl", hash = "sha256:bfed61996d618eb90d6eaae0178002e3466a28b06bfc557a7a3a7266378d8c5a", size = 312745, upload-time = "2026-08-17T08:22:52.232Z" }, - { url = "https://files.pythonhosted.org/packages/45/7a/f64b4a4cc8b51e950709207f55f7f56ae9c5af6631dd31d7fb443312418c/xxhash-4.0.1-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9761ff4a0ffa583fe850731ad24fe82c88cccb7a2294727db0955f3279a4cb3f", size = 301419, upload-time = "2026-08-17T08:35:50.143Z" }, - { url = "https://files.pythonhosted.org/packages/a0/71/bac313b8de073569b8db3152044a7cfcce87a3fa9698c18fe9f914dee6b1/xxhash-4.0.1-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:edccc2ec58435a580f96a48a3ccae8cd0a480824119165dd90108718ad81ae6e", size = 534485, upload-time = "2026-08-17T08:22:11.515Z" }, - { url = "https://files.pythonhosted.org/packages/b9/0c/16b5e419f24e59507ee05626d2bb0deafdb03f9f27783bc0785a9849602e/xxhash-4.0.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4741d42d59e4e5fa1a86c17ab9c27dc8ea459c700d91b6742fdb9138d9a516cb", size = 279605, upload-time = "2026-08-17T08:22:52.934Z" }, - { url = "https://files.pythonhosted.org/packages/5f/55/5787dd6e2d8d5b61256a5039f6b18c2193c7c1de4a2fd2413288d0d9c604/xxhash-4.0.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:440c401e146ce64bdb3beb8ff0c84677b6f21307c28a34779071cecee5d4d70c", size = 358924, upload-time = "2026-08-17T08:35:58.164Z" }, - { url = "https://files.pythonhosted.org/packages/f3/68/89be41991f3b0a2e91f940bdf3128852c3ed571cf560d98ad0f67024afe4/xxhash-4.0.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:5b7979f71d06ae45a769de0699900a246d8cb632db1e8bfdc79ec019063a503c", size = 295305, upload-time = "2026-08-17T08:22:13.683Z" }, - { url = "https://files.pythonhosted.org/packages/e6/5a/52ff0a0cc361aad393ff9a46ffe3aabbcf9c03d6c8f2612da7d553048276/xxhash-4.0.1-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:62198213fc3e0c56e567894b318ba45834e007d065f84ba6dc9165d21546fc56", size = 320228, upload-time = "2026-08-17T08:35:18.946Z" }, - { url = "https://files.pythonhosted.org/packages/0f/b5/91c60ff22c7f6cd5f6d7a5bad5a2cdcb4c33987dfa50bf13f0d856279b2e/xxhash-4.0.1-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:b3bece52127ac20044311ee73567f9f0893b5de64f9028aecc90cc740cfd525a", size = 279414, upload-time = "2026-08-17T08:23:03.212Z" }, - { url = "https://files.pythonhosted.org/packages/b9/94/9685954804d47d0390871a64bec606a0d536406382d71a784df3a5883fb4/xxhash-4.0.1-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:a865d2d470220e659220fdb59d5b6c4422802d8d6098e1324bc4d12444798914", size = 297594, upload-time = "2026-08-17T08:35:57.881Z" }, - { url = "https://files.pythonhosted.org/packages/89/62/b67ac9412907b7a07a2a0c08c3440b9e4480231a7b3de0767e87011e4564/xxhash-4.0.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:8580aab306888224074c7edeec734de0c3c5ccde65b2da4e6c9a5e28f7c0a1bd", size = 348526, upload-time = "2026-08-17T08:22:18.571Z" }, - { url = "https://files.pythonhosted.org/packages/37/ed/6723cc49a9f567d52d01fd7c1741b0f2e3a13e71d15f7ac49d753a20c115/xxhash-4.0.1-cp314-cp314t-musllinux_1_2_s390x.whl", hash = "sha256:2d52dc7c33c1b83082b707f6b7814dc76d2faaa2ea62bd9c5fab4b36f83c087f", size = 499307, upload-time = "2026-08-17T08:22:56.52Z" }, - { url = "https://files.pythonhosted.org/packages/fd/2e/7b10e101ab988d93b791023be7191d7661271d6ab31ac082276b9091042a/xxhash-4.0.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:6a9f98af872355e0c02439e48583958eee00e60b928bb20476460d9d40cb7b4e", size = 274989, upload-time = "2026-08-17T08:36:01.834Z" }, - { url = "https://files.pythonhosted.org/packages/9b/8d/7eabcc8d29cce40621443cff24c07d7306ef574b8956c47ac59f21098005/xxhash-4.0.1-cp314-cp314t-win32.whl", hash = "sha256:a14578102a6081465aec9cf73c76c3cd3f79f0709bdb3b8ae7ab0b54c9d8b089", size = 35482, upload-time = "2026-08-17T08:22:32.336Z" }, - { url = "https://files.pythonhosted.org/packages/ca/89/2a4268e1971f63038b79fb75e3b9c8de942cd77acabbb0c5625352a31940/xxhash-4.0.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c57963970d359a72262f7fe6be88f945e2334d4bc41462b7f08c37b0abf35ca6", size = 37490, upload-time = "2026-08-17T08:35:22.475Z" }, - { url = "https://files.pythonhosted.org/packages/90/7b/950ecab1fe4cf421d0a6211ddd9a0ac82e39e55c45a111ceb90953dc6c9a/xxhash-4.0.1-cp314-cp314t-win_arm64.whl", hash = "sha256:b659fad79c99b0238c7ad7e9d7dbf4eebfea9097c2dba65fa0a4d18a25b29a2f", size = 34596, upload-time = "2026-08-17T08:23:10.001Z" }, -]