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.
"""

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

import typer
from typing_extensions import Annotated

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.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

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)."""


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)

    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)
        for root_file in root_files:
            convert_steps_to_parquet(
                root_file,
                output_path=output,
                batch_size=batch_size,
                tree_name=tree,
                compression=compression_value,
            )
        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,
) -> 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),
    )


@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")
    ],
    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."""
    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,
    )


@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()