From 91f0343e420302569b3c0b729d131ab1880ec57c Mon Sep 17 00:00:00 2001 From: CyrilJl Date: Fri, 20 Mar 2026 00:07:00 +0100 Subject: [PATCH 1/6] Add OptiMask performance tooling --- .gitignore | 6 +- dev/perf/TRIAL_NOTES.md | 16 ++++ dev/perf/benchmark_optimask.py | 154 +++++++++++++++++++++++++++++++++ dev/perf/branch_registry.json | 51 +++++++++++ dev/perf/profile_optimask.py | 43 +++++++++ 5 files changed, 269 insertions(+), 1 deletion(-) create mode 100644 dev/perf/TRIAL_NOTES.md create mode 100644 dev/perf/benchmark_optimask.py create mode 100644 dev/perf/branch_registry.json create mode 100644 dev/perf/profile_optimask.py diff --git a/.gitignore b/.gitignore index f0be60f..0da1289 100644 --- a/.gitignore +++ b/.gitignore @@ -8,7 +8,11 @@ __pycache__/ # development -dev/ +dev/* +!dev/perf/ +!dev/perf/*.py +!dev/perf/*.json +!dev/perf/*.md # Distribution / packaging .Python diff --git a/dev/perf/TRIAL_NOTES.md b/dev/perf/TRIAL_NOTES.md new file mode 100644 index 0000000..f819c8c --- /dev/null +++ b/dev/perf/TRIAL_NOTES.md @@ -0,0 +1,16 @@ +# OptiMask Performance Trials + +Shared baseline for all trial branches: + +- Benchmark target: the large-array scenarios from `tests/test_optimask.py::test_speed` +- Cases: `(100000, 1000)` and `(1000, 100000)` +- Ratio: `0.02` +- Benchmark script: `python dev/perf/benchmark_optimask.py --record` +- Profiling script: `python dev/perf/profile_optimask.py` +- Constraint: keep `n_tries` unchanged so comparisons stay valid + +Trial log: + +- `perf/tooling-base`: baseline + - `(100000, 1000)`: `188.49 ms` mean + - `(1000, 100000)`: `227.33 ms` mean diff --git a/dev/perf/benchmark_optimask.py b/dev/perf/benchmark_optimask.py new file mode 100644 index 0000000..a90c98a --- /dev/null +++ b/dev/perf/benchmark_optimask.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import argparse +import json +import subprocess +from datetime import datetime, timezone +from pathlib import Path +from time import perf_counter + +import numpy as np + +from optimask import OptiMask + + +ROOT = Path(__file__).resolve().parents[2] +REGISTRY_PATH = ROOT / "dev" / "perf" / "branch_registry.json" + + +def generate_random(m: int, n: int, ratio: float, seed: int) -> np.ndarray: + arr = np.zeros((m, n), dtype=np.float32) + nan_count = int(ratio * m * n) + rng = np.random.default_rng(seed) + indices = rng.choice(m * n, nan_count, replace=False) + arr.flat[indices] = np.nan + return arr + + +def get_branch_name() -> str: + result = subprocess.run( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], + check=True, + capture_output=True, + text=True, + cwd=ROOT, + ) + return result.stdout.strip() + + +def benchmark_case( + *, + shape: tuple[int, int], + ratio: float, + data_seed: int, + solve_seed: int, + repeats: int, +) -> dict[str, object]: + rows, cols = shape + x = generate_random(rows, cols, ratio, seed=data_seed) + solver = OptiMask(random_state=solve_seed) + + solver.solve(x) + + timings_ms: list[float] = [] + for _ in range(repeats): + start = perf_counter() + solver.solve(x) + timings_ms.append(1e3 * (perf_counter() - start)) + + return { + "shape": [rows, cols], + "ratio": ratio, + "timings_ms": timings_ms, + "mean_ms": float(np.mean(timings_ms)), + "median_ms": float(np.median(timings_ms)), + "min_ms": float(np.min(timings_ms)), + "max_ms": float(np.max(timings_ms)), + } + + +def load_registry() -> dict[str, object]: + if REGISTRY_PATH.exists(): + return json.loads(REGISTRY_PATH.read_text()) + return {"branches": {}} + + +def write_registry(registry: dict[str, object]) -> None: + REGISTRY_PATH.write_text(json.dumps(registry, indent=2) + "\n") + + +def record_run( + *, + registry: dict[str, object], + branch: str, + cases: list[dict[str, object]], + repeats: int, + ratio: float, + solve_seed: int, + data_seed_base: int, +) -> None: + branch_runs = registry.setdefault("branches", {}).setdefault(branch, []) + branch_runs.append( + { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "repeats": repeats, + "ratio": ratio, + "solve_seed": solve_seed, + "data_seed_base": data_seed_base, + "cases": cases, + } + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Benchmark OptiMask on large arrays.") + parser.add_argument("--ratio", type=float, default=0.02) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--solve-seed", type=int, default=99) + parser.add_argument("--data-seed-base", type=int, default=1000) + parser.add_argument( + "--shape", + action="append", + default=[], + help="Shape formatted as rows,cols. Can be repeated. Defaults to test_speed shapes.", + ) + parser.add_argument("--record", action="store_true", help="Record the run in branch_registry.json.") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + shapes = args.shape or ["100000,1000", "1000,100000"] + + cases: list[dict[str, object]] = [] + for index, item in enumerate(shapes): + rows_str, cols_str = item.split(",") + case = benchmark_case( + shape=(int(rows_str), int(cols_str)), + ratio=args.ratio, + data_seed=args.data_seed_base + index, + solve_seed=args.solve_seed, + repeats=args.repeats, + ) + cases.append(case) + + branch = get_branch_name() + payload = {"branch": branch, "cases": cases} + print(json.dumps(payload, indent=2)) + + if args.record: + registry = load_registry() + record_run( + registry=registry, + branch=branch, + cases=cases, + repeats=args.repeats, + ratio=args.ratio, + solve_seed=args.solve_seed, + data_seed_base=args.data_seed_base, + ) + write_registry(registry) + + +if __name__ == "__main__": + main() diff --git a/dev/perf/branch_registry.json b/dev/perf/branch_registry.json new file mode 100644 index 0000000..a6eb56d --- /dev/null +++ b/dev/perf/branch_registry.json @@ -0,0 +1,51 @@ +{ + "branches": { + "perf/tooling-base": [ + { + "timestamp_utc": "2026-03-19T23:05:29.345043+00:00", + "repeats": 5, + "ratio": 0.02, + "solve_seed": 99, + "data_seed_base": 1000, + "cases": [ + { + "shape": [ + 100000, + 1000 + ], + "ratio": 0.02, + "timings_ms": [ + 185.20669999998063, + 193.7179999658838, + 189.09100000746548, + 184.8823999753222, + 189.54469996970147 + ], + "mean_ms": 188.4885599836707, + "median_ms": 189.09100000746548, + "min_ms": 184.8823999753222, + "max_ms": 193.7179999658838 + }, + { + "shape": [ + 1000, + 100000 + ], + "ratio": 0.02, + "timings_ms": [ + 221.3730999501422, + 225.12890002690256, + 221.29289992153645, + 249.7455000411719, + 219.1337998956442 + ], + "mean_ms": 227.33483996707946, + "median_ms": 221.3730999501422, + "min_ms": 219.1337998956442, + "max_ms": 249.7455000411719 + } + ] + } + ] + } +} diff --git a/dev/perf/profile_optimask.py b/dev/perf/profile_optimask.py new file mode 100644 index 0000000..bea09c9 --- /dev/null +++ b/dev/perf/profile_optimask.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +import argparse +import cProfile +import pstats +from pathlib import Path + +from benchmark_optimask import generate_random +from optimask import OptiMask + + +ROOT = Path(__file__).resolve().parents[2] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Profile a single OptiMask solve call.") + parser.add_argument("--rows", type=int, default=100000) + parser.add_argument("--cols", type=int, default=1000) + parser.add_argument("--ratio", type=float, default=0.02) + parser.add_argument("--data-seed", type=int, default=1000) + parser.add_argument("--solve-seed", type=int, default=99) + parser.add_argument("--sort", default="cumtime", choices=["cumtime", "tottime", "ncalls"]) + parser.add_argument("--limit", type=int, default=25) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + x = generate_random(args.rows, args.cols, args.ratio, seed=args.data_seed) + solver = OptiMask(random_state=args.solve_seed) + solver.solve(x) + + profiler = cProfile.Profile() + profiler.enable() + solver.solve(x) + profiler.disable() + + stats = pstats.Stats(profiler).strip_dirs().sort_stats(args.sort) + stats.print_stats(args.limit) + + +if __name__ == "__main__": + main() From 7d2e7fe5b2b901bd596b4f2339c08fc03cbe690f Mon Sep 17 00:00:00 2001 From: CyrilJl Date: Fri, 20 Mar 2026 00:13:35 +0100 Subject: [PATCH 2/6] Trial 5: track NaN columns during preprocess --- optimask/_optimask.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/optimask/_optimask.py b/optimask/_optimask.py index a7d0d69..a602c5a 100644 --- a/optimask/_optimask.py +++ b/optimask/_optimask.py @@ -153,6 +153,7 @@ def _preprocess(x): iy, ix = np.empty(m * n, dtype=np.uint32), np.empty(m * n, dtype=np.uint32) cols_index_mapper = -np.ones(n, dtype=np.int32) rows_with_nan = np.empty(m, dtype=np.uint32) + cols_with_nan = np.empty(n, dtype=np.uint32) n_rows_with_nan = 0 n_cols_with_nan = 0 cnt = 0 @@ -168,6 +169,7 @@ def _preprocess(x): else: ix[cnt] = n_cols_with_nan cols_index_mapper[j] = n_cols_with_nan + cols_with_nan[n_cols_with_nan] = j n_cols_with_nan += 1 cnt += 1 @@ -177,9 +179,7 @@ def _preprocess(x): iy, ix = iy[:cnt], ix[:cnt] rows_with_nan = rows_with_nan[:n_rows_with_nan] - cols_with_nan = np.flatnonzero(cols_index_mapper >= 0)[ - cols_index_mapper[cols_index_mapper >= 0].argsort() - ].astype(np.uint32) + cols_with_nan = cols_with_nan[:n_cols_with_nan] return iy, ix, rows_with_nan, cols_with_nan def _trial(self, k, rng, m_nan, n_nan, iy, ix, m, n): From 78c5f710292f0fb7d4d77f45c1d801f9034bc634 Mon Sep 17 00:00:00 2001 From: CyrilJl Date: Fri, 20 Mar 2026 00:18:09 +0100 Subject: [PATCH 3/6] Record OptiMask trial benchmark results --- dev/perf/TRIAL_NOTES.md | 32 +++++++++++ dev/perf/branch_registry.json | 100 ++++++++++++++++++++++++++++++++++ 2 files changed, 132 insertions(+) diff --git a/dev/perf/TRIAL_NOTES.md b/dev/perf/TRIAL_NOTES.md index f819c8c..f874690 100644 --- a/dev/perf/TRIAL_NOTES.md +++ b/dev/perf/TRIAL_NOTES.md @@ -14,3 +14,35 @@ Trial log: - `perf/tooling-base`: baseline - `(100000, 1000)`: `188.49 ms` mean - `(1000, 100000)`: `227.33 ms` mean +- `perf/trial-01-numpy-preprocess`: regression + - Strategy: replace numba preprocessing with `np.isnan(...).nonzero()` and `np.unique(...)` + - Result: `641.04 ms` / `647.30 ms` +- `perf/trial-02-preprocess-two-pass`: regression + - Strategy: exact-size two-pass numba preprocess + - Result: `238.73 ms` / `243.52 ms` +- `perf/trial-03-no-parallel-kernels`: regression + - Strategy: remove `parallel=True` from hot numba kernels + - Result: `218.26 ms` / `246.48 ms` +- `perf/trial-04-fused-step-kernels`: regression + - Strategy: fuse alternating row/column update kernels + - Result: `197.77 ms` / `245.14 ms` +- `perf/trial-05-preprocess-direct-cols`: best overall + - Strategy: keep the fast one-pass preprocess but build `cols_with_nan` incrementally instead of reconstructing it with boolean filtering and `argsort` + - Result: `186.94 ms` / `206.32 ms` +- `perf/trial-06-fast-solve-dispatch`: mixed + - Strategy: cache optional `polars` import and avoid redundant `np.asarray(X)` work + - Result: `190.44 ms` / `203.51 ms` +- `perf/trial-07-groupby-manual-compare`: regression + - Strategy: replace `max(...)` inside `groupby_max` with a manual comparison + - Result: `196.53 ms` / `213.05 ms` +- `perf/trial-08-manual-rectangle-scan`: close second + - Strategy: compute the largest rectangle with a scalar scan instead of temporary arrays + - Result: `190.44 ms` / `203.20 ms` +- `perf/trial-09-rectangle-plus-fast-dispatch`: mixed + - Strategy: combine trial 8 with the solve-dispatch cleanup from trial 6 + - Result: `190.78 ms` / `204.17 ms` + +Recommendation: + +- Keep `perf/trial-05-preprocess-direct-cols` as the primary improvement branch. +- Consider cherry-picking trial 8 separately only if follow-up benchmarks on your target hardware confirm the horizontal-array gain is worth the extra code change. diff --git a/dev/perf/branch_registry.json b/dev/perf/branch_registry.json index a6eb56d..c9a777c 100644 --- a/dev/perf/branch_registry.json +++ b/dev/perf/branch_registry.json @@ -1,4 +1,57 @@ { + "best_branch": "perf/trial-05-preprocess-direct-cols", + "trial_summary": { + "perf/tooling-base": { + "status": "baseline", + "vertical_mean_ms": 188.4885599836707, + "horizontal_mean_ms": 227.33483996707946 + }, + "perf/trial-01-numpy-preprocess": { + "status": "regression", + "vertical_mean_ms": 641.0418800078332, + "horizontal_mean_ms": 647.3005800042301 + }, + "perf/trial-02-preprocess-two-pass": { + "status": "regression", + "vertical_mean_ms": 238.73366001062095, + "horizontal_mean_ms": 243.5193199897185 + }, + "perf/trial-03-no-parallel-kernels": { + "status": "regression", + "vertical_mean_ms": 218.25817998033017, + "horizontal_mean_ms": 246.4767999947071 + }, + "perf/trial-04-fused-step-kernels": { + "status": "regression", + "vertical_mean_ms": 197.76684001553804, + "horizontal_mean_ms": 245.14067999552935 + }, + "perf/trial-05-preprocess-direct-cols": { + "status": "best_overall", + "vertical_mean_ms": 186.93744002375752, + "horizontal_mean_ms": 206.32389998063445 + }, + "perf/trial-06-fast-solve-dispatch": { + "status": "mixed", + "vertical_mean_ms": 190.4429600108415, + "horizontal_mean_ms": 203.50568001158535 + }, + "perf/trial-07-groupby-manual-compare": { + "status": "regression", + "vertical_mean_ms": 196.5345600154251, + "horizontal_mean_ms": 213.05025999899954 + }, + "perf/trial-08-manual-rectangle-scan": { + "status": "close_second", + "vertical_mean_ms": 190.44409999623895, + "horizontal_mean_ms": 203.20151997730136 + }, + "perf/trial-09-rectangle-plus-fast-dispatch": { + "status": "mixed", + "vertical_mean_ms": 190.7827200135216, + "horizontal_mean_ms": 204.16883998550475 + } + }, "branches": { "perf/tooling-base": [ { @@ -46,6 +99,53 @@ } ] } + ], + "perf/trial-08-manual-rectangle-scan": [ + { + "timestamp_utc": "2026-03-19T23:17:04.441246+00:00", + "repeats": 5, + "ratio": 0.02, + "solve_seed": 99, + "data_seed_base": 1000, + "cases": [ + { + "shape": [ + 100000, + 1000 + ], + "ratio": 0.02, + "timings_ms": [ + 184.71750000026077, + 203.57519993558526, + 182.02970002312213, + 196.70490000862628, + 185.1932000136003 + ], + "mean_ms": 190.44409999623895, + "median_ms": 185.1932000136003, + "min_ms": 182.02970002312213, + "max_ms": 203.57519993558526 + }, + { + "shape": [ + 1000, + 100000 + ], + "ratio": 0.02, + "timings_ms": [ + 194.40169993322343, + 217.1362000517547, + 202.55619997624308, + 193.83589993230999, + 208.0775999929756 + ], + "mean_ms": 203.20151997730136, + "median_ms": 202.55619997624308, + "min_ms": 193.83589993230999, + "max_ms": 217.1362000517547 + } + ] + } ] } } From 6aead48e8f13388d3c038ae1f9ea002958188927 Mon Sep 17 00:00:00 2001 From: CyrilJl Date: Fri, 20 Mar 2026 00:18:58 +0100 Subject: [PATCH 4/6] Update best-branch benchmark summary --- dev/perf/TRIAL_NOTES.md | 2 +- dev/perf/branch_registry.json | 51 +++++++++++++++++++++++++++++++++-- 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/dev/perf/TRIAL_NOTES.md b/dev/perf/TRIAL_NOTES.md index f874690..5b1301b 100644 --- a/dev/perf/TRIAL_NOTES.md +++ b/dev/perf/TRIAL_NOTES.md @@ -28,7 +28,7 @@ Trial log: - Result: `197.77 ms` / `245.14 ms` - `perf/trial-05-preprocess-direct-cols`: best overall - Strategy: keep the fast one-pass preprocess but build `cols_with_nan` incrementally instead of reconstructing it with boolean filtering and `argsort` - - Result: `186.94 ms` / `206.32 ms` + - Result: `185.92 ms` / `202.79 ms` - `perf/trial-06-fast-solve-dispatch`: mixed - Strategy: cache optional `polars` import and avoid redundant `np.asarray(X)` work - Result: `190.44 ms` / `203.51 ms` diff --git a/dev/perf/branch_registry.json b/dev/perf/branch_registry.json index c9a777c..5315db6 100644 --- a/dev/perf/branch_registry.json +++ b/dev/perf/branch_registry.json @@ -28,8 +28,8 @@ }, "perf/trial-05-preprocess-direct-cols": { "status": "best_overall", - "vertical_mean_ms": 186.93744002375752, - "horizontal_mean_ms": 206.32389998063445 + "vertical_mean_ms": 185.9151199925691, + "horizontal_mean_ms": 202.78805999550968 }, "perf/trial-06-fast-solve-dispatch": { "status": "mixed", @@ -146,6 +146,53 @@ } ] } + ], + "perf/trial-05-preprocess-direct-cols": [ + { + "timestamp_utc": "2026-03-19T23:18:32.653025+00:00", + "repeats": 5, + "ratio": 0.02, + "solve_seed": 99, + "data_seed_base": 1000, + "cases": [ + { + "shape": [ + 100000, + 1000 + ], + "ratio": 0.02, + "timings_ms": [ + 183.97919996641576, + 192.23649997729808, + 180.98800000734627, + 180.4888000478968, + 191.88309996388853 + ], + "mean_ms": 185.9151199925691, + "median_ms": 183.97919996641576, + "min_ms": 180.4888000478968, + "max_ms": 192.23649997729808 + }, + { + "shape": [ + 1000, + 100000 + ], + "ratio": 0.02, + "timings_ms": [ + 198.4950000187382, + 197.51129997894168, + 196.86200004070997, + 221.5913999825716, + 199.48059995658696 + ], + "mean_ms": 202.78805999550968, + "median_ms": 198.4950000187382, + "min_ms": 196.86200004070997, + "max_ms": 221.5913999825716 + } + ] + } ] } } From c91b16b43de6c20f01e2767f2e18f5df28c1e410 Mon Sep 17 00:00:00 2001 From: CyrilJl Date: Fri, 20 Mar 2026 18:02:45 +0100 Subject: [PATCH 5/6] Ignore dev folder on trial 05 branch --- .gitignore | 6 +- dev/perf/TRIAL_NOTES.md | 48 -------- dev/perf/benchmark_optimask.py | 154 ------------------------- dev/perf/branch_registry.json | 198 --------------------------------- dev/perf/profile_optimask.py | 43 ------- 5 files changed, 1 insertion(+), 448 deletions(-) delete mode 100644 dev/perf/TRIAL_NOTES.md delete mode 100644 dev/perf/benchmark_optimask.py delete mode 100644 dev/perf/branch_registry.json delete mode 100644 dev/perf/profile_optimask.py diff --git a/.gitignore b/.gitignore index 0da1289..f0be60f 100644 --- a/.gitignore +++ b/.gitignore @@ -8,11 +8,7 @@ __pycache__/ # development -dev/* -!dev/perf/ -!dev/perf/*.py -!dev/perf/*.json -!dev/perf/*.md +dev/ # Distribution / packaging .Python diff --git a/dev/perf/TRIAL_NOTES.md b/dev/perf/TRIAL_NOTES.md deleted file mode 100644 index 5b1301b..0000000 --- a/dev/perf/TRIAL_NOTES.md +++ /dev/null @@ -1,48 +0,0 @@ -# OptiMask Performance Trials - -Shared baseline for all trial branches: - -- Benchmark target: the large-array scenarios from `tests/test_optimask.py::test_speed` -- Cases: `(100000, 1000)` and `(1000, 100000)` -- Ratio: `0.02` -- Benchmark script: `python dev/perf/benchmark_optimask.py --record` -- Profiling script: `python dev/perf/profile_optimask.py` -- Constraint: keep `n_tries` unchanged so comparisons stay valid - -Trial log: - -- `perf/tooling-base`: baseline - - `(100000, 1000)`: `188.49 ms` mean - - `(1000, 100000)`: `227.33 ms` mean -- `perf/trial-01-numpy-preprocess`: regression - - Strategy: replace numba preprocessing with `np.isnan(...).nonzero()` and `np.unique(...)` - - Result: `641.04 ms` / `647.30 ms` -- `perf/trial-02-preprocess-two-pass`: regression - - Strategy: exact-size two-pass numba preprocess - - Result: `238.73 ms` / `243.52 ms` -- `perf/trial-03-no-parallel-kernels`: regression - - Strategy: remove `parallel=True` from hot numba kernels - - Result: `218.26 ms` / `246.48 ms` -- `perf/trial-04-fused-step-kernels`: regression - - Strategy: fuse alternating row/column update kernels - - Result: `197.77 ms` / `245.14 ms` -- `perf/trial-05-preprocess-direct-cols`: best overall - - Strategy: keep the fast one-pass preprocess but build `cols_with_nan` incrementally instead of reconstructing it with boolean filtering and `argsort` - - Result: `185.92 ms` / `202.79 ms` -- `perf/trial-06-fast-solve-dispatch`: mixed - - Strategy: cache optional `polars` import and avoid redundant `np.asarray(X)` work - - Result: `190.44 ms` / `203.51 ms` -- `perf/trial-07-groupby-manual-compare`: regression - - Strategy: replace `max(...)` inside `groupby_max` with a manual comparison - - Result: `196.53 ms` / `213.05 ms` -- `perf/trial-08-manual-rectangle-scan`: close second - - Strategy: compute the largest rectangle with a scalar scan instead of temporary arrays - - Result: `190.44 ms` / `203.20 ms` -- `perf/trial-09-rectangle-plus-fast-dispatch`: mixed - - Strategy: combine trial 8 with the solve-dispatch cleanup from trial 6 - - Result: `190.78 ms` / `204.17 ms` - -Recommendation: - -- Keep `perf/trial-05-preprocess-direct-cols` as the primary improvement branch. -- Consider cherry-picking trial 8 separately only if follow-up benchmarks on your target hardware confirm the horizontal-array gain is worth the extra code change. diff --git a/dev/perf/benchmark_optimask.py b/dev/perf/benchmark_optimask.py deleted file mode 100644 index a90c98a..0000000 --- a/dev/perf/benchmark_optimask.py +++ /dev/null @@ -1,154 +0,0 @@ -from __future__ import annotations - -import argparse -import json -import subprocess -from datetime import datetime, timezone -from pathlib import Path -from time import perf_counter - -import numpy as np - -from optimask import OptiMask - - -ROOT = Path(__file__).resolve().parents[2] -REGISTRY_PATH = ROOT / "dev" / "perf" / "branch_registry.json" - - -def generate_random(m: int, n: int, ratio: float, seed: int) -> np.ndarray: - arr = np.zeros((m, n), dtype=np.float32) - nan_count = int(ratio * m * n) - rng = np.random.default_rng(seed) - indices = rng.choice(m * n, nan_count, replace=False) - arr.flat[indices] = np.nan - return arr - - -def get_branch_name() -> str: - result = subprocess.run( - ["git", "rev-parse", "--abbrev-ref", "HEAD"], - check=True, - capture_output=True, - text=True, - cwd=ROOT, - ) - return result.stdout.strip() - - -def benchmark_case( - *, - shape: tuple[int, int], - ratio: float, - data_seed: int, - solve_seed: int, - repeats: int, -) -> dict[str, object]: - rows, cols = shape - x = generate_random(rows, cols, ratio, seed=data_seed) - solver = OptiMask(random_state=solve_seed) - - solver.solve(x) - - timings_ms: list[float] = [] - for _ in range(repeats): - start = perf_counter() - solver.solve(x) - timings_ms.append(1e3 * (perf_counter() - start)) - - return { - "shape": [rows, cols], - "ratio": ratio, - "timings_ms": timings_ms, - "mean_ms": float(np.mean(timings_ms)), - "median_ms": float(np.median(timings_ms)), - "min_ms": float(np.min(timings_ms)), - "max_ms": float(np.max(timings_ms)), - } - - -def load_registry() -> dict[str, object]: - if REGISTRY_PATH.exists(): - return json.loads(REGISTRY_PATH.read_text()) - return {"branches": {}} - - -def write_registry(registry: dict[str, object]) -> None: - REGISTRY_PATH.write_text(json.dumps(registry, indent=2) + "\n") - - -def record_run( - *, - registry: dict[str, object], - branch: str, - cases: list[dict[str, object]], - repeats: int, - ratio: float, - solve_seed: int, - data_seed_base: int, -) -> None: - branch_runs = registry.setdefault("branches", {}).setdefault(branch, []) - branch_runs.append( - { - "timestamp_utc": datetime.now(timezone.utc).isoformat(), - "repeats": repeats, - "ratio": ratio, - "solve_seed": solve_seed, - "data_seed_base": data_seed_base, - "cases": cases, - } - ) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Benchmark OptiMask on large arrays.") - parser.add_argument("--ratio", type=float, default=0.02) - parser.add_argument("--repeats", type=int, default=5) - parser.add_argument("--solve-seed", type=int, default=99) - parser.add_argument("--data-seed-base", type=int, default=1000) - parser.add_argument( - "--shape", - action="append", - default=[], - help="Shape formatted as rows,cols. Can be repeated. Defaults to test_speed shapes.", - ) - parser.add_argument("--record", action="store_true", help="Record the run in branch_registry.json.") - return parser.parse_args() - - -def main() -> None: - args = parse_args() - shapes = args.shape or ["100000,1000", "1000,100000"] - - cases: list[dict[str, object]] = [] - for index, item in enumerate(shapes): - rows_str, cols_str = item.split(",") - case = benchmark_case( - shape=(int(rows_str), int(cols_str)), - ratio=args.ratio, - data_seed=args.data_seed_base + index, - solve_seed=args.solve_seed, - repeats=args.repeats, - ) - cases.append(case) - - branch = get_branch_name() - payload = {"branch": branch, "cases": cases} - print(json.dumps(payload, indent=2)) - - if args.record: - registry = load_registry() - record_run( - registry=registry, - branch=branch, - cases=cases, - repeats=args.repeats, - ratio=args.ratio, - solve_seed=args.solve_seed, - data_seed_base=args.data_seed_base, - ) - write_registry(registry) - - -if __name__ == "__main__": - main() diff --git a/dev/perf/branch_registry.json b/dev/perf/branch_registry.json deleted file mode 100644 index 5315db6..0000000 --- a/dev/perf/branch_registry.json +++ /dev/null @@ -1,198 +0,0 @@ -{ - "best_branch": "perf/trial-05-preprocess-direct-cols", - "trial_summary": { - "perf/tooling-base": { - "status": "baseline", - "vertical_mean_ms": 188.4885599836707, - "horizontal_mean_ms": 227.33483996707946 - }, - "perf/trial-01-numpy-preprocess": { - "status": "regression", - "vertical_mean_ms": 641.0418800078332, - "horizontal_mean_ms": 647.3005800042301 - }, - "perf/trial-02-preprocess-two-pass": { - "status": "regression", - "vertical_mean_ms": 238.73366001062095, - "horizontal_mean_ms": 243.5193199897185 - }, - "perf/trial-03-no-parallel-kernels": { - "status": "regression", - "vertical_mean_ms": 218.25817998033017, - "horizontal_mean_ms": 246.4767999947071 - }, - "perf/trial-04-fused-step-kernels": { - "status": "regression", - "vertical_mean_ms": 197.76684001553804, - "horizontal_mean_ms": 245.14067999552935 - }, - "perf/trial-05-preprocess-direct-cols": { - "status": "best_overall", - "vertical_mean_ms": 185.9151199925691, - "horizontal_mean_ms": 202.78805999550968 - }, - "perf/trial-06-fast-solve-dispatch": { - "status": "mixed", - "vertical_mean_ms": 190.4429600108415, - "horizontal_mean_ms": 203.50568001158535 - }, - "perf/trial-07-groupby-manual-compare": { - "status": "regression", - "vertical_mean_ms": 196.5345600154251, - "horizontal_mean_ms": 213.05025999899954 - }, - "perf/trial-08-manual-rectangle-scan": { - "status": "close_second", - "vertical_mean_ms": 190.44409999623895, - "horizontal_mean_ms": 203.20151997730136 - }, - "perf/trial-09-rectangle-plus-fast-dispatch": { - "status": "mixed", - "vertical_mean_ms": 190.7827200135216, - "horizontal_mean_ms": 204.16883998550475 - } - }, - "branches": { - "perf/tooling-base": [ - { - "timestamp_utc": "2026-03-19T23:05:29.345043+00:00", - "repeats": 5, - "ratio": 0.02, - "solve_seed": 99, - "data_seed_base": 1000, - "cases": [ - { - "shape": [ - 100000, - 1000 - ], - "ratio": 0.02, - "timings_ms": [ - 185.20669999998063, - 193.7179999658838, - 189.09100000746548, - 184.8823999753222, - 189.54469996970147 - ], - "mean_ms": 188.4885599836707, - "median_ms": 189.09100000746548, - "min_ms": 184.8823999753222, - "max_ms": 193.7179999658838 - }, - { - "shape": [ - 1000, - 100000 - ], - "ratio": 0.02, - "timings_ms": [ - 221.3730999501422, - 225.12890002690256, - 221.29289992153645, - 249.7455000411719, - 219.1337998956442 - ], - "mean_ms": 227.33483996707946, - "median_ms": 221.3730999501422, - "min_ms": 219.1337998956442, - "max_ms": 249.7455000411719 - } - ] - } - ], - "perf/trial-08-manual-rectangle-scan": [ - { - "timestamp_utc": "2026-03-19T23:17:04.441246+00:00", - "repeats": 5, - "ratio": 0.02, - "solve_seed": 99, - "data_seed_base": 1000, - "cases": [ - { - "shape": [ - 100000, - 1000 - ], - "ratio": 0.02, - "timings_ms": [ - 184.71750000026077, - 203.57519993558526, - 182.02970002312213, - 196.70490000862628, - 185.1932000136003 - ], - "mean_ms": 190.44409999623895, - "median_ms": 185.1932000136003, - "min_ms": 182.02970002312213, - "max_ms": 203.57519993558526 - }, - { - "shape": [ - 1000, - 100000 - ], - "ratio": 0.02, - "timings_ms": [ - 194.40169993322343, - 217.1362000517547, - 202.55619997624308, - 193.83589993230999, - 208.0775999929756 - ], - "mean_ms": 203.20151997730136, - "median_ms": 202.55619997624308, - "min_ms": 193.83589993230999, - "max_ms": 217.1362000517547 - } - ] - } - ], - "perf/trial-05-preprocess-direct-cols": [ - { - "timestamp_utc": "2026-03-19T23:18:32.653025+00:00", - "repeats": 5, - "ratio": 0.02, - "solve_seed": 99, - "data_seed_base": 1000, - "cases": [ - { - "shape": [ - 100000, - 1000 - ], - "ratio": 0.02, - "timings_ms": [ - 183.97919996641576, - 192.23649997729808, - 180.98800000734627, - 180.4888000478968, - 191.88309996388853 - ], - "mean_ms": 185.9151199925691, - "median_ms": 183.97919996641576, - "min_ms": 180.4888000478968, - "max_ms": 192.23649997729808 - }, - { - "shape": [ - 1000, - 100000 - ], - "ratio": 0.02, - "timings_ms": [ - 198.4950000187382, - 197.51129997894168, - 196.86200004070997, - 221.5913999825716, - 199.48059995658696 - ], - "mean_ms": 202.78805999550968, - "median_ms": 198.4950000187382, - "min_ms": 196.86200004070997, - "max_ms": 221.5913999825716 - } - ] - } - ] - } -} diff --git a/dev/perf/profile_optimask.py b/dev/perf/profile_optimask.py deleted file mode 100644 index bea09c9..0000000 --- a/dev/perf/profile_optimask.py +++ /dev/null @@ -1,43 +0,0 @@ -from __future__ import annotations - -import argparse -import cProfile -import pstats -from pathlib import Path - -from benchmark_optimask import generate_random -from optimask import OptiMask - - -ROOT = Path(__file__).resolve().parents[2] - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description="Profile a single OptiMask solve call.") - parser.add_argument("--rows", type=int, default=100000) - parser.add_argument("--cols", type=int, default=1000) - parser.add_argument("--ratio", type=float, default=0.02) - parser.add_argument("--data-seed", type=int, default=1000) - parser.add_argument("--solve-seed", type=int, default=99) - parser.add_argument("--sort", default="cumtime", choices=["cumtime", "tottime", "ncalls"]) - parser.add_argument("--limit", type=int, default=25) - return parser.parse_args() - - -def main() -> None: - args = parse_args() - x = generate_random(args.rows, args.cols, args.ratio, seed=args.data_seed) - solver = OptiMask(random_state=args.solve_seed) - solver.solve(x) - - profiler = cProfile.Profile() - profiler.enable() - solver.solve(x) - profiler.disable() - - stats = pstats.Stats(profiler).strip_dirs().sort_stats(args.sort) - stats.print_stats(args.limit) - - -if __name__ == "__main__": - main() From d90aa0fb295d9086e255c361b041697465cf2646 Mon Sep 17 00:00:00 2001 From: CyrilJl Date: Fri, 20 Mar 2026 18:22:25 +0100 Subject: [PATCH 6/6] Print result sizes in speed test --- tests/test_optimask.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_optimask.py b/tests/test_optimask.py index 0a3fa6b..cb66bde 100644 --- a/tests/test_optimask.py +++ b/tests/test_optimask.py @@ -126,14 +126,14 @@ def test_speed(opti_mask_instance): print("\nVertical arrays") for _ in range(5): start = perf_counter() - opti_mask_instance.solve(X=x) - print(f"{1e3 * (perf_counter() - start):.2f}ms") + rows, cols = opti_mask_instance.solve(X=x) + print(f"{1e3 * (perf_counter() - start):.2f}ms rows={rows.size} cols={cols.size}") x = generate_random(m=1_000, n=100_000, ratio=0.02) print("Horizontal arrays") for _ in range(5): start = perf_counter() - opti_mask_instance.solve(X=x) - print(f"{1e3 * (perf_counter() - start):.2f}ms") + rows, cols = opti_mask_instance.solve(X=x) + print(f"{1e3 * (perf_counter() - start):.2f}ms rows={rows.size} cols={cols.size}") def test_large_arrays(opti_mask_instance):