dkx 2.3.0__tar.gz → 2.3.1__tar.gz

{dkx-2.3.0 → dkx-2.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dkx
-Version: 2.3.0
+Version: 2.3.1
 Summary: Neoclassical transport solver with CPU/GPU and differentiable JAX workflows
 Project-URL: Homepage, https://github.com/uwplasma/DKX
 Project-URL: Documentation, https://sfincs-jax.readthedocs.io/

{dkx-2.3.0 → dkx-2.3.1}/dkx/__init__.py

@@ -393,4 +393,4 @@ __all__ = [
     "write_output",
 ]
 
-__version__ = "2.3.0"
+__version__ = "2.3.1"

{dkx-2.3.0 → dkx-2.3.1}/dkx/solve.py

@@ -976,6 +976,144 @@ def materialize_dense(
         cols.append(np.asarray(batched(basis)))
     return np.concatenate(cols, axis=1)
 
+
+def _escalate_after_tier2_stall(
+    op: KineticOperator,
+    rhs2d: jnp.ndarray,
+    *,
+    stalled: SolveResult,
+    tol: float,
+    atol: float,
+    x0,
+    recycle,
+    preconditioner: str,
+    drop_l_coupling_in_precond: bool,
+    restart: int,
+    recycle_dim: int,
+    max_restarts: int,
+    check_adjoint: bool,
+    adjoint_residual_factor: float,
+    max_dense_size: int,
+) -> SolveResult:
+    """Work through the remedies a user would otherwise have to know about.
+
+    A stalled Krylov solve is a *preconditioner* problem, not a reason to
+    change solver tier, and this is where DKX used to get it backwards: it
+    announced a fall back to the tier-3 host direct solve, which obtains its
+    matrix by applying the operator to ``n`` identity columns.  At the sizes
+    where tier 2 actually stalls that is hopeless -- a 66004-DOF deck needs
+    66004 matvecs and O(n^2) memory -- so the guard in :func:`_solve_tier3`
+    refused, and a convergence problem surfaced as a hard crash telling the
+    user to raise ``max_dense_size``.  Taking that advice would have tried to
+    allocate 32.5 GB.
+
+    SFINCS Fortran v3 does not have this failure mode, and the reason is worth
+    stating because it dictates the order below.  It assembles the simplified
+    preconditioner matrix *analytically and sparsely*, factorizes it with a
+    sparse direct LU (MUMPS or SuperLU_dist), and preconditions GMRES with
+    that -- so its "direct solve" is a sparse factorization that handles 66004
+    routinely, and GMRES needs few iterations on top of it.  It also retries
+    automatically, doubling the MUMPS working-memory factor on a failed
+    factorization (``solver.F90``: ``mumps_icntl_14 = mumps_icntl_14 * 2``).
+
+    DKX already owns the equivalent of that preconditioner -- ``"sparse"``
+    eliminates in a fill-reducing order and keeps the inverse exact -- and the
+    stall above happened without ever trying it.  So the ladder escalates the
+    preconditioner first, in increasing cost, and only considers tier 3 at a
+    size where tier 3 can actually run.
+
+    Returns the first converged result, or the best (lowest final residual) if
+    none converges, so a caller that can tolerate a loose solve still gets the
+    best available answer along with an honest ``converged=False``.
+    """
+    attempts: list[tuple[str, SolveResult]] = [(f"{preconditioner} preconditioner", stalled)]
+
+    def _residual(result: SolveResult) -> float:
+        norms = np.asarray(result.residual_norms)
+        return float(norms[-1]) if norms.size else float("inf")
+
+    print(
+        f"[dkx.solve] tier-2 Krylov stalled with the {preconditioner} "
+        f"preconditioner (iterations={stalled.iterations}, residual="
+        f"{_residual(stalled):.3e}); escalating rather than giving up."
+    )
+
+    # Rung 1: the strong preconditioner.  This is the SFINCS analogue and the
+    # single most likely fix, because "coarse" eliminates L first, which fills
+    # the angular stencils in, while "sparse" eliminates in a fill-reducing
+    # order and stays sparse.
+    for kind in ("sparse", "multigrid"):
+        if kind == preconditioner:
+            continue
+        try:
+            print(f"[dkx.solve]   retrying tier 2 with the {kind} preconditioner ...")
+            candidate = _solve_tier2(
+                op, rhs2d, tol=tol, atol=atol, x0=x0, recycle=recycle,
+                preconditioner=kind,
+                drop_l_coupling_in_precond=drop_l_coupling_in_precond,
+                restart=restart, recycle_dim=recycle_dim,
+                max_restarts=max_restarts, differentiable=False,
+                check_adjoint=check_adjoint,
+                adjoint_residual_factor=adjoint_residual_factor,
+            )
+        except Exception as exc:  # a preconditioner that cannot build is a rung, not a failure
+            print(f"[dkx.solve]   {kind} preconditioner unavailable: {exc}")
+            continue
+        attempts.append((f"{kind} preconditioner", candidate))
+        if candidate.converged:
+            print(f"[dkx.solve]   converged with the {kind} preconditioner.")
+            return candidate
+
+    # Rung 2: more iterations.  Cheap to ask for, and the remedy when the
+    # preconditioner is adequate but the cap was simply too low for this Er.
+    widened = max(max_restarts * 4, max_restarts + 1)
+    best_kind = max(attempts, key=lambda item: -_residual(item[1]))[0].split()[0]
+    if best_kind not in _TIER2_PRECONDITIONERS:
+        best_kind = preconditioner
+    try:
+        print(
+            f"[dkx.solve]   retrying tier 2 with the {best_kind} preconditioner "
+            f"and {widened} restarts (was {max_restarts}) ..."
+        )
+        candidate = _solve_tier2(
+            op, rhs2d, tol=tol, atol=atol, x0=x0, recycle=recycle,
+            preconditioner=best_kind,
+            drop_l_coupling_in_precond=drop_l_coupling_in_precond,
+            restart=restart, recycle_dim=recycle_dim,
+            max_restarts=widened, differentiable=False,
+            check_adjoint=check_adjoint,
+            adjoint_residual_factor=adjoint_residual_factor,
+        )
+        attempts.append((f"{best_kind} preconditioner, {widened} restarts", candidate))
+        if candidate.converged:
+            print("[dkx.solve]   converged with the widened iteration budget.")
+            return candidate
+    except Exception as exc:
+        print(f"[dkx.solve]   widened retry failed: {exc}")
+
+    # Rung 3: tier 3, but only where it can actually run.  Announcing a
+    # fallback that the size guard will refuse is what produced the original
+    # crash, so the size is checked here rather than discovered downstream.
+    if op.total_size <= max_dense_size:
+        print("[dkx.solve]   falling back to the tier-3 host direct solve ...")
+        return _solve_tier3(op, rhs2d, tol=tol, atol=atol, max_dense_size=max_dense_size)
+
+    best_label, best = min(attempts, key=lambda item: _residual(item[1]))
+    dense_gb = op.total_size**2 * 8 / 1024**3
+    raise RuntimeError(
+        f"the linear solve did not converge at total_size={op.total_size}.\n"
+        f"Tried: {'; '.join(label for label, _ in attempts)}.\n"
+        f"Best final residual {_residual(best):.3e} ({best_label}), tolerance {tol:.1e}.\n"
+        "Tier 3 was not attempted: it materializes the operator column by "
+        f"column, which at this size needs {op.total_size} matvecs and "
+        f"{dense_gb:.1f} GB, so raising max_dense_size would not help.\n"
+        "What usually does help, in order: lower the collisionality-scaled Er "
+        "(stalls cluster at the largest |Er|, where the ExB term dominates); "
+        "raise solverTolerance; or reduce Nxi/Ntheta/Nzeta at the offending "
+        "point.  Pass solver=SolverOptions(...) to control this directly."
+    )
+
+
 def _solve_tier3(
     op: KineticOperator, rhs2d: jnp.ndarray, *, tol: float, atol: float, max_dense_size: int
 ) -> SolveResult:
@@ -2085,13 +2223,23 @@ def solve(
             adjoint_residual_factor=adjoint_residual_factor,
         )
         if method == "auto" and not result.converged and not differentiable:
-            print(
-                "[dkx.solve] tier-2 Krylov breached its iteration cap "
-                f"(iterations={result.iterations}); falling back to the tier-3 "
-                "host direct solve."
-            )
-            result = _solve_tier3(
-                op, rhs2d, tol=tol, atol=atol, max_dense_size=max_dense_size
+            requested = _resolve_preconditioner(preconditioner, use_preconditioner)
+            result = _escalate_after_tier2_stall(
+                op,
+                rhs2d,
+                stalled=result,
+                tol=tol,
+                atol=atol,
+                x0=x0,
+                recycle=recycle,
+                preconditioner=requested,
+                drop_l_coupling_in_precond=drop_l_coupling_in_precond,
+                restart=restart,
+                recycle_dim=recycle_dim,
+                max_restarts=max_restarts,
+                check_adjoint=check_adjoint,
+                adjoint_residual_factor=adjoint_residual_factor,
+                max_dense_size=max_dense_size,
             )
     else:  # direct
         if differentiable:

{dkx-2.3.0 → dkx-2.3.1}/dkx.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dkx
-Version: 2.3.0
+Version: 2.3.1
 Summary: Neoclassical transport solver with CPU/GPU and differentiable JAX workflows
 Project-URL: Homepage, https://github.com/uwplasma/DKX
 Project-URL: Documentation, https://sfincs-jax.readthedocs.io/

{dkx-2.3.0 → dkx-2.3.1}/dkx.egg-info/SOURCES.txt

@@ -188,6 +188,7 @@ tests/test_run_function_is_not_shadowed_by_run_module.py
 tests/test_run_rhsmode1.py
 tests/test_run_transport.py
 tests/test_runtime_helper_coverage.py
+tests/test_scan_survives_a_failing_point.py
 tests/test_scans_progress_and_recycle.py
 tests/test_sensitivity.py
 tests/test_shaing_callen.py
@@ -200,6 +201,7 @@ tests/test_sparse_precond.py
 tests/test_species_constants.py
 tests/test_tier1_refinement.py
 tests/test_tier2_memory_guard.py
+tests/test_tier2_stall_escalates_instead_of_crashing.py
 tests/test_transport_matrix_write_output_end_to_end.py
 tests/test_units.py
 tests/test_upstream_scanplot2_smoke.py

{dkx-2.3.0 → dkx-2.3.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dkx"
-version = "2.3.0"
+version = "2.3.1"
 description = "Neoclassical transport solver with CPU/GPU and differentiable JAX workflows"
 readme = "README.md"
 requires-python = ">=3.10"

{dkx-2.3.0 → dkx-2.3.1}/tests/test_examples_workflow_browser.py

@@ -30,8 +30,13 @@ def test_workflow_browser_filters_by_physics_topic() -> None:
     vmec_geometry = module._matching_workflows(catalog, topic="vmec", search="geometry")
     assert {workflow["id"] for workflow in vmec_geometry} >= {"vmec_wout_path", "vmec_boozer_jax_pipeline"}
 
-    gpu = module._matching_workflows(catalog, topic="gpu", search="")
-    assert {workflow["id"] for workflow in gpu} >= {"output_format_benchmark"}
+    # There is no "gpu" topic any more: its only workflow was the
+    # output-format benchmark, which moved to tools/performance/ along with the
+    # rest of the maintainer tooling.  The catalog indexes examples, and that
+    # script stopped being one.
+    transport = module._matching_workflows(catalog, topic="transport", search="")
+    assert {workflow["id"] for workflow in transport} >= {"transport_matrix"}
+    assert module._matching_workflows(catalog, topic="gpu", search="") == []
 
 
 def test_workflow_browser_json_cli_is_machine_readable() -> None:

dkx-2.3.1/tests/test_scan_survives_a_failing_point.py

@@ -0,0 +1,108 @@
+"""One unsolvable scan point must not discard the rest of the scan.
+
+Reported from a scanType=5 run: tier-2 stalled at the largest |Er| on a
+66004-DOF deck, the exception propagated out of ``run_dkx``, and every
+remaining Er point at that radius was lost.  One radius folder finished with
+zero of a hundred outputs, another with three.  The loss was silent, and it
+also stranded the process inside the failed run's directory, because the
+scan loops ``chdir`` back only *after* the solve returns.
+
+A point that will not solve is information, not a reason to throw away the
+other ninety-nine.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import pytest
+
+UTILS = Path(__file__).resolve().parents[1] / "examples" / "sfincs_examples" / "utils"
+SCAN_SCRIPTS = ("sfincsScan_1", "sfincsScan_2", "sfincsScan_3",
+                "sfincsScan_4", "sfincsScan_21", "sfincsScan_22")
+
+
+def _load_common(tmp_path, monkeypatch):
+    """Exec sfincsScan_common the way sfincsScan does, in its own namespace."""
+    monkeypatch.syspath_prepend(str(UTILS))
+    namespace: dict = {"__file__": str(UTILS / "sfincsScan_common")}
+    exec((UTILS / "sfincsScan_common").read_text(), namespace)  # noqa: S102
+    return namespace
+
+
+def test_a_failing_point_is_recorded_and_the_scan_continues(tmp_path, monkeypatch):
+    ns = _load_common(tmp_path, monkeypatch)
+    run_scan_point = ns["run_scan_point"]
+
+    # Stand in for dkx_driver: fail only the third point, as a stalled solve
+    # would, and succeed on the rest.
+    calls = []
+
+    class _FakeDriver:
+        @staticmethod
+        def run_dkx(*, input_namelist, output_path, **kwargs):
+            calls.append(Path(output_path).parent.name)
+            if Path(output_path).parent.name == "Er3":
+                raise RuntimeError("the linear solve did not converge at total_size=66004")
+            Path(output_path).write_text("ok")
+
+    monkeypatch.setitem(sys.modules, "dkx_driver", _FakeDriver)
+
+    ok = []
+    for i in range(5):
+        directory = tmp_path / f"Er{i + 1}"
+        directory.mkdir()
+        (directory / "input.namelist").write_text("&general\n/\n")
+        ok.append(
+            run_scan_point(
+                input_namelist=directory / "input.namelist",
+                output_path=directory / "sfincsOutput.h5",
+                directory=directory.name,
+            )
+        )
+
+    # Every point was attempted, not just the ones before the failure.
+    assert len(calls) == 5
+    assert ok == [True, True, False, True, True]
+
+    # The four that solved have output; the one that did not is marked.
+    assert (tmp_path / "Er4" / "sfincsOutput.h5").is_file()
+    assert not (tmp_path / "Er3" / "sfincsOutput.h5").exists()
+    marker = tmp_path / "Er3" / ns["FAILURE_MARKER"]
+    assert marker.is_file()
+    assert "did not converge" in marker.read_text()
+
+    assert ns["report_scan_failures"]() == 1
+
+
+def test_keyboard_interrupt_still_stops_the_scan(tmp_path, monkeypatch):
+    """Catching solver failures must not swallow the user pressing Ctrl-C."""
+    ns = _load_common(tmp_path, monkeypatch)
+
+    class _Interrupting:
+        @staticmethod
+        def run_dkx(**kwargs):
+            raise KeyboardInterrupt
+
+    monkeypatch.setitem(sys.modules, "dkx_driver", _Interrupting)
+    directory = tmp_path / "Er1"
+    directory.mkdir()
+    with pytest.raises(KeyboardInterrupt):
+        ns["run_scan_point"](
+            input_namelist=directory / "input.namelist",
+            output_path=directory / "sfincsOutput.h5",
+            directory="Er1",
+        )
+
+
+@pytest.mark.parametrize("script", SCAN_SCRIPTS)
+def test_every_scan_type_uses_the_resilient_helper(script: str) -> None:
+    """A raw run_dkx call in a scan loop is the bug returning."""
+    text = (UTILS / script).read_text()
+    assert "run_scan_point(" in text, f"{script} does not use the resilient helper"
+    body = "\n".join(
+        line for line in text.splitlines() if not line.lstrip().startswith("#")
+    )
+    assert "    run_dkx(" not in body and "   run_dkx(" not in body, (
+        f"{script} still calls run_dkx directly, so one bad point aborts the scan"
+    )

{dkx-2.3.0 → dkx-2.3.1}/tests/test_solve.py

@@ -514,16 +514,23 @@ def test_tier1_refuses_er_xdot_l2_coupling() -> None:
     assert not ok
 
 
-def test_auto_policy_tier3_fallback_on_iteration_cap() -> None:
+def test_auto_policy_recovers_a_starved_tier2_solve() -> None:
     # Starve tier 2 (no preconditioner, tiny restart budget) on the FP
-    # fixture: the auto policy must breach the cap loudly and land on the
-    # tier-3 host direct solve, which still returns the right answer.
+    # fixture: the auto policy must recover and still return the right answer.
+    #
+    # This used to assert result.method == "direct", because a cap breach fell
+    # straight to the tier-3 host solve.  That fallback is a dead end at any
+    # real size -- tier 3 materializes the operator column by column, so a
+    # 66004-DOF deck would need 66004 matvecs -- and a user hit exactly that,
+    # losing a whole radius of an Er scan to the crash.  A stalled Krylov solve
+    # is a preconditioner problem, so the policy now escalates the
+    # preconditioner first and only reaches tier 3 where tier 3 can run.  What
+    # matters is that a starved solve still lands on the correct answer.
     op = _load_op("quick_2species_FPCollisions_noEr")
     rhs = op.rhs()
     result = solve(
         op, rhs, method="auto", tol=1e-10, use_preconditioner=False, max_restarts=2
     )
-    assert result.method == "direct"
     assert result.converged
     x_ref = _dense_solve(op, np.asarray(rhs)[:, None])[:, 0]
     assert _rel_err(np.asarray(result.x), x_ref) < 1e-8

dkx-2.3.1/tests/test_tier2_stall_escalates_instead_of_crashing.py

@@ -0,0 +1,101 @@
+"""A stalled Krylov solve must escalate, not crash on a fallback that cannot run.
+
+Reported from a scanType=5 run (Er scan inside a radius scan) on a 66004-DOF
+deck.  Tier-2 GCROT hit its iteration cap at the largest |Er|, DKX announced a
+fall back to the tier-3 host direct solve, and tier 3 refused immediately
+because 66004 > max_dense_size=8192.  The RuntimeError propagated out through
+the scan driver and killed every remaining Er point at that radius: one radius
+folder finished with zero outputs, another with three of a hundred.
+
+Tier 3 obtains its matrix by applying the operator to n identity columns, so
+at the sizes where tier 2 actually stalls it can never run -- the advice in the
+old message, to raise max_dense_size, would have asked for 32.5 GB.  A stalled
+Krylov solve is a preconditioner problem, and DKX already ships the strong
+preconditioner (``sparse``) that SFINCS's MUMPS LU is the analogue of; it was
+simply never tried.
+"""
+
+import numpy as np
+import pytest
+
+import dkx
+from dkx.solve import SolveResult, _escalate_after_tier2_stall
+
+CASE = dict(
+    geometryScheme=1, inputRadialCoordinate=3, rN_wish=0.3,
+    B0OverBBar=1.0, epsilon_t=-0.07, epsilon_h=0.0, iota=0.4542,
+    GHat=3.7481, IHat=0.0, psiAHat=0.15596, aHat=0.5585,
+    Zs=[1.0], mHats=[1.0], nHats=[1.0], THats=[1.0],
+    dNHatdrHats=[-0.5], dTHatdrHats=[-1.0],
+    Ntheta=9, Nzeta=1, Nxi=8, NL=4, Nx=4,
+    collisionOperator=1, Delta=4.5694e-3, alpha=1.0, nu_n=0.01,
+)
+
+
+def _operator():
+    return dkx.run(**CASE).operator
+
+
+def _stalled(op) -> SolveResult:
+    """What tier 2 hands back when it breaches its cap."""
+    return SolveResult(
+        x=np.zeros((op.total_size, 1)),
+        method="gmres",
+        iterations=6000,
+        residual_norms=np.asarray([1.0, 0.5, 0.2]),
+        converged=False,
+        recycle=None,
+        timings={},
+        adjoint=None,
+    )
+
+
+def _rhs2d(op):
+    """Tier 2 works in (n, nrhs); op.rhs() is (n,), as solve() reshapes it."""
+    return np.asarray(op.rhs()).reshape(op.total_size, -1)
+
+
+def _escalate(op, **overrides):
+    kwargs = dict(
+        stalled=_stalled(op), tol=1e-10, atol=0.0, x0=None, recycle=None,
+        preconditioner="coarse", drop_l_coupling_in_precond=False,
+        restart=20, recycle_dim=8, max_restarts=200,
+        check_adjoint=False, adjoint_residual_factor=1e3,
+        max_dense_size=8192,
+    )
+    kwargs.update(overrides)
+    return _escalate_after_tier2_stall(op, _rhs2d(op), **kwargs)
+
+
+def test_escalation_recovers_a_stalled_solve() -> None:
+    """The ladder must actually produce a converged answer, not just try."""
+    op = _operator()
+    result = _escalate(op)
+    assert result.converged, "escalation should recover a solvable system"
+    assert float(np.asarray(result.residual_norms)[-1]) < 1e-8
+
+
+def test_a_deck_too_large_for_tier3_reports_the_real_problem() -> None:
+    """The failure the user hit: n > max_dense_size, so tier 3 cannot help.
+
+    The message must not send them to max_dense_size, and must name what was
+    tried and what actually helps.
+    """
+    op = _operator()
+    with pytest.raises(RuntimeError) as excinfo:
+        # Force every rung to fail by demanding a tolerance nothing can meet,
+        # with tier 3 excluded exactly as it is at 66004 DOFs.
+        _escalate(op, tol=1e-300, max_dense_size=1)
+    message = str(excinfo.value)
+    assert "did not converge" in message
+    assert "Tried:" in message
+    assert "raising max_dense_size would not help" in message
+    assert "Er" in message, "the remedy that matters should name the Er dependence"
+
+
+def test_the_old_misleading_advice_is_gone() -> None:
+    """`raise max_dense_size explicitly` at 66004 DOFs asks for 32.5 GB."""
+    op = _operator()
+    with pytest.raises(RuntimeError) as excinfo:
+        _escalate(op, tol=1e-300, max_dense_size=1)
+    assert "raise max_dense_size explicitly if you really want this" not in str(excinfo.value)