diff --git a/strategy/cli.py b/strategy/cli.py index c3c1501..d84a6d4 100644 --- a/strategy/cli.py +++ b/strategy/cli.py @@ -38,6 +38,10 @@ def build_parser() -> argparse.ArgumentParser: 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="report reconstruction error vs a float64 reference") @@ -55,6 +59,8 @@ def main(argv=None) -> int: 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.spectral_alpha < 0: + raise ValueError(f"--spectral-alpha must be >= 0, got {args.spectral_alpha}") cfg = Config( device=args.device, dtype=args.dtype, @@ -70,7 +76,7 @@ def main(argv=None) -> int: 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) + 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. @@ -79,7 +85,8 @@ def main(argv=None) -> int: 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) + 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 diff --git a/strategy/runner.py b/strategy/runner.py index a3f80fb..51e3104 100644 --- a/strategy/runner.py +++ b/strategy/runner.py @@ -44,7 +44,8 @@ def _flop_ratio_line(ratio: float) -> str: def run(n: int, cfg: Config, fill: str = "lowrank", verify: bool = False, - keep: bool = False, data_rank: int | None = None) -> dict: + keep: bool = False, data_rank: int | None = None, + spectral_alpha: float = 1.0) -> dict: """Generate A, B, compute C = A @ B with the subspace strategy, report. Default fill is 'lowrank' because that is the regime the strategy targets; @@ -67,9 +68,9 @@ def run(n: int, cfg: Config, fill: str = "lowrank", verify: bool = False, try: A = storage.generate(n, dt, on_disk, pa if on_disk else None, cfg.seed, fill, - data_rank=data_rank) + data_rank=data_rank, spectral_alpha=spectral_alpha) B = storage.generate(n, dt, on_disk, pb if on_disk else None, cfg.seed + 1, fill, - data_rank=data_rank) + data_rank=data_rank, spectral_alpha=spectral_alpha) C = storage.allocate(n, dt, on_disk, pc if on_disk else None) info = {} @@ -134,7 +135,8 @@ def _rel_frobenius_streamed(Ce, Cs, block_bytes: int = 256 * 1024**2) -> float: def compare(n: int, cfg: Config, fill: str = "lowrank", - data_rank: int | None = None, keep: bool = False) -> dict: + data_rank: int | None = None, keep: bool = False, + spectral_alpha: float = 1.0) -> dict: """Run the exact baseline and the subspace strategy on the SAME inputs and report time, throughput and the smart strategy's reconstruction error.""" backend = Backend(cfg.device, cfg.verbose) @@ -148,9 +150,9 @@ def compare(n: int, cfg: Config, fill: str = "lowrank", try: A = storage.generate(n, dt, on_disk, pa if on_disk else None, cfg.seed, fill, - data_rank=data_rank) + data_rank=data_rank, spectral_alpha=spectral_alpha) B = storage.generate(n, dt, on_disk, pb if on_disk else None, cfg.seed + 1, fill, - data_rank=data_rank) + data_rank=data_rank, spectral_alpha=spectral_alpha) Ce = storage.allocate(n, dt, on_disk, pe if on_disk else None) Cs = storage.allocate(n, dt, on_disk, ps if on_disk else None) diff --git a/tests/test_cli_errors.py b/tests/test_cli_errors.py index 2a1a48c..83a3596 100644 --- a/tests/test_cli_errors.py +++ b/tests/test_cli_errors.py @@ -47,6 +47,22 @@ def test_bad_data_rank_exits_cleanly(argv, capsys): assert "error:" in capsys.readouterr().err +# --spectral-alpha is the k^-alpha decay exponent for --fill decaying-spectrum; +# a negative exponent would *grow* the tail (not a decaying spectrum), so it is +# rejected on CPU before any device work, like the other knobs. +STRATEGY_BAD_SPECTRAL_ALPHA_ARGS = [ + ["--n", "8", "--spectral-alpha", "-0.5"], + ["--n", "8", "--spectral-alpha", "-1"], +] + + +@pytest.mark.parametrize("argv", STRATEGY_BAD_SPECTRAL_ALPHA_ARGS, ids=lambda a: " ".join(a)) +def test_bad_spectral_alpha_exits_cleanly(argv, capsys): + rc = strategy_cli.main(argv) + assert rc == 2, f"expected exit 2 for {argv}, got {rc}" + assert "error:" in capsys.readouterr().err + + STRATEGY_BAD_RANK_M_ARGS = [ ["--n", "8", "--rank-m", "0"], # non-positive: rank_m must be >= 1 ["--n", "8", "--rank-m", "-2"], # negative: previously hit GPU before subspace diff --git a/tests/test_spectral_alpha.py b/tests/test_spectral_alpha.py new file mode 100644 index 0000000..302bf4d --- /dev/null +++ b/tests/test_spectral_alpha.py @@ -0,0 +1,56 @@ +"""``--spectral-alpha`` exposes storage's k^-alpha decay exponent -- previously +the only ``storage.generate`` knob with no CLI path, so ``--fill +decaying-spectrum`` was stuck at the default alpha=1.0. These pure-NumPy checks +pin (a) that the flag is parsed and (b) that the exponent actually steepens the +singular-value decay, so a larger alpha concentrates more energy in the leading +components. +""" +from __future__ import annotations + +import numpy as np + +from strategy import storage +from strategy.cli import build_parser + + +def test_spectral_alpha_is_parsed(): + args = build_parser().parse_args(["--fill", "decaying-spectrum", + "--spectral-alpha", "2.5"]) + assert args.spectral_alpha == 2.5 + # default is 1.0 (matches storage.generate's default) + assert build_parser().parse_args([]).spectral_alpha == 1.0 + + +def _leading_energy_fraction(mat, k): + s = np.linalg.svd(np.asarray(mat, dtype=np.float64), compute_uv=False) + return float((s[:k] ** 2).sum() / (s ** 2).sum()) + + +def test_larger_alpha_steepens_the_spectrum(): + n, R = 96, 32 + gentle = storage.generate(n, np.float64, False, None, seed=0, + fill="decaying-spectrum", data_rank=R, + spectral_alpha=0.5) + steep = storage.generate(n, np.float64, False, None, seed=0, + fill="decaying-spectrum", data_rank=R, + spectral_alpha=3.0) + # A steeper exponent must pack more of the energy into the top few components. + assert _leading_energy_fraction(steep, 4) > _leading_energy_fraction(gentle, 4) + + +def test_alpha_zero_is_flat_like_lowrank_scale(): + # alpha=0 -> k^0 = 1 for every component -> uniform weights (no decay). + n, R = 64, 16 + flat = storage.generate(n, np.float64, False, None, seed=1, + fill="decaying-spectrum", data_rank=R, + spectral_alpha=0.0) + s = np.linalg.svd(np.asarray(flat, dtype=np.float64), compute_uv=False) + # rank is exactly R and the top-R singular values are of comparable scale + # (no k^-alpha taper), unlike a steep spectrum. + assert np.count_nonzero(s > 1e-9 * s[0]) == R + assert s[R - 1] / s[0] > 0.05 + + +if __name__ == "__main__": + import pytest + raise SystemExit(pytest.main([__file__, "-q"]))