sciencehub-team/constellarationPublic

def load_boundary(source) -> surface_rz_fourier.SurfaceRZFourier:
    """A plasma boundary from a JSON string, a dict, or a path to a JSON file.

    The format is the one the ConStellaration dataset and leaderboard use:
    `r_cos` and `z_sin` as (m_max + 1) × (2 n_max + 1) arrays, `n_field_periods`,
    and `is_stellarator_symmetric: true`.
    """
    if isinstance(source, surface_rz_fourier.SurfaceRZFourier):
        return source
    if isinstance(source, dict):
        return surface_rz_fourier.SurfaceRZFourier.model_validate(source)
    if isinstance(source, Path) or (isinstance(source, str) and not source.lstrip().startswith("{")):
        source = Path(source).read_text()
    return surface_rz_fourier.SurfaceRZFourier.model_validate_json(source)

def rotating_ellipse(aspect_ratio: float, elongation: float, rotational_transform: float, n_field_periods: int) -> surface_rz_fourier.SurfaceRZFourier:
    """The rotating-ellipse boundary the paper's baselines start from, with major radius 1."""
    return initial_guess.generate_rotating_ellipse(
        aspect_ratio=aspect_ratio,
        elongation=elongation,
        rotational_transform=rotational_transform,
        n_field_periods=n_field_periods,
    )

def compute_metrics(boundary, problem: str = "simple_to_build", fidelity: str = "high_fidelity", corrected_flux_compression: bool = False) -> dict:
    """Run the forward model with the settings the official evaluator uses for `problem`.

    Returns the metrics as a dict. The geometrical problem skips the Boozer
    transform and the QI residual, as the evaluator does; the other two
    compute everything. VMEC++ raises RuntimeError when the equilibrium does
    not converge.

    `fidelity` selects the VMEC++ preset: "high_fidelity" is what the
    leaderboard uses; "low_fidelity" and "very_low_fidelity" are cheaper
    presets for the inner loop of an optimiser, whose results should be
    checked at high fidelity. With `corrected_flux_compression` the dict also
    holds `flux_compression_pest`, the proxy averaged with the field-aligned
    Jacobian as constellaration 0.3.0 does, from the same equilibrium.
    """
    settings = forward_model.ConstellarationSettings
    chosen = settings.default_high_fidelity_skip_qi() if problem == "geometrical" else settings.default_high_fidelity()
    if fidelity != "high_fidelity":
        chosen = chosen.model_copy(update={"vmec_preset_settings": chosen.vmec_preset_settings.model_copy(update={"fidelity": fidelity})})
    metrics, equilibrium = forward_model.forward_model(load_boundary(boundary), settings=chosen)
    out = metrics.model_dump(mode="json")
    if corrected_flux_compression:
        out["flux_compression_pest"] = flux_compression_both_weights(equilibrium, chosen.turbulent_settings)[1]
    return out

def flux_compression_both_weights(equilibrium, settings) -> tuple[float, float]:
    """Flux compression in bad curvature averaged with the VMEC and with the PEST Jacobian.

    The first value is what constellaration 0.2.6 reports. The second changes
    only the averaging weight, as constellaration 0.3.0 does
    (proximafusion/constellaration#109), on the same equilibrium.
    """
    from constellaration.mhd import turbulent_transport, vmec_utils
    from simsopt import mhd

    vmec = vmec_utils.as_simsopt_vmec(equilibrium)
    alpha, theta = turbulent_transport._set_up_surface_grid(
        settings=settings, n_field_periods=equilibrium.nfp, is_stellarator_symmetric=not equilibrium.lasym
    )
    data = mhd.vmec_fieldlines(vs=vmec, s=settings.normalized_toroidal_flux, alpha=alpha, theta1d=theta)
    grad_r_squared = equilibrium.Aminor_p**2 / (4 * settings.normalized_toroidal_flux[:, None, None]) * data.grad_s_dot_grad_s
    bad_curvature = np.heaviside(data.B_cross_kappa_dot_grad_alpha * np.sign(equilibrium.phipf[-1]), 0)

    def averaged(jacobian):
        weights = np.abs(jacobian) / np.abs(jacobian).mean(axis=(1, 2))[:, None, None]
        return float(np.mean(bad_curvature * grad_r_squared * weights, axis=(1, 2))[0])

    return averaged(data.sqrt_g_vmec), averaged(data.edge_toroidal_flux_over_2pi / data.B_sup_theta_pest)

