drbx 2.0.0.dev0__py3-none-any.whl

drbx/geometry/fci_maps.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import jax.numpy as jnp
+import numpy as np
+from netCDF4 import Dataset
+
+
+@dataclass(frozen=True)
+class FciMaps:
+    """Forward/backward traced-field-line maps in logical index coordinates."""
+
+    forward_x: jnp.ndarray
+    forward_z: jnp.ndarray
+    backward_x: jnp.ndarray
+    backward_z: jnp.ndarray
+    forward_boundary: jnp.ndarray
+    backward_boundary: jnp.ndarray
+    dphi: float
+
+    @property
+    def shape(self) -> tuple[int, int, int]:
+        return tuple(int(value) for value in self.forward_x.shape)
+
+
+def identity_fci_maps(*, nx: int, ny: int, nz: int, dphi: float = 1.0) -> FciMaps:
+    """Build maps whose field lines connect to the same `(x,z)` point."""
+
+    x = jnp.arange(nx, dtype=jnp.float64)[:, None, None]
+    z = jnp.arange(nz, dtype=jnp.float64)[None, None, :]
+    forward_x = jnp.broadcast_to(x, (nx, ny, nz))
+    forward_z = jnp.broadcast_to(z, (nx, ny, nz))
+    boundary = jnp.zeros((nx, ny, nz), dtype=bool)
+    return FciMaps(
+        forward_x=forward_x,
+        forward_z=forward_z,
+        backward_x=forward_x,
+        backward_z=forward_z,
+        forward_boundary=boundary,
+        backward_boundary=boundary,
+        dphi=float(dphi),
+    )
+
+
+def load_fci_maps_netcdf(path: str | Path) -> FciMaps:
+    """Load a compact FCI map bundle from a NetCDF grid file."""
+
+    resolved = Path(path)
+    with Dataset(resolved) as dataset:
+        forward_x = _read_variable(dataset, ("forward_xt_prime", "forward_x", "forward_R_index"))
+        forward_z = _read_variable(dataset, ("forward_zt_prime", "forward_z", "forward_Z_index"))
+        backward_x = _read_variable(dataset, ("backward_xt_prime", "backward_x", "backward_R_index"))
+        backward_z = _read_variable(dataset, ("backward_zt_prime", "backward_z", "backward_Z_index"))
+        dphi = _infer_dphi(dataset, forward_x.shape[1])
+
+    forward_boundary = (forward_x < 0.0) | (forward_x > float(forward_x.shape[0] - 1))
+    backward_boundary = (backward_x < 0.0) | (backward_x > float(backward_x.shape[0] - 1))
+    return FciMaps(
+        forward_x=jnp.asarray(forward_x, dtype=jnp.float64),
+        forward_z=jnp.asarray(forward_z, dtype=jnp.float64),
+        backward_x=jnp.asarray(backward_x, dtype=jnp.float64),
+        backward_z=jnp.asarray(backward_z, dtype=jnp.float64),
+        forward_boundary=jnp.asarray(forward_boundary),
+        backward_boundary=jnp.asarray(backward_boundary),
+        dphi=float(dphi),
+    )
+
+
+def _read_variable(dataset: Dataset, candidates: tuple[str, ...]) -> np.ndarray:
+    for name in candidates:
+        if name in dataset.variables:
+            return np.asarray(dataset.variables[name][:], dtype=np.float64)
+    raise KeyError(f"None of the FCI map variables were found: {', '.join(candidates)}")
+
+
+def _infer_dphi(dataset: Dataset, ny: int) -> float:
+    if "dy" in dataset.variables:
+        dy = np.asarray(dataset.variables["dy"][:], dtype=np.float64)
+        if dy.size:
+            return float(np.nanmean(dy))
+    if "yperiod" in dataset.ncattrs():
+        return float(dataset.getncattr("yperiod")) / float(ny)
+    return float(2.0 * np.pi / max(ny, 1))

drbx/geometry/island_divertor.py

