cli.py

"""Command-line interface:  python -m strategy --n 8192 --transform rsvd ..."""
from __future__ import annotations

import argparse
import math
import sys

from .config import Config, DTYPES
from . import runner
from .transforms import available as available_transforms, get_transform


def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        prog="strategy",
        description="Smart (subspace) matrix x matrix: compress -> compute -> "
                    "reconstruct. Approximate; for low-rank / smooth data.",
    )
    p.add_argument("--n", type=int, default=8192, help="matrix dimension n (n x n). default 8192")
    p.add_argument("--dtype", choices=list(DTYPES), default="fp32")
    p.add_argument("--device", type=int, default=0, help="CUDA device index")
    p.add_argument("--rank-m", type=int, default=None,
                   help="subspace dimension M (default min(n, max(64, n//8))). "
                        "Smaller = faster + less accurate")
    p.add_argument("--transform", default="rsvd",
                   help=f"subspace basis (the core tech). one of: "
                        f"{available_transforms()}")
    p.add_argument("--transform-seed", type=int, default=0)
    p.add_argument("--vram-fraction", type=float, default=0.6)
    p.add_argument("--compare", action="store_true",
                   help="run BOTH exact and smart on the same inputs and "
                        "report speed + accuracy")
    p.add_argument("--storage", choices=["auto", "ram", "disk"], default="auto")
    p.add_argument("--workdir", default="./_strategy_data")
    p.add_argument("--fill", choices=["lowrank", "random", "decaying-spectrum", "iota", "zeros"],
                   default="lowrank",
                   help="test-matrix content. 'lowrank' = compressible data "
                        "where the strategy is accurate (default)")
    p.add_argument("--data-rank", type=int, default=None,
                   help="rank of generated matrices for --fill lowrank / "
                        "decaying-spectrum (default max(1, n//32))")
    p.add_argument("--spectral-alpha", type=float, default=1.0,
                   help="singular-value decay exponent k^-alpha for "
                        "--fill decaying-spectrum (default 1.0; larger = faster "
                        "decay). Ignored by other fills.")
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--verify", action="store_true",
                   help="check reconstruction error vs a float64 reference "
                        "(exit 1 if above the dtype tolerance)")
    p.add_argument("--keep", action="store_true", help="keep on-disk files")
    p.add_argument("--quiet", action="store_true")
    return p


def main(argv=None) -> int:
    args = build_parser().parse_args(argv)
    try:
        if args.n < 1:
            raise ValueError(f"--n must be a positive integer, got {args.n}")
        if args.data_rank is not None and args.data_rank < 1:
            raise ValueError(f"--data-rank must be a positive integer, got {args.data_rank}")
        if args.rank_m is not None and args.rank_m < 1:
            raise ValueError(f"--rank-m must be a positive integer, got {args.rank_m}")
        if args.rank_m is not None and args.rank_m > args.n:
            raise ValueError(f"--rank-m must be <= --n ({args.n}), got {args.rank_m}")
        if not math.isfinite(args.spectral_alpha) or args.spectral_alpha < 0:
            raise ValueError(
                f"--spectral-alpha must be a finite number >= 0, got {args.spectral_alpha}"
            )
        cfg = Config(
            device=args.device,
            dtype=args.dtype,
            rank_m=args.rank_m,
            transform=args.transform,
            transform_seed=args.transform_seed,
            vram_fraction=args.vram_fraction,
            storage=args.storage,
            workdir=args.workdir,
            seed=args.seed,
            verbose=not args.quiet,
        )
        get_transform(cfg.transform, cfg.transform_seed)
        if args.compare:
            out = runner.compare(args.n, cfg, fill=args.fill, data_rank=args.data_rank,
                                 keep=args.keep, spectral_alpha=args.spectral_alpha)
            if args.quiet:
                # compare() only prints when verbose (i.e. not --quiet), so
                # without this a --compare --quiet run produced NO output at all.
                print(f"exact {out['exact_seconds']:.4f}s  "
                      f"smart {out['smart_seconds']:.4f}s  "
                      f"speedup {out['speedup']:.2f}x  rel_err {out['rel_err']:.2e}")
            return 0
        info = runner.run(args.n, cfg, fill=args.fill, verify=args.verify,
                          keep=args.keep, data_rank=args.data_rank,
                          spectral_alpha=args.spectral_alpha)
    except (ValueError, RuntimeError, MemoryError, KeyError) as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    if args.quiet:
        err = info.get("verify", {}).get("max_rel_err")
        tail = f"  rel_err={err:.2e}" if err is not None else ""
        print(f"{info['mode']}  {info['seconds']:.4f}s{tail}")
    # Fail only on an explicit mismatch. A skipped verify (large / disk-backed n)
    # returns {"skipped": ...} with no "ok" key -- not a failure. Mirrors matmul.
    v = info.get("verify")
    if v and v.get("ok") is False:
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())