def constraint_report(problem: str, metrics: dict) -> pd.DataFrame:
    """Each constraint of `problem` with its value, bound and normalised violation.

    The violation is (value − bound) / |bound|, signed so that a positive
    number means the constraint is broken; the evaluator accepts violations up
    to 0.01. For the vacuum well, whose bound is 0, it divides by 0.1. The
    numbers come from constellaration's own code.
    """
    spec = PROBLEMS[problem]
    evaluator = spec["evaluator"]()
    typed = forward_model.ConstellarationMetrics.model_validate({k: v for k, v in metrics.items() if k in forward_model.ConstellarationMetrics.model_fields})
    violations = evaluator._normalized_constraint_violations(typed)
    rows = []
    for (name, relation, bound), violation in zip(spec["constraints"], violations):
        if name.startswith("|"):
            value = abs(metrics[name.strip("|")])
        elif name == "log10(qi)":
            value = float(np.log10(metrics["qi"]))
        else:
            value = metrics[name]
        rows.append({
            "constraint": f"{name} {relation} {bound:g}",
            "value": value,
            "normalised violation": float(violation),
            "satisfied": bool(violation <= TOLERANCE),
        })
    return pd.DataFrame(rows)

def hypervolume(points) -> float:
    """The MHD-stable score of feasible (gradient scale length, aspect ratio) points.

    The area in the plane of −L∇B and A that the points dominate, up to the
    reference point (−L∇B, A) = (1, 20). Pass only feasible points.
    """
    if len(points) == 0:
        return 0.0
    from pymoo.indicators import hv

    X = np.array([(-float(L), float(A)) for L, A in points])
    return float(hv.Hypervolume(ref_point=np.array(HYPERVOLUME_REFERENCE))(X))

def added_hypervolume(candidate, front) -> float:
    """How much a feasible (L∇B, A) point would add to the hypervolume of `front`.

    `front` is a list of feasible (L∇B, A) points, for example the leading
    submission's. Zero when the candidate is dominated.
    """
    return hypervolume([*front, candidate]) - hypervolume(front)

def evaluate(problem: str, boundary, force: bool = False) -> dict:
    """Score a boundary (or, for "mhd_stable", a set of boundaries) as the leaderboard does.

    For "geometrical" and "simple_to_build" the result holds objective,
    feasibility, feasible, score, metrics, constraints (a list of dicts as in
    `constraint_report`) and seconds. For "mhd_stable", `boundary` is a list
    of boundaries or the leaderboard's submission format (a JSON array of
    boundary JSON strings, as text or as a path), and the result holds one
    such record per boundary under "points" and the set's score, the
    hypervolume of its feasible points. A boundary that does not parse, or
    whose evaluation raises, usually because VMEC++ did not converge, gets a
    record with "error" and feasible False, where the leaderboard would reject the whole submission;
    for the single-objective problems the record has score 0. The calls are
    those of constellaration 0.2.6's `evaluate`, with one forward-model run
    per boundary.

    The two QI problems raise MemoryError before any work when the machine
    has less memory than one QI evaluation needs (`memory_shortfall`), because
    running out of memory would end the Python process. `force=True` runs
    them anyway.
    """
    if problem not in PROBLEMS:
        raise ValueError(f"problem must be one of {', '.join(PROBLEMS)}")
    refuse_without_memory(problem, force)
    if problem == "mhd_stable":
        points = [evaluate_single(problem, b) for b in boundary_list(boundary)]
        return {"problem": problem, "evaluator": EVALUATOR_VERSION, "points": points, "score": hypervolume(feasible_points(points))}
    return {"problem": problem, "evaluator": EVALUATOR_VERSION, **evaluate_single(problem, boundary)}

def boundary_list(boundaries) -> list:
    """The boundaries of a list, or of the leaderboard's submission format.

    The submission format is a JSON array of boundary JSON strings, given as
    text or as a path to a file. A single boundary becomes a list of one.
    """
    if isinstance(boundaries, Path) or (isinstance(boundaries, str) and not boundaries.lstrip().startswith("{")):
        text = boundaries if isinstance(boundaries, str) and boundaries.lstrip().startswith("[") else Path(boundaries).read_text()
        boundaries = json.loads(text)
    if isinstance(boundaries, (str, dict, surface_rz_fourier.SurfaceRZFourier)):
        return [boundaries]
    return list(boundaries)

def feasible_points(records) -> list[tuple[float, float]]:
    """The (L∇B, A) points of the feasible records among `records`, as `hypervolume` takes them."""
    return [
        (r["metrics"]["minimum_normalized_magnetic_gradient_scale_length"], r["metrics"]["aspect_ratio"])
        for r in records
        if r["feasible"]
    ]

