regular_isd_rgbp_points.py

"""Compute regular-GBP parameter scans for regular ISD algorithms in mpf precision."""

from __future__ import annotations

import argparse
from datetime import datetime
from pathlib import Path
import time

import mpmath as mp

from mp_optimizer import optimize_problem, strategies_for_algorithm
from mp_regular_isd import REGULAR_ALGORITHMS
from mp_regular_sampler import build_regular_sampler_config
from mp_transform import regular_gbp
from mp_utils import Float, dump_precise_json, format_fixed, is_finite_number, iter_float_range
from mp_validator import array_to_input_params, validate_point


def _point_format_digits(c):
    """Return per-field decimal digits for one progress line."""
    if Float(c) == Float("0.01"):
        return {"omega": 32, "kappa": 32, "fun": 16}
    return {"omega": 16, "kappa": 16, "fun": 16}


def _output_path(algorithm, seed, output):
    if output:
        return Path(output)
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    return Path(__file__).resolve().parent / "json" / f"regular_{algorithm}_rgbp_{seed}_{timestamp}.json"


def _run_direct(spec, kappa, omega):
    value = Float(spec["objective"](Float(kappa), Float(omega)))
    return {
        "fun": value,
        "input_params": None,
        "success": bool(is_finite_number(value)),
        "method": "direct",
        "tol": Float("0"),
        "timed_out": False,
    }


def compute_regular_rgbp_points(start, end, step, *, algorithm, seed=0, max_single_opt_time=None):
    """Compute one regular-GBP parameter scan for a regular algorithm."""
    del max_single_opt_time
    entries = []
    for idx, c in iter_float_range(start, end, step):
        point_started_at = time.perf_counter()
        omega, kappa = regular_gbp(c)
        entry = {"c": Float(c), "omega": Float(omega), "kappa": Float(kappa), "algorithm": algorithm}

        if omega != omega or kappa != kappa:
            entry["result"] = {
                "fun": mp.inf,
                "input_params": None,
                "feasible": False,
                "success": False,
                "method": "N/A",
                "tol": Float("nan"),
            }
            entries.append(entry)
            continue

        config = build_regular_sampler_config(algorithm, kappa, omega)
        if config["direct"]:
            result = _run_direct(config, kappa, omega)
            validation = {
                "feasible": bool(result["success"]),
                "min_bound_margin": mp.inf,
                "min_ineq": mp.inf,
                "max_eq_abs": Float("0"),
            }
        else:
            result = optimize_problem(
                objective=config["objective_vector"],
                constraints=config["wrapped_constraints"],
                bounds=config["bounds"],
                sample_func=config["sample_func"],
                strategies=strategies_for_algorithm("regular", algorithm),
                seed=int(seed) + idx,
            )
            validation = validate_point(
                kappa=kappa,
                omega=omega,
                input_params=result.get("input_params"),
                var_names=config["free_var_names"],
                vars_type=config["vars_type"],
                objective=config["objective_vector"],
                constraints=config["constraints"],
                bounds=config["bounds"],
                stored_fun=result.get("fun"),
                tol=result.get("tol"),
                penalty=config["penalty"],
            )

        normalized_fun = Float(result["fun"]) / max(Float("1") - Float(kappa), Float("1e-30"))
        entry["result"] = {
            "fun": normalized_fun,
            "input_params": array_to_input_params(result.get("input_params"), config["free_var_names"]),
            "feasible": bool(validation["feasible"]),
            "success": bool(result.get("success", False)),
            "method": result.get("method", "N/A"),
            "tol": Float(result.get("tol", Float("nan"))),
        }
        entries.append(entry)
        digits = _point_format_digits(c)
        print(
            f"c={format_fixed(c, 2)}, algorithm={algorithm}, omega={format_fixed(omega, digits['omega'])}, "
            f"kappa={format_fixed(kappa, digits['kappa'])}, fun={format_fixed(normalized_fun, digits['fun'])}, "
            f"feasible={entry['result']['feasible']}, elapsed={time.perf_counter() - point_started_at:.3f}s"
        )
    return entries


def main():
    parser = argparse.ArgumentParser(description="Compute regular-GBP parameter scans for regular ISD algorithms.")
    parser.add_argument("--algorithm", default="enum", choices=sorted(REGULAR_ALGORITHMS))
    parser.add_argument("--start", type=str, default="0.01")
    parser.add_argument("--end", type=str, default="1.00")
    parser.add_argument("--step", type=str, default="0.01")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--single-opt-timeout", type=float, default=None, help="Accepted for CLI compatibility; ignored by mpoptimize.")
    parser.add_argument("--output", type=str, default=None)
    args = parser.parse_args()

    records = compute_regular_rgbp_points(
        Float(args.start),
        Float(args.end),
        Float(args.step),
        algorithm=args.algorithm,
        seed=args.seed,
        max_single_opt_time=args.single_opt_timeout,
    )
    output_path = _output_path(args.algorithm, args.seed, args.output)
    dump_precise_json(output_path, records)
    print(f"wrote {output_path}")


if __name__ == "__main__":
    main()