drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/vorticity.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+import numpy as np
+from jax.experimental.ode import odeint
+
+from solvax import FourierHelmholtzOperator, build_fourier_helmholtz_operator, solve_fourier_helmholtz
+from .limiters import monotonic_centered_edges_jax as _monotonic_centered_edges
+from .mesh import StructuredMesh, apply_zero_dirichlet_x_guards, communicate_y_guards
+from .metrics import StructuredMetrics
+
+
+VorticityOperator = FourierHelmholtzOperator
+
+
+@dataclass(frozen=True)
+class VorticityRhsResult:
+    vorticity: jnp.ndarray
+    potential: jnp.ndarray
+
+
+@dataclass(frozen=True)
+class VorticityHistoryResult:
+    vorticity_history: jnp.ndarray
+    potential_history: jnp.ndarray
+
+
+def build_vorticity_operator(
+    *,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    average_atomic_mass: float,
+) -> VorticityOperator:
+    _require_supported_vorticity_mesh(mesh, metrics=metrics)
+
+    x_slice = slice(mesh.xstart, mesh.xend + 1)
+    y_index = mesh.ystart
+
+    dx = jnp.asarray(metrics.dx[x_slice, y_index, 0], dtype=jnp.float64)
+    dz = jnp.asarray(metrics.dz[x_slice, y_index, 0], dtype=jnp.float64)
+    g11 = jnp.asarray(metrics.g11[x_slice, y_index, 0], dtype=jnp.float64)
+    g33 = jnp.asarray(metrics.g33[x_slice, y_index, 0], dtype=jnp.float64)
+    rhs_scale = jnp.asarray(metrics.Bxy[x_slice, y_index, 0], dtype=jnp.float64)
+    rhs_scale = (rhs_scale * rhs_scale) / jnp.asarray(average_atomic_mass, dtype=jnp.float64)
+
+    return VorticityOperator(
+        **build_fourier_helmholtz_operator(
+            dx=dx,
+            dz=dz,
+            g11=g11,
+            g33=g33,
+            rhs_scale=rhs_scale,
+            nz=mesh.nz,
+        ).__dict__,
+    )
+
+
+def apply_vorticity_boundaries(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    result = apply_zero_dirichlet_x_guards(field, mesh)
+    if mesh.myg > 0:
+        result = communicate_y_guards(result, mesh)
+    return result
+
+
+def solve_potential(
+    vorticity: jnp.ndarray,
+    *,
+    mesh: StructuredMesh,
+    operator: VorticityOperator,
+) -> jnp.ndarray:
+    guarded = apply_vorticity_boundaries(vorticity, mesh)
+    y_index = mesh.ystart
+    interior = jnp.asarray(guarded[mesh.xstart : mesh.xend + 1, y_index, :], dtype=jnp.float64)
+    interior_phi = solve_fourier_helmholtz(interior, operator=operator)
+
+    potential = jnp.zeros_like(guarded, dtype=jnp.float64)
+    potential = potential.at[mesh.xstart : mesh.xend + 1, y_index, :].set(interior_phi)
+    potential = _apply_potential_boundaries(potential, mesh)
+    if mesh.myg > 0:
+        potential = communicate_y_guards(potential, mesh)
+    return potential
+
+
+def compute_vorticity_rhs(
+    vorticity: jnp.ndarray,
+    *,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    operator: VorticityOperator,
+) -> VorticityRhsResult:
+    guarded = apply_vorticity_boundaries(vorticity, mesh)
+    potential = solve_potential(guarded, mesh=mesh, operator=operator)
+    result = jnp.zeros_like(guarded, dtype=jnp.float64)
+
+    xs, xe = mesh.xstart, mesh.xend
+    interior = slice(xs, xe + 1)
+
+    def roll_km(array: jnp.ndarray) -> jnp.ndarray:
+        return jnp.roll(array, 1, axis=-1)
+
+    def roll_kp(array: jnp.ndarray) -> jnp.ndarray:
+        return jnp.roll(array, -1, axis=-1)
+
+    def shift_from_left(array: jnp.ndarray) -> jnp.ndarray:
+        return jnp.concatenate([jnp.zeros_like(array[:1]), array[:-1]], axis=0)
+
+    def shift_from_right(array: jnp.ndarray) -> jnp.ndarray:
+        return jnp.concatenate([array[1:], jnp.zeros_like(array[:1])], axis=0)
+
+    for j in range(mesh.ystart, mesh.yend + 1):
+        phi = potential[:, j, :]
+        phi_c = phi[interior]
+        phi_m = phi[xs - 1 : xe]
+        phi_p = phi[xs + 1 : xe + 2]
+
+        # Corner-averaged potential at the four cell corners (i∓1/2, k∓1/2).
+        fmm = 0.25 * (phi_c + phi_m + roll_km(phi_c) + roll_km(phi_m))
+        fmp = 0.25 * (phi_c + roll_kp(phi_c) + phi_m + roll_kp(phi_m))
+        fpp = 0.25 * (phi_c + roll_kp(phi_c) + phi_p + roll_kp(phi_p))
+        fpm = 0.25 * (phi_c + phi_p + roll_km(phi_c) + roll_km(phi_p))
+
+        dx_c = jnp.asarray(metrics.dx[interior, j, 0], dtype=jnp.float64)[:, None]
+        dz_c = jnp.asarray(metrics.dz[interior, j, 0], dtype=jnp.float64)[:, None]
+        J_c = jnp.asarray(metrics.J[interior, j, 0], dtype=jnp.float64)[:, None]
+        J_m = jnp.asarray(metrics.J[xs - 1 : xe, j, 0], dtype=jnp.float64)[:, None]
+        J_p = jnp.asarray(metrics.J[xs + 1 : xe + 2, j, 0], dtype=jnp.float64)[:, None]
+
+        v_up = J_c * (fmp - fpp) / dx_c
+        v_down = J_c * (fmm - fpm) / dx_c
+        v_right = (0.5 * (J_c + J_p)) * (fpp - fpm) / dz_c
+        v_left = (0.5 * (J_c + J_m)) * (fmp - fmm) / dz_c
+
+        w_c = guarded[interior, j, :]
+        x_left_face, x_right_face = _mc_cell_edges(w_c, guarded[xs - 1 : xe, j, :], guarded[xs + 1 : xe + 2, j, :])
+        z_left_face, z_right_face = _mc_cell_edges(w_c, roll_km(w_c), roll_kp(w_c))
+
+        # Upwinded face fluxes; x-fluxes vanish at the outer faces (bndry_flux = false).
+        flux_right = jnp.where(v_right > 0.0, v_right * x_right_face, 0.0).at[-1].set(0.0)
+        flux_left = jnp.where(v_left < 0.0, v_left * x_left_face, 0.0).at[0].set(0.0)
+        flux_up = jnp.where(v_up > 0.0, v_up * z_right_face / (J_c * dz_c), 0.0)
+        flux_down = jnp.where(v_down < 0.0, v_down * z_left_face / (J_c * dz_c), 0.0)
+
+        # Flux-conservative scatter pairs written as differences of shifted arrays.
+        dxJ = dx_c * J_c
+        total = (
+            -shift_from_left(flux_right) / dxJ
+            - roll_km(flux_up)
+            + flux_right / dxJ
+            - flux_left / dxJ
+            + flux_up
+            - flux_down
+            + roll_kp(flux_down)
+            + shift_from_right(flux_left) / dxJ
+        )
+        result = result.at[interior, j, :].set(total)
+
+    return VorticityRhsResult(
+        vorticity=-result,
+        potential=potential,
+    )
+
+
+def advance_vorticity_history(
+    initial_vorticity: jnp.ndarray,
+    *,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    operator: VorticityOperator,
+    timestep: float,
+    steps: int,
+    start_time: float = 0.0,
+    rtol: float = 1e-6,
+    atol: float = 1e-8,
+    mxstep: int = 20000,
+) -> VorticityHistoryResult:
+    if steps < 0:
+        raise ValueError("steps must be non-negative")
+
+    y0 = jnp.ravel(jnp.asarray(initial_vorticity[mesh.xstart : mesh.xend + 1, mesh.ystart, :], dtype=jnp.float64))
+    times = jnp.asarray([start_time + timestep * index for index in range(steps + 1)], dtype=jnp.float64)
+
+    def rhs_flat(interior_state: jnp.ndarray, time: jnp.ndarray) -> jnp.ndarray:
+        del time
+        full_state = _assemble_state(interior_state, mesh)
+        rhs = compute_vorticity_rhs(full_state, mesh=mesh, metrics=metrics, operator=operator).vorticity
+        return jnp.ravel(rhs[mesh.xstart : mesh.xend + 1, mesh.ystart, :])
+
+    interior_history = odeint(rhs_flat, y0, times, rtol=rtol, atol=atol, mxstep=mxstep)
+
+    vorticity_history = []
+    potential_history = []
+    for index, interior_state in enumerate(interior_history):
+        full_state = _assemble_state(interior_state, mesh)
+        vorticity_history.append(full_state)
+        if index == 0:
+            potential_history.append(jnp.zeros_like(full_state, dtype=jnp.float64))
+        else:
+            potential_history.append(solve_potential(full_state, mesh=mesh, operator=operator))
+
+    return VorticityHistoryResult(
+        vorticity_history=jnp.stack(vorticity_history, axis=0),
+        potential_history=jnp.stack(potential_history, axis=0),
+    )
+
+
+def _mc_cell_edges(center: jnp.ndarray, minus: jnp.ndarray, plus: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
+    return _monotonic_centered_edges(center, minus, plus)
+
+
+def _apply_potential_boundaries(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    if mesh.mxg != 2:
+        raise NotImplementedError("Native electrostatic potential boundaries currently require MXG = 2.")
+
+    result = jnp.asarray(field, dtype=jnp.float64)
+    rows = slice(mesh.ystart, mesh.yend + 1)
+    result = result.at[mesh.xstart - 1, rows, :].set(-result[mesh.xstart, rows, :])
+    result = result.at[mesh.xend + 1, rows, :].set(-result[mesh.xend, rows, :])
+    result = result.at[mesh.xstart - 2, rows, :].set(result[mesh.xstart - 1, rows, :])
+    result = result.at[mesh.xend + 2, rows, :].set(result[mesh.xend + 1, rows, :])
+    return result
+
+
+def _assemble_state(interior_state: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    interior = jnp.reshape(interior_state, (mesh.xend - mesh.xstart + 1, mesh.nz))
+    full = jnp.zeros((mesh.nx, mesh.local_ny, mesh.nz), dtype=jnp.float64)
+    full = full.at[mesh.xstart : mesh.xend + 1, mesh.ystart, :].set(interior)
+    return apply_vorticity_boundaries(full, mesh)
+
+
+def _require_supported_vorticity_mesh(mesh: StructuredMesh, *, metrics: StructuredMetrics) -> None:
+    if mesh.mxg != 2:
+        raise NotImplementedError("Native electrostatic vorticity currently requires MXG = 2.")
+    if mesh.ny != 1 or mesh.myg != 0:
+        raise NotImplementedError("Native electrostatic vorticity currently requires ny = 1 with MYG = 0.")
+    if not np.allclose(np.asarray(metrics.g23), 0.0, rtol=1e-12, atol=1e-12):
+        raise NotImplementedError("Native electrostatic vorticity currently requires g23 = 0.")
+
+    x_slice = slice(mesh.xstart, mesh.xend + 1)
+    y_index = mesh.ystart
+    for name, array in {
+        "dx": metrics.dx[x_slice, y_index, 0],
+        "dz": metrics.dz[x_slice, y_index, 0],
+        "J": metrics.J[x_slice, y_index, 0],
+        "g11": metrics.g11[x_slice, y_index, 0],
+        "g33": metrics.g33[x_slice, y_index, 0],
+        "Bxy": metrics.Bxy[x_slice, y_index, 0],
+    }.items():
+        values = np.asarray(array, dtype=np.float64)
+        if not np.allclose(values, values[:1], rtol=1e-12, atol=1e-12):
+            raise NotImplementedError(f"Native electrostatic vorticity currently requires uniform {name}.")

drbx/runtime/__init__.py

@@ -0,0 +1,53 @@
+from .artifacts import (
+    ARTIFACT_BASE_URL,
+    ARTIFACT_RELEASE_TAG,
+    DOCS_MEDIA_ASSET,
+    ensure_docs_media,
+)
+from .performance import (
+    configure_jax_runtime,
+    resolve_host_device_count,
+    resolve_runtime_precision,
+    runtime_jax_dtype,
+    runtime_numpy_dtype,
+    runtime_parallel_summary,
+)
+from .memory import PeakRssMeasurement, bytes_to_mebibytes, measure_peak_rss, process_tree_rss_bytes
+from .output import RestartBundle, build_run_log_payload, format_run_log_text, load_restart_bundle, print_run_log, write_restart_bundle, write_run_log_payload
+from .scheduler import ComponentRequest, Scheduler, SupportsSchedulerHooks, expand_component_requests
+from .run_config import MeshScalarConfig, ParallelTransformConfig, RunConfiguration, SolverConfig, TimeConfig
+from .state import SimulationState
+
+__all__ = [
+    "ComponentRequest",
+    "ARTIFACT_BASE_URL",
+    "ARTIFACT_RELEASE_TAG",
+    "RestartBundle",
+    "build_run_log_payload",
+    "configure_jax_runtime",
+    "DOCS_MEDIA_ASSET",
+    "format_run_log_text",
+    "load_restart_bundle",
+    "MeshScalarConfig",
+    "ParallelTransformConfig",
+    "PeakRssMeasurement",
+    "bytes_to_mebibytes",
+    "ensure_docs_media",
+    "print_run_log",
+    "process_tree_rss_bytes",
+    "resolve_host_device_count",
+    "resolve_runtime_precision",
+    "RunConfiguration",
+    "measure_peak_rss",
+    "runtime_jax_dtype",
+    "runtime_numpy_dtype",
+    "runtime_parallel_summary",
+    "Scheduler",
+    "SimulationState",
+    "SolverConfig",
+    "SupportsSchedulerHooks",
+    "TimeConfig",
+    "expand_component_requests",
+    "write_restart_bundle",
+    "write_run_log_payload",
+]

drbx/runtime/artifacts.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import time
+import urllib.request
+import zipfile
+
+from .paths import repo_root
+
+
+ARTIFACT_RELEASE_TAG = "validation-artifacts-2026-04-28"
+ARTIFACT_BASE_URL = f"https://github.com/uwplasma/drbx/releases/download/{ARTIFACT_RELEASE_TAG}"
+DOCS_MEDIA_ASSET = "drbx_docs_media.zip"
+
+DOCS_MEDIA_SENTINELS = (
+    Path(
+        "docs/data/diverted_tokamak_turbulence_artifacts/movies/"
+        "diverted_tokamak_turbulence.gif"
+    ),
+    Path(
+        "docs/data/tokamak_tcv_x21_toroidal_movie_artifacts/movies/"
+        "tokamak_tcv_x21_toroidal.gif"
+    ),
+    Path(
+        "docs/data/stellarator_fci_validation_artifacts/showcase/movies/"
+        "stellarator_sol_showcase.gif"
+    ),
+    Path(
+        "docs/data/essos_imported_drb_movie_stationarity_jacobi_media/movies/"
+        "movie_compact.gif"
+    ),
+)
+
+
+def ensure_docs_media(
+    *,
+    root: str | Path | None = None,
+    base_url: str | None = None,
+    force: bool = False,
+) -> Path:
+    """Restore release-backed documentation figures, movies, and NPZ arrays."""
+
+    resolved_root = Path(root) if root is not None else repo_root()
+    missing = [
+        relative_path
+        for relative_path in DOCS_MEDIA_SENTINELS
+        if not (resolved_root / relative_path).exists()
+    ]
+    if not force and not missing:
+        return resolved_root / "docs" / "data"
+    if os.environ.get("DRBX_OFFLINE_ARTIFACTS", "").lower() in {"1", "true", "yes"}:
+        raise FileNotFoundError(
+            "Docs media are not present and DRBX_OFFLINE_ARTIFACTS is enabled."
+        )
+
+    cache_dir = _artifact_cache_dir(resolved_root)
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    archive_path = cache_dir / DOCS_MEDIA_ASSET
+    if force or not archive_path.exists():
+        resolved_base_url = (
+            base_url or os.environ.get("DRBX_ARTIFACT_BASE_URL") or ARTIFACT_BASE_URL
+        ).rstrip("/")
+        _download_release_asset(
+            f"{resolved_base_url}/{DOCS_MEDIA_ASSET}",
+            archive_path,
+            asset_name=DOCS_MEDIA_ASSET,
+        )
+    with zipfile.ZipFile(archive_path) as archive:
+        archive.extractall(resolved_root)
+    remaining = [
+        str(relative_path)
+        for relative_path in DOCS_MEDIA_SENTINELS
+        if not (resolved_root / relative_path).exists()
+    ]
+    if remaining:
+        raise FileNotFoundError(
+            "Docs media artifact archive did not restore expected files under "
+            f"{resolved_root}: {', '.join(remaining)}"
+        )
+    return resolved_root / "docs" / "data"
+
+
+def _download_release_asset(url: str, destination: Path, *, asset_name: str) -> None:
+    if _download_with_gh(asset_name, destination):
+        return
+    _download_with_urllib(url, destination)
+
+
+def _artifact_cache_dir(root: Path) -> Path:
+    configured = (
+        os.environ.get("DRBX_ARTIFACT_CACHE_DIR")
+        or os.environ.get("DRBX_ARTIFACT_CACHE")
+    )
+    if configured:
+        return Path(configured).expanduser()
+    return root / ".drbx_artifact_cache"
+
+
+def _download_with_gh(asset_name: str, destination: Path) -> bool:
+    gh = shutil.which("gh")
+    if gh is None:
+        return False
+    completed = subprocess.run(
+        [
+            gh,
+            "release",
+            "download",
+            ARTIFACT_RELEASE_TAG,
+            "--repo",
+            "uwplasma/drbx",
+            "--pattern",
+            asset_name,
+            "--dir",
+            str(destination.parent),
+            "--clobber",
+        ],
+        check=False,
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+        text=True,
+    )
+    downloaded = destination.parent / asset_name
+    if completed.returncode == 0 and downloaded.exists():
+        if downloaded != destination:
+            downloaded.replace(destination)
+        return True
+    return False
+
+
+def _download_with_urllib(url: str, destination: Path) -> None:
+    temporary = destination.with_suffix(destination.suffix + ".tmp")
+    request = urllib.request.Request(url)
+    token = os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN")
+    if token:
+        request.add_header("Authorization", f"Bearer {token}")
+    timeout = float(os.environ.get("DRBX_ARTIFACT_DOWNLOAD_TIMEOUT", "120"))
+    attempts = max(1, int(os.environ.get("DRBX_ARTIFACT_DOWNLOAD_ATTEMPTS", "3")))
+    last_error: Exception | None = None
+    for attempt in range(1, attempts + 1):
+        try:
+            with urllib.request.urlopen(request, timeout=timeout) as response, temporary.open(
+                "wb"
+            ) as output:
+                while True:
+                    chunk = response.read(1024 * 1024)
+                    if not chunk:
+                        break
+                    output.write(chunk)
+            temporary.replace(destination)
+            return
+        except Exception as error:  # pragma: no cover - type depends on urllib backend
+            last_error = error
+            if temporary.exists():
+                temporary.unlink()
+            if attempt < attempts:
+                time.sleep(min(2.0, 0.25 * attempt))
+    assert last_error is not None
+    raise last_error

drbx/runtime/memory.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import os
+import subprocess
+import threading
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TypeVar
+
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class PeakRssMeasurement:
+    start_rss_bytes: int | None
+    end_rss_bytes: int | None
+    peak_rss_bytes: int | None
+    peak_rss_delta_bytes: int | None
+    sample_count: int
+    sampling_interval_seconds: float
+    status: str
+
+
+def measure_peak_rss(
+    function: Callable[[], T],
+    *,
+    sampling_interval_seconds: float = 0.10,
+) -> tuple[T, PeakRssMeasurement]:
+    samples: list[int] = []
+    sample_failures = 0
+    stop_event = threading.Event()
+    root_pid = os.getpid()
+
+    def append_sample() -> None:
+        nonlocal sample_failures
+        value = process_tree_rss_bytes(root_pid)
+        if value is None:
+            sample_failures += 1
+            return
+        samples.append(value)
+
+    append_sample()
+
+    def sample_loop() -> None:
+        while not stop_event.wait(sampling_interval_seconds):
+            append_sample()
+
+    sampler = threading.Thread(target=sample_loop, name="drbx-rss-sampler", daemon=True)
+    sampler.start()
+    try:
+        result = function()
+    finally:
+        stop_event.set()
+        sampler.join(timeout=max(1.0, 4.0 * sampling_interval_seconds))
+    append_sample()
+    start_rss_bytes = samples[0] if samples else None
+    end_rss_bytes = samples[-1] if samples else None
+    peak_rss_bytes = max(samples) if samples else None
+    peak_rss_delta_bytes = (
+        None
+        if peak_rss_bytes is None or start_rss_bytes is None
+        else max(int(peak_rss_bytes - start_rss_bytes), 0)
+    )
+    status = "sampled_process_tree_rss"
+    if not samples:
+        status = "unavailable"
+    elif sample_failures:
+        status = "sampled_process_tree_rss_with_partial_failures"
+    return result, PeakRssMeasurement(
+        start_rss_bytes=start_rss_bytes,
+        end_rss_bytes=end_rss_bytes,
+        peak_rss_bytes=peak_rss_bytes,
+        peak_rss_delta_bytes=peak_rss_delta_bytes,
+        sample_count=len(samples),
+        sampling_interval_seconds=float(sampling_interval_seconds),
+        status=status,
+    )
+
+
+def process_tree_rss_bytes(root_pid: int) -> int | None:
+    pids = process_tree_pids(root_pid)
+    rss_kib = 0
+    observed = False
+    for pid in pids:
+        pid_rss = process_rss_kib(pid)
+        if pid_rss is None:
+            continue
+        rss_kib += pid_rss
+        observed = True
+    return int(rss_kib * 1024) if observed else None
+
+
+def process_tree_pids(root_pid: int) -> list[int]:
+    pids: list[int] = [int(root_pid)]
+    frontier: list[int] = [int(root_pid)]
+    while frontier:
+        parent = frontier.pop()
+        try:
+            completed = subprocess.run(
+                ["pgrep", "-P", str(parent)],
+                check=False,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.DEVNULL,
+                text=True,
+            )
+        except OSError:
+            continue
+        if completed.returncode not in (0, 1):
+            continue
+        children = [int(value) for value in completed.stdout.split() if value.isdigit()]
+        for child in children:
+            if child not in pids:
+                pids.append(child)
+                frontier.append(child)
+    return pids
+
+
+def process_rss_kib(pid: int) -> int | None:
+    try:
+        completed = subprocess.run(
+            ["ps", "-o", "rss=", "-p", str(pid)],
+            check=False,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.DEVNULL,
+            text=True,
+        )
+    except OSError:
+        return None
+    if completed.returncode != 0:
+        return None
+    text = completed.stdout.strip()
+    if not text:
+        return None
+    try:
+        return int(float(text.splitlines()[-1].strip()))
+    except ValueError:
+        return None
+
+
+def bytes_to_mebibytes(value: int | None) -> float | None:
+    if value is None:
+        return None
+    return float(value / (1024.0 * 1024.0))