@@ -0,0 +1,291 @@
+"""Analytic island-divertor stellarator field (B8).
+
+The rotating-ellipse embedding carrying an *island-divertor* magnetic field: the
+rotational transform is sheared, ``iota(x)`` rising from axis to edge, and
+resonant radial field perturbations ``B^x ~ eps sin(m theta - n zeta)`` open
+magnetic island chains at the rational surfaces ``iota = n/m``. With two
+overlapping chains the edge becomes stochastic: edge field lines wander
+radially and strike the wall — genuinely **open** field lines whose endpoint
+masks come from the field-line tracer, not from a hand-placed limiter — while
+the core below the resonances stays closed. This is the analytic island-divertor
+configuration of the B8 benchmark ladder rung.
+
+The field components (contravariant, with ``J`` the metric Jacobian):
+
+    B^zeta  = c / J
+    B^theta = iota(x) * c / J
+    B^x     = (c / J) * sum_k eps_k * sin(m_k theta - n_k zeta)
+
+so field lines obey ``dx/dzeta = sum_k eps_k sin(m_k theta - n_k zeta)`` and
+``dtheta/dzeta = iota(x)`` — the standard perturbed-twist-map form whose island
+half-width at a resonance is ``sqrt(eps m / iota')``. The helper
+:func:`island_divertor_field_line_rhs` exposes exactly that reduced system for
+Poincare / connection-length studies, and the full :class:`FciGeometry3D` uses
+the same components so the FCI tracer sees the same field.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+
+from .embedding import metric_from_position_fn
+from .fci_geometry import (
+    BFieldGeometry,
+    CellCenteredGrid3D,
+    FaceBFieldGeometry,
+    FaceMetricGeometry,
+    FciGeometry3D,
+    FciMaps3D,
+    Grid1D,
+    MetricGeometry,
+    Spacing3D,
+    _bmag_from_contravariant_components,
+    build_fci_maps_from_b_contravariant,
+    logical_grid_from_axis_vectors,
+)
+from .rotating_ellipse import rotating_ellipse_position
+
+__all__ = [
+    "IslandDivertorField",
+    "build_island_divertor_geometry",
+    "island_divertor_connection_length",
+    "island_divertor_field_line_rhs",
+]
+
+
+@dataclass(frozen=True)
+class IslandDivertorField:
+    """The analytic island-divertor field parameters.
+
+    ``resonances`` is a tuple of ``(m, n, epsilon)`` triplets; each opens an
+    island chain at the surface where ``iota(x) = n / m``. ``iota`` varies
+    linearly from ``iota_axis`` at ``x_min`` to ``iota_edge`` at ``x_max``.
+    """
+
+    x_min: float = 0.2
+    x_max: float = 1.0
+    iota_axis: float = 0.55
+    iota_edge: float = 0.82
+    resonances: tuple[tuple[int, int, float], ...] = ((3, 2, 0.005), (4, 3, 0.005), (5, 4, 0.005))
+
+    def iota(self, x):
+        x_norm = (x - self.x_min) / (self.x_max - self.x_min)
+        return self.iota_axis + (self.iota_edge - self.iota_axis) * x_norm
+
+    def radial_perturbation(self, theta, zeta):
+        perturbation = jnp.zeros_like(jnp.asarray(theta, dtype=jnp.float64) + jnp.asarray(zeta))
+        for m, n, epsilon in self.resonances:
+            perturbation = perturbation + float(epsilon) * jnp.sin(m * theta - n * zeta)
+        return perturbation
+
+
+def island_divertor_field_line_rhs(field: IslandDivertorField, x, theta, zeta):
+    """Field-line equations ``(dx/dzeta, dtheta/dzeta)`` of the reduced system."""
+
+    return field.radial_perturbation(theta, zeta), field.iota(x)
+
+
+def island_divertor_connection_length(
+    field: IslandDivertorField,
+    x: jnp.ndarray,
+    theta: jnp.ndarray,
+    zeta: jnp.ndarray,
+    *,
+    direction: int = 1,
+    max_transits: int = 40,
+    steps_per_transit: int = 64,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """Trace field lines and return ``(connection_length_transits, is_open)``.
+
+    Integrates the reduced field-line system with RK4 from every starting point
+    (arrays broadcast together) for up to ``max_transits`` toroidal transits in
+    the given ``direction``. A line is **open** if it leaves the radial domain;
+    its connection length is the number of transits to the exit (open lines) or
+    ``max_transits`` (closed lines). Pure JAX (``vmap``/``scan``), so a full 3D
+    grid of starting points traces in one shot.
+    """
+
+    x0, theta0, zeta0 = jnp.broadcast_arrays(
+        jnp.asarray(x, dtype=jnp.float64),
+        jnp.asarray(theta, dtype=jnp.float64),
+        jnp.asarray(zeta, dtype=jnp.float64),
+    )
+    shape = x0.shape
+    dz = float(direction) * 2.0 * jnp.pi / float(steps_per_transit)
+
+    def rhs(x_value, theta_value, zeta_value):
+        return field.radial_perturbation(theta_value, zeta_value), field.iota(x_value)
+
+    def step(carry, _):
+        x_value, theta_value, zeta_value, alive, steps_alive = carry
+        k1x, k1t = rhs(x_value, theta_value, zeta_value)
+        k2x, k2t = rhs(x_value + 0.5 * dz * k1x, theta_value + 0.5 * dz * k1t, zeta_value + 0.5 * dz)
+        k3x, k3t = rhs(x_value + 0.5 * dz * k2x, theta_value + 0.5 * dz * k2t, zeta_value + 0.5 * dz)
+        k4x, k4t = rhs(x_value + dz * k3x, theta_value + dz * k3t, zeta_value + dz)
+        x_next = x_value + dz / 6.0 * (k1x + 2 * k2x + 2 * k3x + k4x)
+        theta_next = theta_value + dz / 6.0 * (k1t + 2 * k2t + 2 * k3t + k4t)
+        # Only the outer boundary is material (the divertor wall); the inner
+        # boundary is the artificial core cutoff, so reflect there.
+        x_next = jnp.maximum(x_next, field.x_min)
+        alive_next = alive & (x_next < field.x_max)
+        x_next = jnp.where(alive, x_next, x_value)
+        theta_next = jnp.where(alive, theta_next, theta_value)
+        return (x_next, theta_next, zeta_value + dz, alive_next, steps_alive + alive_next), None
+
+    initial = (
+        x0.ravel(),
+        theta0.ravel(),
+        zeta0.ravel(),
+        jnp.ones(x0.size, dtype=bool),
+        jnp.zeros(x0.size, dtype=jnp.int32),
+    )
+    (final_x, _, _, alive, steps_alive), _ = jax.lax.scan(
+        step, initial, None, length=int(max_transits) * int(steps_per_transit)
+    )
+    del final_x
+    connection_transits = steps_alive.astype(jnp.float64) / float(steps_per_transit)
+    return connection_transits.reshape(shape), (~alive).reshape(shape)
+
+
+def build_island_divertor_geometry(
+    shape: tuple[int, int, int],
+    *,
+    field: IslandDivertorField = IslandDivertorField(),
+    r0: float = 3.0,
+    elongation: float = 0.35,
+    n_field_periods: int = 1,
+    c_phi: float = 3.0,
+    construct_fci_maps: bool = False,
+    map_substeps: int = 8,
+    open_field_line_masks: bool = False,
+    mask_max_transits: int = 30,
+) -> FciGeometry3D:
+    """Build the island-divertor geometry on the rotating-ellipse embedding.
+
+    With ``open_field_line_masks=True`` the endpoint masks come from
+    multi-transit field-line tracing (:func:`island_divertor_connection_length`):
+    a cell is an open endpoint if its field line leaves the radial domain within
+    ``mask_max_transits`` toroidal transits — the island-divertor scrape-off
+    layer *emerges from the field itself* (stochastic-edge lines are open with
+    finite connection length, core lines stay closed). ``construct_fci_maps``
+    additionally traces the one-cell FCI parallel maps for the operators.
+    """
+
+    nx, ny, nz = shape
+    x_faces = jnp.linspace(field.x_min, field.x_max, nx + 1, dtype=jnp.float64)
+    theta_faces = jnp.linspace(0.0, 2.0 * jnp.pi, ny + 1, dtype=jnp.float64)
+    zeta_faces = jnp.linspace(0.0, 2.0 * jnp.pi, nz + 1, dtype=jnp.float64)
+    grid = CellCenteredGrid3D(
+        x=Grid1D(centers=0.5 * (x_faces[:-1] + x_faces[1:]), faces=x_faces),
+        y=Grid1D(centers=0.5 * (theta_faces[:-1] + theta_faces[1:]), faces=theta_faces),
+        z=Grid1D(centers=0.5 * (zeta_faces[:-1] + zeta_faces[1:]), faces=zeta_faces),
+    )
+    target_shape = grid.shape
+
+    def _position(u: jax.Array) -> jax.Array:
+        return rotating_ellipse_position(
+            u[0], u[1], u[2], r0=r0, elongation=elongation, n_field_periods=n_field_periods
+        )
+
+    def _metric(logical_grid: jax.Array) -> MetricGeometry:
+        return metric_from_position_fn(_position, logical_grid)
+
+    def _bfield(logical_grid: jax.Array, metric: MetricGeometry) -> BFieldGeometry:
+        x = logical_grid[..., 0]
+        theta = logical_grid[..., 1]
+        zeta = logical_grid[..., 2]
+        toroidal = float(c_phi) / metric.J
+        B_contra = jnp.stack(
+            (
+                field.radial_perturbation(theta, zeta) * toroidal,
+                field.iota(x) * toroidal,
+                toroidal,
+            ),
+            axis=-1,
+        )
+        return BFieldGeometry(
+            B_contra=B_contra,
+            Bmag=_bmag_from_contravariant_components(B_contra, metric.g_cov),
+        )
+
+    def _logical(x_axis, y_axis, z_axis):
+        return logical_grid_from_axis_vectors(x_axis, y_axis, z_axis)
+
+    cell_logical = _logical(grid.x.centers, grid.y.centers, grid.z.centers)
+    cell_metric = _metric(cell_logical)
+    cell_bfield = _bfield(cell_logical, cell_metric)
+    face_metric = FaceMetricGeometry(
+        x=_metric(_logical(grid.x.faces, grid.y.centers, grid.z.centers)),
+        y=_metric(_logical(grid.x.centers, grid.y.faces, grid.z.centers)),
+        z=_metric(_logical(grid.x.centers, grid.y.centers, grid.z.faces)),
+    )
+    face_bfield = FaceBFieldGeometry(
+        x=_bfield(_logical(grid.x.faces, grid.y.centers, grid.z.centers), face_metric.x),
+        y=_bfield(_logical(grid.x.centers, grid.y.faces, grid.z.centers), face_metric.y),
+        z=_bfield(_logical(grid.x.centers, grid.y.centers, grid.z.faces), face_metric.z),
+    )
+
+    if construct_fci_maps:
+        map_fields = build_fci_maps_from_b_contravariant(
+            grid,
+            cell_bfield.B_contra,
+            cell_bfield.Bmag,
+            periodic_axes=(False, True, True),
+            substeps=int(map_substeps),
+        )
+    else:
+        ones = jnp.ones(target_shape, dtype=jnp.float64)
+        zeros = jnp.zeros(target_shape, dtype=jnp.float64)
+        map_fields = {
+            "forward_x": zeros,
+            "forward_y": zeros,
+            "backward_x": zeros,
+            "backward_y": zeros,
+            "forward_endpoint_x": zeros,
+            "forward_endpoint_y": zeros,
+            "forward_endpoint_z": zeros,
+            "backward_endpoint_x": zeros,
+            "backward_endpoint_y": zeros,
+            "backward_endpoint_z": zeros,
+            "forward_length": ones,
+            "backward_length": ones,
+            "forward_boundary": zeros.astype(bool),
+            "backward_boundary": zeros.astype(bool),
+        }
+
+    if open_field_line_masks:
+        forward_transits, forward_open = island_divertor_connection_length(
+            field, cell_logical[..., 0], cell_logical[..., 1], cell_logical[..., 2],
+            direction=+1, max_transits=mask_max_transits,
+        )
+        backward_transits, backward_open = island_divertor_connection_length(
+            field, cell_logical[..., 0], cell_logical[..., 1], cell_logical[..., 2],
+            direction=-1, max_transits=mask_max_transits,
+        )
+        del forward_transits, backward_transits
+        map_fields["forward_boundary"] = jnp.asarray(map_fields["forward_boundary"], dtype=bool) | forward_open
+        map_fields["backward_boundary"] = jnp.asarray(map_fields["backward_boundary"], dtype=bool) | backward_open
+
+    maps = FciMaps3D(**{name: map_fields[name] for name in (
+        "forward_x", "forward_y", "backward_x", "backward_y",
+        "forward_endpoint_x", "forward_endpoint_y", "forward_endpoint_z",
+        "backward_endpoint_x", "backward_endpoint_y", "backward_endpoint_z",
+        "forward_length", "backward_length", "forward_boundary", "backward_boundary",
+    )})
+    spacing = Spacing3D(
+        dx=jnp.broadcast_to(grid.x.widths[:, None, None], target_shape),
+        dy=jnp.broadcast_to(grid.y.widths[None, :, None], target_shape),
+        dz=jnp.broadcast_to(grid.z.widths[None, None, :], target_shape),
+    )
+    return FciGeometry3D(
+        grid=grid,
+        maps=maps,
+        spacing=spacing,
+        cell_metric=cell_metric,
+        face_metric=face_metric,
+        cell_bfield=cell_bfield,
+        face_bfield=face_bfield,
+    )

drbx/geometry/metric_tensor.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+import numpy as np
+
+
+@dataclass(frozen=True)
+class MetricTensor3D:
+    """Complete 3D metric payload used by non-axisymmetric geometry kernels."""
+
+    dx: jnp.ndarray
+    dy: jnp.ndarray
+    dz: jnp.ndarray
+    J: jnp.ndarray
+    Bxy: jnp.ndarray
+    g11: jnp.ndarray
+    g22: jnp.ndarray
+    g33: jnp.ndarray
+    g12: jnp.ndarray
+    g13: jnp.ndarray
+    g23: jnp.ndarray
+    g_11: jnp.ndarray
+    g_22: jnp.ndarray
+    g_33: jnp.ndarray
+    g_12: jnp.ndarray
+    g_13: jnp.ndarray
+    g_23: jnp.ndarray
+
+    @property
+    def shape(self) -> tuple[int, int, int]:
+        return tuple(int(value) for value in self.J.shape)
+
+
+def metric_inverse_residual(metric: MetricTensor3D) -> jnp.ndarray:
+    """Return `max(abs(g^ik g_kj - delta^i_j))` over the grid."""
+
+    contravariant = jnp.stack(
+        [
+            jnp.stack([metric.g11, metric.g12, metric.g13], axis=-1),
+            jnp.stack([metric.g12, metric.g22, metric.g23], axis=-1),
+            jnp.stack([metric.g13, metric.g23, metric.g33], axis=-1),
+        ],
+        axis=-2,
+    )
+    covariant = jnp.stack(
+        [
+            jnp.stack([metric.g_11, metric.g_12, metric.g_13], axis=-1),
+            jnp.stack([metric.g_12, metric.g_22, metric.g_23], axis=-1),
+            jnp.stack([metric.g_13, metric.g_23, metric.g_33], axis=-1),
+        ],
+        axis=-2,
+    )
+    product = jnp.einsum("...ik,...kj->...ij", contravariant, covariant)
+    identity = jnp.eye(3, dtype=product.dtype)
+    return jnp.max(jnp.abs(product - identity))
+
+
+def build_metric_report(metric: MetricTensor3D) -> dict[str, object]:
+    """Build finite/positive/inverse-consistency diagnostics for a metric payload."""
+
+    report: dict[str, object] = {
+        "shape": list(metric.shape),
+        "inverse_residual_linf": float(metric_inverse_residual(metric)),
+        "fields": {},
+    }
+    for name in (
+        "J",
+        "Bxy",
+        "g11",
+        "g22",
+        "g33",
+        "g12",
+        "g13",
+        "g23",
+        "g_11",
+        "g_22",
+        "g_33",
+        "g_12",
+        "g_13",
+        "g_23",
+    ):
+        values = np.asarray(getattr(metric, name), dtype=np.float64)
+        report["fields"][name] = {
+            "finite": bool(np.all(np.isfinite(values))),
+            "minimum": float(np.nanmin(values)),
+            "maximum": float(np.nanmax(values)),
+            "mean": float(np.nanmean(values)),
+            "std": float(np.nanstd(values)),
+        }
+    report["passed"] = (
+        bool(report["fields"]["J"]["finite"])
+        and bool(report["fields"]["Bxy"]["finite"])
+        and float(report["fields"]["J"]["minimum"]) > 0.0
+        and float(report["fields"]["Bxy"]["minimum"]) > 0.0
+        and float(report["inverse_residual_linf"]) < 1.0e-8
+    )
+    return report

drbx/geometry/open_slab.py

@@ -0,0 +1,150 @@
+"""Open-field-line slab SOL flux-tube geometry.
+
+This module provides :func:`build_open_slab_geometry`, a Cartesian scrape-off-layer
+(SOL) flux tube with a straight magnetic field along the parallel coordinate
+``z`` and **open** field lines: every field line terminates on a material target
+plate at ``z = 0`` and ``z = L_parallel``. The endpoint masks
+(``forward_boundary`` at the ``z = L`` target, ``backward_boundary`` at the
+``z = 0`` target) are exactly the open-field-line target masks the sheath /
+recycling closure (:mod:`drbx.native.fci_sheath_recycling`) consumes.
+
+The logical coordinates are ``(x, y, z)`` with ``x`` the radial SOL coordinate,
+``y`` the periodic binormal, and ``z`` the parallel (connection-length)
+coordinate bounded by the two targets. The metric is Cartesian/identity — the
+point of this geometry is the *open* parallel topology, not curvature — so the
+FCI parallel gradient reduces to ``d/dz`` and the target plates carry the Bohm
+sheath. It is the open-field-line counterpart to the closed flux tubes
+(rotating ellipse, shifted torus) elsewhere in the package.
+"""
+
+from __future__ import annotations
+
+import jax.numpy as jnp
+
+from .fci_geometry import (
+    BFieldGeometry,
+    CellCenteredGrid3D,
+    FaceBFieldGeometry,
+    FaceMetricGeometry,
+    FciGeometry3D,
+    FciMaps3D,
+    Grid1D,
+    MetricGeometry,
+    Spacing3D,
+)
+
+__all__ = ["build_open_slab_geometry"]
+
+
+def _identity_metric(shape: tuple[int, int, int]) -> MetricGeometry:
+    ones = jnp.ones(shape, dtype=jnp.float64)
+    zeros = jnp.zeros(shape, dtype=jnp.float64)
+    return MetricGeometry(
+        J=ones,
+        g11=ones,
+        g22=ones,
+        g33=ones,
+        g12=zeros,
+        g13=zeros,
+        g23=zeros,
+        g_11=ones,
+        g_22=ones,
+        g_33=ones,
+        g_12=zeros,
+        g_13=zeros,
+        g_23=zeros,
+    )
+
+
+def _parallel_bfield(shape: tuple[int, int, int]) -> BFieldGeometry:
+    return BFieldGeometry(
+        B_contra=jnp.broadcast_to(jnp.array([0.0, 0.0, 1.0], dtype=jnp.float64), shape + (3,)),
+        Bmag=jnp.ones(shape, dtype=jnp.float64),
+    )
+
+
+def build_open_slab_geometry(
+    shape: tuple[int, int, int],
+    *,
+    radial_extent: float = 1.0,
+    binormal_extent: float = 1.0,
+    parallel_length: float = 40.0,
+) -> FciGeometry3D:
+    """Build an open-field-line slab SOL flux tube of the requested ``shape``.
+
+    Args:
+        shape: ``(nx, ny, nz)`` cell-centered grid. ``nz`` is the number of cells
+            along the field between the two targets; ``nx``/``ny`` are the radial
+            and binormal extents (use small values for a flux-tube study).
+        radial_extent, binormal_extent: physical sizes of the ``x`` and ``y``
+            axes.
+        parallel_length: connection length ``L_parallel`` between the ``z = 0``
+            and ``z = L`` target plates.
+
+    Returns:
+        A :class:`FciGeometry3D` whose field lines are open: ``forward_boundary``
+        is set on the ``z = L`` target plane and ``backward_boundary`` on the
+        ``z = 0`` target plane, so :func:`drbx.native.fci_sheath_recycling`.
+        `build_fci_target_masks` marks exactly the two target planes.
+    """
+
+    nx, ny, nz = shape
+    x_faces = jnp.linspace(0.0, float(radial_extent), nx + 1, dtype=jnp.float64)
+    y_faces = jnp.linspace(0.0, float(binormal_extent), ny + 1, dtype=jnp.float64)
+    z_faces = jnp.linspace(0.0, float(parallel_length), nz + 1, dtype=jnp.float64)
+    grid = CellCenteredGrid3D(
+        x=Grid1D(centers=0.5 * (x_faces[:-1] + x_faces[1:]), faces=x_faces),
+        y=Grid1D(centers=0.5 * (y_faces[:-1] + y_faces[1:]), faces=y_faces),
+        z=Grid1D(centers=0.5 * (z_faces[:-1] + z_faces[1:]), faces=z_faces),
+    )
+    target_shape = (nx, ny, nz)
+    zeros = jnp.zeros(target_shape, dtype=jnp.float64)
+    ones = jnp.ones(target_shape, dtype=jnp.float64)
+
+    cell_metric = _identity_metric(target_shape)
+    face_metric = FaceMetricGeometry(
+        x=_identity_metric((nx + 1, ny, nz)),
+        y=_identity_metric((nx, ny + 1, nz)),
+        z=_identity_metric((nx, ny, nz + 1)),
+    )
+    cell_bfield = _parallel_bfield(target_shape)
+    face_bfield = FaceBFieldGeometry(
+        x=_parallel_bfield((nx + 1, ny, nz)),
+        y=_parallel_bfield((nx, ny + 1, nz)),
+        z=_parallel_bfield((nx, ny, nz + 1)),
+    )
+
+    # Open field lines: the forward map exits the domain at the z = L target and
+    # the backward map exits at the z = 0 target. Interior planes stay closed.
+    forward_boundary = jnp.zeros(target_shape, dtype=bool).at[:, :, -1].set(True)
+    backward_boundary = jnp.zeros(target_shape, dtype=bool).at[:, :, 0].set(True)
+    maps = FciMaps3D(
+        forward_x=zeros,
+        forward_y=zeros,
+        backward_x=zeros,
+        backward_y=zeros,
+        forward_endpoint_x=zeros,
+        forward_endpoint_y=zeros,
+        forward_endpoint_z=zeros,
+        backward_endpoint_x=zeros,
+        backward_endpoint_y=zeros,
+        backward_endpoint_z=zeros,
+        forward_length=ones,
+        backward_length=ones,
+        forward_boundary=forward_boundary,
+        backward_boundary=backward_boundary,
+    )
+    spacing = Spacing3D(
+        dx=jnp.broadcast_to(grid.x.widths[:, None, None], target_shape),
+        dy=jnp.broadcast_to(grid.y.widths[None, :, None], target_shape),
+        dz=jnp.broadcast_to(grid.z.widths[None, None, :], target_shape),
+    )
+    return FciGeometry3D(
+        grid=grid,
+        maps=maps,
+        spacing=spacing,
+        cell_metric=cell_metric,
+        face_metric=face_metric,
+        cell_bfield=cell_bfield,
+        face_bfield=face_bfield,
+    )