def error_record(problem: str, message: str, seconds: float) -> dict:
    """The record of a boundary whose evaluation failed: not feasible, and score 0 for the single-objective problems."""
    record = {"error": message, "feasible": False, "feasibility": None, "metrics": None, "constraints": None}
    if problem != "mhd_stable":
        record.update(objective=None, score=0.0)
    record["seconds"] = round(seconds, 1)
    return record

def evaluate_single(problem: str, boundary) -> dict:
    """One boundary's record for `evaluate`: one forward-model run, then the evaluator's own feasibility and score."""
    started = time.perf_counter()
    try:
        metrics = compute_metrics(load_boundary(boundary), problem)
    except Exception as error:
        return error_record(problem, f"{type(error).__name__}: {error}", time.perf_counter() - started)
    evaluator = PROBLEMS[problem]["evaluator"]()
    typed = forward_model.ConstellarationMetrics.model_validate(metrics)
    feasible = evaluator.is_feasible(typed)
    record = {
        "feasibility": evaluator.compute_feasibility(typed),
        "feasible": feasible,
        "metrics": metrics,
        "constraints": constraint_report(problem, metrics).to_dict("records"),
    }
    if problem != "mhd_stable":
        record["objective"] = evaluator.get_objective(typed)[0]
        record["score"] = evaluator._score(typed) if feasible else 0.0
    record["seconds"] = round(time.perf_counter() - started, 1)
    return record

def evaluate_many(problem: str, boundaries, jobs: int = 1, timeout: float = 600.0, force: bool = False) -> list[dict]:
    """Score each boundary in its own Python process, `jobs` at a time.

    `boundaries` is a list or the leaderboard's submission format, as
    `boundary_list` reads it. Returns one record per boundary, in order, as
    `evaluate_single` makes them, except that "seconds" is the wall time of
    the whole process, imports included. A boundary that does not parse, a
    process that runs longer than `timeout` seconds, and a process that
    crashes each give an error record with feasible False; the other
    boundaries still run. The timeout counts wall time, so keep `jobs` at or
    below the number of CPUs. For "mhd_stable",
    `hypervolume(feasible_points(records))` is the set's score.

    The process boundary does not protect the caller from running out of
    memory where one limit covers the caller and its children, as on
    ScienceHub compute: a child that exceeds the kernel's memory ends the
    whole kernel. Use `jobs=1` there. For the same reason the two QI problems
    raise MemoryError, as in `evaluate`, unless `force=True`. Each process
    imports this file again, so the file must exist on disk.
    """
    if problem not in PROBLEMS:
        raise ValueError(f"problem must be one of {', '.join(PROBLEMS)}")
    refuse_without_memory(problem, force)
    source = globals().get("__file__")
    if not source or not Path(source).is_file():
        raise RuntimeError("evaluate_many runs this notebook's file in new processes; import it from evaluate.py on disk.")
    worker = (
        "import importlib.util, json, sys\n"
        "spec = importlib.util.spec_from_file_location('constellaration_evaluate', sys.argv[1])\n"
        "module = importlib.util.module_from_spec(spec)\n"
        "spec.loader.exec_module(module)\n"
        "task = json.loads(open(sys.argv[2]).read())\n"
        "record = module.evaluate_single(task['problem'], task['boundary'])\n"
        "open(sys.argv[3], 'w').write(json.dumps(record))\n"
    )
    texts, records = [], []
    for b in boundary_list(boundaries):
        try:
            texts.append(load_boundary(b).model_dump_json())
            records.append(None)
        except Exception as error:
            texts.append(None)
            records.append(error_record(problem, f"{type(error).__name__}: {error}", 0.0))
    waiting = [i for i in reversed(range(len(texts))) if records[i] is None]
    running = {}
    with tempfile.TemporaryDirectory() as scratch:
        scratch = Path(scratch)

        def finished(i, child, seconds):
            out = scratch / f"{i}.out"
            if child.returncode == 0 and out.is_file():
                return {**json.loads(out.read_text()), "seconds": round(seconds, 1)}
            if child.returncode < 0:
                reason = f"killed by signal {-child.returncode}"
                if child.returncode == -9:
                    reason += ", usually because it ran out of memory"
            else:
                lines = (scratch / f"{i}.err").read_text().strip().splitlines()
                reason = lines[-1] if lines else f"exit status {child.returncode}"
            return error_record(problem, f"ProcessError: {reason}", seconds)

        try:
            while waiting or running:
                while waiting and len(running) < max(1, jobs):
                    i = waiting.pop()
                    (scratch / f"{i}.json").write_text(json.dumps({"problem": problem, "boundary": texts[i]}))
                    with open(scratch / f"{i}.err", "w") as err:
                        command = [sys.executable, "-c", worker, str(source), str(scratch / f"{i}.json"), str(scratch / f"{i}.out")]
                        running[i] = (subprocess.Popen(command, stdout=subprocess.DEVNULL, stderr=err), time.perf_counter())
                time.sleep(0.2)
                for i, (child, started) in list(running.items()):
                    seconds = time.perf_counter() - started
                    if child.poll() is not None:
                        records[i] = finished(i, child, seconds)
                    elif seconds > timeout:
                        child.kill()
                        child.wait()
                        records[i] = error_record(problem, f"TimeoutError: no result after {timeout:g} s", seconds)
                    else:
                        continue
                    del running[i]
        finally:
            for child, _ in running.values():
                child.kill()
                child.wait()
    return records

