heptokens 0.1.0__tar.gz

heptokens-0.1.0/.gitattributes

@@ -0,0 +1,2 @@
+# SCM syntax highlighting & preventing 3-way merges
+pixi.lock merge=binary linguist-language=YAML linguist-generated=true

heptokens-0.1.0/.github/workflows/cd.yml

@@ -0,0 +1,36 @@
+---
+name: CD
+on:
+  workflow_dispatch:
+  pull_request:
+  push:
+    branches: [main]
+  release:
+    types: [published]
+jobs:
+  dist:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v7
+        with:
+          fetch-depth: 0
+      - uses: hynek/build-and-inspect-python-package@v3.0.1
+  publish:
+    needs: [dist]
+    environment: pypi
+    permissions:
+      id-token: write
+      attestations: write
+      contents: read
+    runs-on: ubuntu-latest
+    if: github.event_name == 'release' && github.event.action == 'published'
+    steps:
+      - uses: actions/download-artifact@v8
+        with:
+          name: Packages
+          path: dist
+      - name: Generate artifact attestation for sdist and wheel
+        uses: actions/attest@v4
+        with:
+          subject-path: dist/*
+      - uses: pypa/gh-action-pypi-publish@release/v1

heptokens-0.1.0/.gitignore

@@ -0,0 +1,35 @@
+# pixi environments
+.pixi/*
+!.pixi/config.toml
+
+# Snakemake files
+.snakemake/
+
+# Python cache files
+__pycache__/
+*.py[cod]
+*$py.class
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+
+# VsCode settings
+.vscode/
+
+# Plots and other results
+plot/
+results/

heptokens-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,17 @@
+---
+repos:
+  - repo: https://github.com/charliermarsh/ruff-pre-commit
+    rev: v0.11.13
+    hooks:
+      - id: ruff
+        args: [--fix, --show-fixes, --exit-non-zero-on-fix]
+      - id: ruff-format
+  - repo: https://github.com/lyz-code/yamlfix
+    rev: 1.16.1
+    hooks:
+      - id: yamlfix
+  - repo: https://github.com/adrienverge/yamllint.git
+    rev: v1.37.1
+    hooks:
+      - id: yamllint
+        args: [-d, '{extends: relaxed, rules: {line-length: disable}}']

heptokens-0.1.0/PKG-INFO

@@ -0,0 +1,446 @@
+Metadata-Version: 2.5
+Name: heptokens
+Version: 0.1.0
+Author: Samuel Klein, Jeffrey Krupa, Michael Kagan
+Requires-Python: >=3.11
+Requires-Dist: awkward>=2.8
+Requires-Dist: h5py>=3.13
+Requires-Dist: hydra-core>=1.3
+Requires-Dist: lightning>=2.5
+Requires-Dist: omegaconf>=2.3
+Requires-Dist: rich>=14
+Requires-Dist: scikit-learn>=1.3
+Requires-Dist: torch>=2.4
+Requires-Dist: uproot>=5.7
+Requires-Dist: vector-quantize-pytorch
+Description-Content-Type: text/markdown
+
+# heptokens
+
+[![pytorch](https://img.shields.io/badge/-PyTorch-EE4C2C?logo=pytorch&logoColor=white)](https://pytorch.org/)
+[![lightning](https://img.shields.io/badge/-Lightning-792EE5?logo=lightning&logoColor=white)](https://lightning.ai/)
+[![hydra](https://img.shields.io/badge/-Hydra-89b8cd?logo=hydra&logoColor=white)](https://hydra.cc/)
+[![wandb](https://img.shields.io/badge/-WandB-orange?logo=weightsandbiases&logoColor=white)](https://wandb.ai)
+[![pixi](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/prefix-dev/pixi/main/assets/badge/v0.json)](https://pixi.sh)
+
+Tokenize HEP physics objects (jets, tracks, calorimeter clusters, …) with residual VQ-VAEs.
+
+`heptokens` is an installable Python library. It ships:
+- **Data classes** — generic HDF5 datasets and Lightning DataModules for structured (object + constituent) and flat (dense, maskless) data.
+- **VQ-VAE model** — `LitVqVae` wrapping `ResidualVQ` with pluggable encoder/decoder.
+- **CLI entry points** — `heptokens-fit-preprocessors`, `heptokens-train`, and `heptokens-export` for Hydra-configured preprocessing, training, and tokenization runs.
+
+## Quick Start
+
+### 1. Install
+
+```bash
+pip install heptokens
+```
+
+Or, for development with a [Pixi](https://pixi.sh/latest/installation/)-managed environment:
+
+```bash
+git clone <repo>
+cd heptokens
+pixi shell          # activates the environment
+```
+
+### 2. Write your datamodule config
+
+Create a YAML file describing your data. This is the only config you must supply — the library bundles everything else (model, callbacks, trainer).
+
+```yaml
+# my_configs/datamodule/my_jets.yaml
+_target_: heptokens.data.structured_array.StructuredArrayModule
+data_path: /path/to/data.h5
+obj_group: jets
+set_group: tracks
+obj_features: [pt, eta, phi]
+set_features: [pt, deta, dphi, d0, z0]
+label_key: HadronConeExclTruthLabelID
+mask_key: valid
+train_frac: 0.7
+val_frac: 0.15
+test_frac: 0.15
+batch_size: 1024
+num_workers: 4
+```
+
+Your HDF5 must have this layout (see `StructuredArrayHDF5Dataset` for full docs):
+
+```
+<obj_group>/
+    <feature>: float [N]
+    <label_key>: int [N]
+<set_group>/
+    <feature>: float [N, max_elements]
+    <mask_key>: bool [N, max_elements]
+```
+
+### 3. Fit preprocessors
+
+The VQ-VAE expects normalised inputs. Use the `heptokens-fit-preprocessors` CLI to fit sklearn scalers on your training data and save them as joblib files:
+
+```bash
+heptokens-fit-preprocessors -cd /path/to/my_configs \
+  datamodule=my_jets \
+  output_dir=resources
+```
+
+Dataset-specific options (which modes to fit, which features to log-scale, etc.) live in a `preprocessing:` block inside your datamodule config so they don't need to be repeated on every invocation:
+
+```yaml
+# my_configs/datamodule/my_jets.yaml
+_target_: heptokens.data.structured_array.StructuredArrayModule
+set_features: [pt, deta, dphi, d0, z0]
+...
+preprocessing:
+  cst_modes: [quantile, log_quantile]
+  cst_log_feature_indices: [pt]   # resolved via set_features — no separate feature list needed
+  fit_jets: true
+```
+
+Then just run:
+
+```bash
+heptokens-fit-preprocessors -cd /path/to/my_configs \
+  datamodule=my_jets \
+  output_dir=resources
+```
+
+Individual keys can still be overridden on the CLI with `datamodule.preprocessing.<key>=...`:
+
+```bash
+heptokens-fit-preprocessors -cd /path/to/my_configs \
+  datamodule=my_jets \
+  datamodule.preprocessing.cst_modes=[quantile,log_quantile]
+```
+
+This saves `resources/cst_<mode>.joblib` (and `resources/jet_<mode>.joblib` when `fit_jets=true`) and writes `resources/preprocessor_config.yaml` with the file paths. Then add them to your datamodule config:
+
+```yaml
+# append to my_configs/datamodule/my_jets.yaml
+transforms:
+  preprocess:
+    _target_: heptokens.data.collation.preprocess_batch
+    _partial_: true
+    cst_fn:
+      _target_: joblib.load
+      filename: resources/cst_quantiles.joblib
+    jet_fn:
+      _target_: joblib.load
+      filename: resources/jet_quantiles.joblib
+```
+
+See [Preprocessing](#preprocessing) for available modes and details.
+
+### 4. Train
+
+```bash
+heptokens-train -cd /path/to/my_configs \
+  datamodule=my_jets \
+  project_name=my_project \
+  network_name=run_01 \
+  output_dir=results
+```
+
+### 5. Export tokens
+
+```bash
+heptokens-export -cd /path/to/my_configs \
+  datamodule=my_jets \
+  ckpt_path=results/my_project/run_01/checkpoints/last.ckpt \
+  output_dir=results/tokens
+```
+
+Output is an `.npz` with `indices [N, n_csts, num_quantizers]`, `labels [N]`, and `codebooks [Q, K, D]`.
+
+### Python API
+
+For scripting or notebooks without Hydra:
+
+```python
+from heptokens.data.structured_array import StructuredArrayHDF5Dataset
+from heptokens.models.vq_vae import LitVqVae
+from heptokens.models.coders import Encoder, Decoder
+import lightning as L
+from torch.utils.data import DataLoader
+
+ds = StructuredArrayHDF5Dataset(
+    "data.h5",
+    obj_group="jets",      set_group="tracks",
+    obj_features=["pt", "eta", "phi"],
+    set_features=["pt", "eta", "phi"],
+    mask_key="valid",      label_key=None,
+)
+loader = DataLoader(ds, batch_size=256, shuffle=True)
+
+sample = next(iter(loader))["csts"]          # [B, N, 3]
+model = LitVqVae(
+    encoder=Encoder, decoder=Decoder,
+    codebook_size=512, codebook_dim=64, num_quantizers=4,
+    data_sample=sample,
+)
+
+trainer = L.Trainer(max_epochs=10, logger=False)
+trainer.fit(model, loader)
+
+batch = next(iter(loader))
+z_q, indices, _ = model.encode(batch)   # indices: [B, N, num_quantizers]
+```
+
+## CLI Reference
+
+All three commands are Hydra-based. Use `-cd` to point at your experiment's config directory containing the `datamodule/` folder. All other configs (model, callbacks, trainer) are bundled.
+
+### Full workflow
+
+```bash
+# 1. Fit scalers on training data
+heptokens-fit-preprocessors -cd /path/to/my_configs \
+  datamodule=my_jets \
+  output_dir=resources
+
+# 2. Wire the saved joblib paths into your datamodule config
+#    (see resources/preprocessor_config.yaml for the generated snippet)
+
+# 3. Train the VQ-VAE
+heptokens-train -cd /path/to/my_configs \
+  datamodule=my_jets \
+  project_name=my_project \
+  network_name=run_01 \
+  output_dir=results
+
+# 4. Export tokens from a trained checkpoint
+heptokens-export -cd /path/to/my_configs \
+  datamodule=my_jets \
+  ckpt_path=results/my_project/run_01/checkpoints/last.ckpt \
+  output_dir=results/tokens
+```
+
+```bash
+# Without installing (pixi environment):
+pixi run python -m heptokens.fit_preprocessors -cd /path/to/my/configs datamodule=my_jets ...
+pixi run python -m heptokens.train -cd /path/to/my/configs datamodule=my_jets ...
+pixi run python -m heptokens.export_tokens -cd /path/to/my/configs datamodule=my_jets ...
+```
+
+### Bundled configs — `heptokens-train` / `heptokens-export`
+
+| Group | Location | Description |
+|---|---|---|
+| `model` | `src/heptokens/conf/model/` | `vqvae`, `transformer_vqvae`, `classifier`, … |
+| `callbacks` | `src/heptokens/conf/callbacks/` | `pretrain`, `encode`, `classify`, … |
+
+`datamodule` is not bundled — it is experiment-specific. See `StructuredArrayModule` in `src/heptokens/data/structured_array.py` for the constructor args.
+
+### Bundled configs — `heptokens-fit-preprocessors`
+
+The command ships `src/heptokens/conf/fit_preprocessors.yaml` with these key options:
+
+Top-level keys (set on the CLI or in your own `fit_preprocessors.yaml`):
+
+| Key | Default | Description |
+|---|---|---|
+| `output_dir` | `resources` | Directory to write joblib files and the config snippet |
+| `max_batches` | — | Limit training batches used for fitting (null = full dataset) |
+
+`preprocessing:` block keys (set inside your datamodule config, or via `datamodule.preprocessing.<key>=...` on the CLI):
+
+| Key | Default | Description |
+|---|---|---|
+| `cst_modes` | `[quantile, standard]` | Scaler modes to fit for constituents |
+| `jet_modes` | `[quantile, standard]` | Scaler modes to fit for jet-level features |
+| `fit_jets` | `false` | Set `true` to also fit jet-level scalers |
+| `cst_log_feature_indices` | — | Features to log-scale before the main transform (names resolved via `set_features`, or indices) |
+| `jet_log_feature_indices` | — | Same for jet features (names resolved via `obj_features`) |
+| `n_quantiles` | `500` | Number of quantiles for `QuantileTransformer` |
+| `log_offset` | `1.0` | Offset in `log(x + offset)` |
+
+### Output layout
+
+```
+<output_dir>/<project_name>/<network_name>/
+    checkpoints/
+    full_config.yaml      # saved for reproducibility / resume
+```
+
+### Resuming a run
+
+```bash
+heptokens-train -cd /path/to/my/configs \
+  full_resume=true ckpt_flag=last.ckpt \
+  output_dir=... project_name=... network_name=...
+```
+
+Or from an explicit checkpoint path:
+
+```bash
+heptokens-train -cd /path/to/my/configs \
+  full_resume=false ckpt_path=/abs/path/to/last.ckpt \
+  datamodule=my_atlas output_dir=... project_name=... network_name=...
+```
+
+## Repository Layout
+
+```
+src/heptokens/
+    data/
+        structured_array.py     # StructuredArrayHDF5Dataset / StructuredArrayModule
+        flat_array.py           # FlatArrayHDF5Dataset / FlatArrayModule (dense, no mask)
+        base.py                 # BaseMapModule (shared DataLoader logic)
+        collation.py, transforms.py
+    models/
+        vq_vae.py               # LitVqVae
+        coders.py               # Encoder, Decoder (base classes + default MLP impl)
+        new_modality_coders.py  # template for custom encoders/decoders
+        classifier.py, transformer.py, …
+    conf/                       # bundled Hydra configs (model, callbacks)
+    callbacks/                  # Lightning callbacks
+    utils/
+    train.py                    # heptokens-train entry point
+    export_tokens.py            # heptokens-export entry point
+    fit_preprocessors.py        # heptokens-fit-preprocessors entry point
+tests/
+profiles/                       # Slurm executor profiles for pixi tasks
+```
+
+## Adding a New Modality
+
+A modality is any input representation with its own data format, token definition, and reconstruction metric. The VQ-VAE core (`LitVqVae`, `ResidualVQ`) is modality-agnostic and never changes.
+
+### Token strategy
+
+| Data type | Token strategy |
+|---|---|
+| Sparse set (tracks, clusters, PFOs) | One token per object. Boolean validity mask for variable-length padding. |
+| Dense grid (η-φ image, calorimeter towers) | Fixed-size patching in `__getitem__`. Each patch is one token. No mask needed if all patches are always valid. |
+
+### Checklist
+
+**1. Dataset / DataModule** — `src/heptokens/data/<modality>.py`
+
+See `StructuredArrayModule` or `FlatArrayModule` for reference implementations.
+- `MyDataset.__getitem__` returns a dict: constituent tensor + optional bool mask.
+- `MyDataModule` extends `BaseMapModule`; implement `setup()` and `get_data_sample()`.
+
+**2. Encoder / Decoder** — `src/heptokens/models/<modality>_coders.py`
+
+Start from `src/heptokens/models/new_modality_coders.py`.
+- Extend `Encoder` and `Decoder` from `heptokens.models.coders`.
+- `Encoder.forward(batch)` → `[B, N, codebook_dim]`, masked positions zeroed.
+- `Decoder.compute_loss(z_q, batch)` → scalar loss.
+- Set `self.input_key` and `self.mask_key` to match your dataset's output keys.
+
+**3. Model config** — `src/heptokens/conf/model/<modality>_vqvae.yaml`
+
+Start from `src/heptokens/conf/model/new_modality_vqvae.yaml`.
+Point `encoder._target_` and `decoder._target_` at your new classes.
+
+**4. Datamodule config** — in your experiment repo's config directory
+
+```yaml
+# my_configs/datamodule/my_modality.yaml
+_target_: heptokens.data.<modality>.MyDataModule
+data_path: /path/to/data.h5
+batch_size: 512
+num_workers: 4
+# ... your fields
+```
+
+**5. Callbacks config** *(optional)* — `src/heptokens/conf/callbacks/<modality>_encode.yaml`
+
+For physical-space reconstruction metrics, subclass `BaseReconstructionMonitor`:
+
+```python
+from heptokens.callbacks.recon import BaseReconstructionMonitor
+
+class MyMonitor(BaseReconstructionMonitor):
+    def __init__(self, inverse_fn, **kwargs):
+        super().__init__(input_key="my_key", mask_key=None, **kwargs)
+        self.inverse_fn = inverse_fn
+
+    def compute_and_log_metrics(self, trainer, pl_module, batch, batch_idx):
+        import torch as T
+        with T.no_grad():
+            z_q = pl_module.encode(batch)[0]
+            recon = pl_module.decode(z_q, batch)
+        orig = self.inverse_fn(batch["my_key"].cpu())
+        reco = self.inverse_fn(recon.cpu())
+        pl_module.log("val/my_mae_unscaled", T.mean(T.abs(orig - reco)))
+```
+
+**Train:**
+
+```bash
+heptokens-train -cd /path/to/my/configs \
+  model=<modality>_vqvae callbacks=<modality>_encode \
+  datamodule=<modality> trainer.max_epochs=30
+```
+
+### Common pitfalls
+
+- **`input_key` mismatch**: The encoder reads `batch[self.input_key]`. Set it consistently in the dataset and encoder.
+- **`n_classes` not needed for VQ-VAE**: Omit or set to `null` in your datamodule config.
+- **Mask dtype**: The mask must be `bool`, shape `[B, N]`. Float/int masks will silently produce wrong results.
+- **Dense grids**: Make sure `N` (number of patches) is fixed across samples so batches stack without a mask.
+
+---
+
+## Preprocessing
+
+The VQ-VAE expects normalised inputs. The library provides sklearn-based scalers in `heptokens.data.transforms` and a batch-level application function in `heptokens.data.collation`. The workflow is: **fit offline → save as joblib → load in your Hydra config** (shown in [Quick Start step 3](#3-fit-preprocessors)).
+
+### How it connects
+
+`BaseMapModule` accepts a `transforms` kwarg. When set, it wraps the DataLoader's collate function so transforms run on every batch automatically — no manual intervention at train time.
+
+### Available modes (`create_preprocessing_transformer`)
+
+| Mode | Description |
+|------|-------------|
+| `"quantile"` | `QuantileTransformer` on all features |
+| `"standard"` | `StandardScaler` on all features |
+| `"log_quantile"` | Log-scale selected features, then quantile-transform all |
+| `"log_standard"` | Log-scale selected features, then standard-scale all |
+
+Parameters: `mode`, `log_feature_indices` (which features to log-scale), `n_quantiles`, `log_offset`.
+
+### Inverse transforms
+
+For evaluating reconstruction quality in physical units:
+
+```python
+from heptokens.data.collation import inverse_preprocess_batch
+
+original = inverse_preprocess_batch(batch, cst_fn=cst_scaler, jet_fn=jet_scaler)
+```
+
+---
+
+## Troubleshooting
+
+### DataLoader workers exit unexpectedly
+- Set `num_workers=0` to confirm worker issue.
+- Set `pin_memory=False` if worker shutdown appears unstable.
+
+### Training becomes slow
+- Tune `num_workers`, `batch_size`, `pin_memory`.
+- Avoid `reload_dataloaders_every_n_epochs=1` on iterable datasets unless necessary.
+
+### W&B run lands under wrong account
+- Override `logger.entity=<username_or_team>` when calling `heptokens-train`.
+
+### Debugger uses wrong Python interpreter
+- Use `.vscode/launch.json` with `"python": "${workspaceFolder}/.pixi/envs/default/bin/python"`.
+
+---
+
+## Resources
+- [Pixi documentation](https://pixi.sh/)
+- [Reproducible Machine Learning Workflows for Scientists, Matthew Feickert, 2025](https://carpentries-incubator.github.io/reproducible-ml-workflows/)
+- [Good ML project structure](https://github.com/mattcleigh/JetSSL-Lite/tree/master)
+
+## Contributors
+Jeffrey Krupa, Samuel Klein, Michael Kagan