giant/scripts/dwarf.py

"""dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.

Unifies the standalone scripts/*.py conversion, migration, versioning, and
simulation-fanout tools into one Typer app so there's a single command name
(and `--help`) to remember instead of five differently-hyphenated ones.
"""

import os
from enum import Enum
from pathlib import Path
from typing import Optional

import typer
from typing_extensions import Annotated

from giant.config import Conditioning
from scripts.bump_dataset_version import (
    run_bump_gen,
    run_bump_schema,
    run_create_manifest,
    run_status,
    run_update_manifest,
)
from scripts.create_root_files import run_make_root
from scripts.geometry_oracle import run_build_geometry_oracle
from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
from scripts.migrate_geant_steps import run_migration
from scripts.steps_to_parquet import convert_steps_to_parquet
from scripts.steps_to_parquet_parallel import run_parallel_job
from scripts.warm_setup_cache import run_warm_setup_cache

app = typer.Typer(no_args_is_help=True)

_DATASET_ROOT_DEFAULT = Path("/ceph/lbogner/geant_steps")


@app.callback()
def _main() -> None:
    """dwarf — dataset/tooling CLI (ROOT<->parquet conversion, dataset versioning, sim fanout)."""


def _warn_if_exceeds_shared_quota(n: int, flag: str) -> None:
    """Soft warning (never blocks) when a worker/job count looks likely to
    grab more than this repo's documented shared-portal-machine etiquette
    (CLAUDE.md's Compute environment: stay within ~1/4 of CPU/RAM and a
    single GPU, since portal1/deepthought{,2}/bms{1..3} are shared with
    other users). Not a hard cap — a legitimate big machine or a
    deliberately aggressive run is still the caller's call.
    """
    cpu_count = os.cpu_count() or 1
    quota = max(1, cpu_count // 4)
    if n > quota:
        typer.echo(
            f"warning: {flag}={n} exceeds ~1/4 of this machine's "
            f"{cpu_count} CPU(s) ({quota}) — portal machines are shared "
            "with other users (see CLAUDE.md's Compute environment section)",
            err=True,
        )


class Compression(str, Enum):
    snappy = "snappy"
    lz4 = "lz4"
    zstd = "zstd"
    gzip = "gzip"
    none = "none"


class PoolType(str, Enum):
    full = "full"
    holdout = "holdout"
    dev = "dev"


@app.command()
def convert(
    root_files: Annotated[list[Path], typer.Argument(help="Input ROOT file(s)")],
    output: Annotated[
        Optional[Path],
        typer.Option(
            "--output",
            "-o",
            help="Output Parquet file (default: <input>.parquet). Only valid with a single input file and --jobs 1.",
        ),
    ] = None,
    batch_size: Annotated[
        str,
        typer.Option(
            "--batch-size",
            help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
        ),
    ] = "100 MB",
    tree: Annotated[str, typer.Option("--tree", help="Tree name inside the ROOT file")] = "Steps",
    compression: Annotated[
        Compression, typer.Option("--compression", help="Parquet compression codec")
    ] = Compression.snappy,
    jobs: Annotated[
        int,
        typer.Option(
            "--jobs",
            "-j",
            help="Convert N files in parallel, resolving each destination from "
            "--dataset-root/--schema (default: 1, sequential, any file layout)",
        ),
    ] = 1,
    dataset_root: Annotated[
        Path,
        typer.Option(
            "--dataset-root",
            help="Dataset root containing raw/ and processed/ (only used with --jobs > 1)",
        ),
    ] = _DATASET_ROOT_DEFAULT,
    schema: Annotated[
        Optional[str],
        typer.Option(
            "--schema",
            help="Schema tag to write parquets under, e.g. schema2 (only used "
            "with --jobs > 1; default: highest schemaN already under "
            "processed/<kind>/<gen>/)",
        ),
    ] = None,
) -> None:
    """Convert ROOT Steps tree(s) to Parquet."""
    if jobs < 1:
        typer.echo("error: --jobs must be >= 1", err=True)
        raise typer.Exit(1)
    _warn_if_exceeds_shared_quota(jobs, "--jobs")

    compression_value = "uncompressed" if compression is Compression.none else compression.value

    if jobs == 1:
        if output is not None and len(root_files) > 1:
            typer.echo("error: --output can only be used with a single input file", err=True)
            raise typer.Exit(1)
        total_orphaned = 0
        for root_file in root_files:
            _, n_orphaned = convert_steps_to_parquet(
                root_file,
                output_path=output,
                batch_size=batch_size,
                tree_name=tree,
                compression=compression_value,
            )
            total_orphaned += n_orphaned
        if total_orphaned:
            typer.echo(f"\n{total_orphaned} orphaned child track(s) dropped across {len(root_files)} file(s).")
        return

    if output is not None:
        typer.echo(
            "error: --output cannot be combined with --jobs > 1 "
            "(destinations are derived from --dataset-root/--schema)",
            err=True,
        )
        raise typer.Exit(1)
    run_parallel_job(
        [str(f) for f in root_files],
        jobs=jobs,
        dataset_root=dataset_root,
        schema=schema,
        batch_size=batch_size,
        tree=tree,
        compression=compression_value,
    )


@app.command()
def migrate(
    root: Annotated[Path, typer.Argument(help="Dataset root to migrate in place")] = _DATASET_ROOT_DEFAULT,
    execute: Annotated[
        bool,
        typer.Option(
            "--execute",
            help="Actually move/copy files and write manifests (default: dry run)",
        ),
    ] = False,
    copy: Annotated[
        bool,
        typer.Option(
            "--copy",
            help="Copy instead of move, leaving the originals in place (e.g. if another process is still reading them)",
        ),
    ] = False,
) -> None:
    """One-time migration into the versioned raw/processed/pools/derived layout."""
    run_migration(str(root), execute=execute, copy=copy)


@app.command("bump-gen")
def bump_gen(
    reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
        Optional[str],
        typer.Option("--date", help="Override date (default: today, ISO)"),
    ] = None,
    to: Annotated[
        Optional[str],
        typer.Option(
            "--to",
            metavar="genN",
            help="Target gen tag (default: one past the current highest)",
        ),
    ] = None,
    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
    """Cut a new raw generation."""
    run_bump_gen(
        kind=kind,
        reason=reason,
        by=by,
        date=date,
        execute=execute,
        root=str(root),
        to=to,
    )


@app.command("bump-schema")
def bump_schema(
    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
    reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
        Optional[str],
        typer.Option("--date", help="Override date (default: today, ISO)"),
    ] = None,
    to: Annotated[
        Optional[str],
        typer.Option(
            "--to",
            metavar="schemaN",
            help="Target schema tag (default: one past the current highest)",
        ),
    ] = None,
    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
    """Cut a new schema within a gen."""
    run_bump_schema(
        kind=kind,
        gen=gen,
        reason=reason,
        by=by,
        date=date,
        execute=execute,
        root=str(root),
        to=to,
    )


@app.command()
def status(
    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
    """List existing gens/schemas per kind."""
    run_status(str(root))


@app.command("update-manifest")
def update_manifest(
    manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
    schema: Annotated[
        Optional[str],
        typer.Option(
            "--schema",
            metavar="schemaN",
            help="Target schema tag (default: highest schema found in the same gen dir)",
        ),
    ] = None,
    gen: Annotated[
        Optional[str],
        typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
    ] = None,
    execute: Annotated[
        bool,
        typer.Option("--execute", help="Write updated manifests (default: dry run)"),
    ] = False,
) -> None:
    """Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
    run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)


@app.command("create-manifest")
def create_manifest(
    files: Annotated[list[Path], typer.Argument(help="Parquet files to include")],
    output: Annotated[
        Optional[Path],
        typer.Option("--output", "-o", help="Explicit path for the new .manifest file"),
    ] = None,
    pool: Annotated[
        Optional[str],
        typer.Option(
            "--pool",
            metavar="DETECTOR",
            help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
        ),
    ] = None,
    type_: Annotated[
        Optional[PoolType],
        typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
    ] = None,
    root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
    execute: Annotated[bool, typer.Option("--execute", help="Write the manifest (default: dry run)")] = False,
    force: Annotated[
        bool,
        typer.Option("--force", help="Overwrite the manifest if it already exists"),
    ] = False,
) -> None:
    """Create a new manifest from a list of parquet files."""
    run_create_manifest(
        [str(f) for f in files],
        execute=execute,
        output=str(output) if output is not None else None,
        pool=pool,
        type_=type_.value if type_ is not None else None,
        root=str(root),
        force=force,
    )


@app.command("make-root")
def make_root(
    executable: Annotated[Path, typer.Option("--executable", help="Built minicalosim run_* executable")],
    detector: Annotated[
        list[str],
        typer.Option(
            "--detector",
            metavar="NAME[:CONFIG]",
            help="Dataset detector label, optionally with ':CONFIG' to pass as the "
            "executable's config-name argument (e.g. sampling_pb_scint:pb_scint). "
            "Omit ':CONFIG' for executables that take no config selector (e.g. run_pbwo4). "
            "Repeatable.",
        ),
    ],
    num_files: Annotated[int, typer.Option("--num-files", help="New shards to create per detector")],
    events_per_file: Annotated[int, typer.Option("--events-per-file", help="nEvents passed to the executable")],
    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")],
    energy_gev: Annotated[
        float | None,
        typer.Option(
            "--energy-gev",
            help="energy_GeV passed to the executable (default: executable's own "
            "default, currently 1.0). Note the dataset detector label is not "
            "derived from this — e.g. use '--detector pbwo4_10gev --energy-gev 10' "
            "to name the dataset accordingly.",
        ),
    ] = None,
    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
    dataset_root: Annotated[Path, typer.Option("--dataset-root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
    jobs: Annotated[int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")] = 4,
    execute: Annotated[
        bool,
        typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)"),
    ] = False,
) -> None:
    """Generate new ROOT shards via a minicalosim executable."""
    _warn_if_exceeds_shared_quota(jobs, "--jobs")
    run_make_root(
        executable=executable,
        detector=detector,
        num_files=num_files,
        events_per_file=events_per_file,
        kind=kind,
        gen=gen,
        dataset_root=str(dataset_root),
        jobs=jobs,
        execute=execute,
        energy_gev=energy_gev,
    )


class OracleMethod(str, Enum):
    slab = "slab"
    knn = "knn"
    svm = "svm"


@app.command("build-geometry-oracle")
def build_geometry_oracle(
    data: Annotated[Path, typer.Argument(help="Steps parquet file or directory of steps files")],
    out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
    method: Annotated[
        OracleMethod,
        typer.Option(
            "--method",
            help=(
                "Lookup strategy: slab (default; exact O(log #segments) fast "
                "path for the layered-slab detector geometry), knn, or svm "
                "(generic fallbacks for non-slab geometries)"
            ),
        ),
    ] = OracleMethod.slab,
    k: Annotated[int, typer.Option("--k", help="Neighbours for the knn classifier")] = 1,
    subsample: Annotated[
        int,
        typer.Option("--subsample", help="Max reference points sampled from the data"),
    ] = 500_000,
    escape_factor: Annotated[
        float,
        typer.Option(
            "--escape-factor",
            help="escape_threshold = this x median spacing of reference points",
        ),
    ] = 5.0,
    seed: Annotated[int, typer.Option("--seed", help="Sampling seed")] = 0,
    depth_axis: Annotated[
        int,
        typer.Option(
            "--depth-axis",
            help="0/1/2 -> x/y/z axis the layers stack along (method=slab only)",
        ),
    ] = 2,
    n_bins: Annotated[
        int,
        typer.Option(
            "--n-bins",
            help="Depth-axis resolution, finer than the thinnest layer (method=slab only)",
        ),
    ] = 2000,
) -> None:
    """Fit a position -> (material, layer_id) oracle for `giant rollout`."""
    run_build_geometry_oracle(
        data=data,
        out=out,
        method=method.value,
        k=k,
        subsample=subsample,
        escape_factor=escape_factor,
        seed=seed,
        depth_axis=depth_axis,
        n_bins=n_bins,
    )


@app.command("warm-cache")
def warm_cache(
    data: Annotated[
        Path,
        typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
    ],
    val_fraction: Annotated[
        float,
        typer.Option(
            "--val-fraction",
            "-f",
            help="Must match the `giant train` run(s) to warm for",
        ),
    ] = 0.1,
    seed: Annotated[
        int,
        typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
    ] = 0,
    particle_conditioning: Annotated[
        Conditioning,
        typer.Option(
            "--particle-conditioning",
            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for",
        ),
    ] = Conditioning.physical,
    material_conditioning: Annotated[
        Conditioning,
        typer.Option(
            "--material-conditioning",
            help="Must match the `giant train` run(s)' conditioning.material.type "
            "to warm for — independent of --particle-conditioning "
            "(the two axes may differ)",
        ),
    ] = Conditioning.physical,
    router: Annotated[
        bool,
        typer.Option(
            "--router/--no-router",
            help="Warm the process vocabulary too (only takes effect with --router-type process)",
        ),
    ] = False,
    router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
    n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
    rebuild: Annotated[
        bool,
        typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
    ] = False,
) -> None:
    """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.

    Warms the vocab maps, event-id split index, and the normalizer entry for
    the given --val-fraction/--seed/--particle-conditioning/
    --material-conditioning, so a later `giant train` run (or a `dwarf
    hparam-scan` sweep, which shares one such entry across every run) skips
    straight to training. See giant/data/setup_cache.py.
    """
    run_warm_setup_cache(
        data=str(data),
        val_fraction=val_fraction,
        seed=seed,
        particle_conditioning=particle_conditioning.value,
        material_conditioning=material_conditioning.value,
        router_enabled=router,
        router_type=router_type,
        n_experts=n_experts,
        rebuild=rebuild,
        echo=typer.echo,
    )


@app.command("hparam-scan")
def hparam_scan(
    data: Annotated[str, typer.Option("--data")] = DATA_DEFAULT,
    scan_dir: Annotated[str, typer.Option("--scan-dir")] = SCAN_DIR_DEFAULT,
    seed: Annotated[int, typer.Option("--seed")] = 0,
    dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
) -> None:
    """Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
    run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)


if __name__ == "__main__":
    app()