def leader_front(path=None) -> list[tuple[float, float]]:
    """The feasible (L∇B, A) points of the leading MHD-stable submission.

    Read from `data/leading-front.json` beside this notebook, or from `path`.
    The leaderboard notebook's `build_data.py` writes that file from the
    leaderboard as recorded on 25 October 2026; the points' hypervolume is
    the leader's score. Pass them to `added_hypervolume` as the front.
    """
    if path is None:
        path = Path(globals()["__file__"]).parent / "data" / "leading-front.json"
    return [tuple(point) for point in json.loads(Path(path).read_text())["points"]]

def memory_limit_gib() -> float | None:
    """The memory available to this process in GiB, or None when it cannot tell.

    Reads the container's limit from cgroup v2, then cgroup v1, which reports
    no limit as a number near 2**63. Without a limit there, or without
    cgroup files, it falls back to MemTotal in /proc/meminfo. macOS has
    neither and gives None.
    """
    for name in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"):
        try:
            raw = Path(name).read_text().strip()
        except OSError:
            continue
        if raw != "max" and int(raw) < 2**60:
            return int(raw) / 2**30
        break
    try:
        for line in Path("/proc/meminfo").read_text().splitlines():
            if line.startswith("MemTotal:"):
                return int(line.split()[1]) / 2**20
    except (OSError, ValueError, IndexError):
        pass
    return None

def memory_shortfall(problem: str) -> str | None:
    """Why `problem` cannot run in this machine's memory, or None when it can.

    Only the two QI problems need more than a ScienceHub kernel's 1.5 GiB.
    """
    limit = memory_limit_gib()
    if problem == "geometrical" or limit is None or limit >= QI_MEMORY_GIB:
        return None
    return (
        f"This machine has {limit:.1f} GiB of memory. The QI residual of constellaration 0.2.6 needs "
        f"about {QI_MEMORY_GIB:g} GiB, so only the geometrical problem runs here. Run the other two on "
        "a computer with at least 4 GiB."
    )

def refuse_without_memory(problem: str, force: bool = False) -> None:
    """Raise MemoryError with `memory_shortfall`'s reason, unless `force` is set."""
    reason = None if force else memory_shortfall(problem)
    if reason:
        raise MemoryError(f"{reason} Pass force=True to run it anyway.")

def cross_sections(boundary, toroidal_angles=None, n_theta: int = 256):
    """R and Z of the boundary in planes of constant toroidal angle φ.

    Evaluates R = Σ r_cos[m, n] cos(mθ − n N φ) and Z = Σ z_sin[m, n] sin(mθ − n N φ)
    on the stored coefficients, without solving anything. Returns a list of
    (φ, R, Z) with φ in radians.
    """
    b = load_boundary(boundary)
    nfp = b.n_field_periods
    if toroidal_angles is None:
        toroidal_angles = [0.0, np.pi / (2 * nfp), np.pi / nfp]
    r_cos, z_sin = np.asarray(b.r_cos), np.asarray(b.z_sin)
    m = np.arange(r_cos.shape[0])[:, None]
    n = np.arange(r_cos.shape[1])[None, :] - (r_cos.shape[1] - 1) // 2
    theta = np.linspace(0, 2 * np.pi, n_theta)
    out = []
    for phi in toroidal_angles:
        angle = m[..., None] * theta - (n * nfp * phi)[..., None]
        out.append((phi, (r_cos[..., None] * np.cos(angle)).sum((0, 1)), (z_sin[..., None] * np.sin(angle)).sum((0, 1))))
    return out