drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/__init__.py

@@ -0,0 +1,224 @@
+from ..runtime import configure_jax_runtime
+
+configure_jax_runtime()
+
+from .deck_runner import (
+    NativeRestartState,
+    NativeRunResult,
+    build_restart_state,
+    run_config_case,
+    run_input_case,
+)
+
+__all__ = [
+    "NativeRestartState",
+    "NativeRunResult",
+    "build_restart_state",
+    "run_config_case",
+    "run_input_case",
+]
+
+# --- FCI stack incorporated from PR #3 (Aiken Xie, branch 3D_fci) ---
+from .fci_halo import (
+    FciCutWallValueEvaluator,
+    GhostFillWeights1D,
+    HaloExchange3D,
+    LocalFciCutWallValueEvaluator,
+    LocalPeriodicTopologyRule3D,
+    LocalStateAndBoundaryPreparer3D,
+    PhysicalGhostCellFiller3D,
+    PolarAxisRegularScalarRule3D,
+    PolarAxisRegularVectorRule3D,
+    PreparedLocalState3D,
+    RemoteFciDependencyExchange,
+    TopologyHaloFiller3D,
+    make_default_topology_halo_filler_3d,
+)
+from .fci_boundaries import (
+    ConservativeStencil3D,
+    FaceGradientStencil3D,
+    FourFieldBoundaryConditions,
+    LocalBoundaryConditionBuilder,
+    LocalBoundaryData3D,
+    LocalCoordinateNormalDerivativeConstructor3D,
+    LocalCoordinateSideValues1D,
+    LocalCoordinateSideValues3D,
+    LocalBoundaryFaceBC3D,
+    LocalCutWallBC3D,
+    LocalCutWallGeometry3D,
+    LocalCutWallNormalDerivativeConstructor3D,
+    LocalCutWallValueReconstructor3D,
+    LocalStencil1D,
+    LocalStencil3D,
+)
+from .fci_model import (
+    FciFieldBundle,
+    FciModelState,
+    assert_matching_field_names,
+    extract_owned_field_from_halo,
+    extract_owned_state_from_halo,
+    inject_owned_field_to_halo,
+    inject_owned_vector_field_to_halo,
+    inject_owned_state_to_halo,
+    update_halo_owned_slice,
+    update_state_halo_owned_slices,
+)
+from .fci_time_integrator import Rk4StepResult, rk4_step, sum_stage_outputs
+from .fci_sharding import (
+    Sharded2FieldStepInfo,
+    ShardedFciGeometry3D,
+    assemble_local_fci_geometry,
+    build_local_fci_geometries,
+    make_shard_mesh,
+    make_sharded_2field_step,
+)
+from .fci_operators import (
+    PerpLaplacianMgHierarchy,
+    PerpLaplacianMgLevel,
+    build_perp_laplacian_mg_hierarchy,
+    build_conservative_stencil_from_field,
+    build_local_conservative_stencil_from_field,
+    build_local_perp_laplacian_stencil,
+    build_local_projected_laplacian_flux_stencil,
+    build_perp_laplacian_stencil,
+    build_perp_laplacian_face_projectors,
+    build_local_perp_laplacian_face_projectors,
+    divergence_conservative_op,
+    local_divergence_conservative_op,
+    curvature_op,
+    local_grad_parallel_op_direct,
+    local_grad_parallel_op_fci,
+    local_grad_perp_op_direct,
+    local_parallel_laplacian_direct_op,
+    local_perp_laplacian_local_op,
+    local_perp_laplacian_conservative_op,
+    grad_parallel_op_fci,
+    grad_perp_op,
+    mg_apply_preconditioner,
+    perp_laplacian_local_op,
+    perp_laplacian_conservative_op,
+    parallel_laplacian_direct_op,
+    parallel_laplacian_conservative_op,
+    local_poisson_bracket_op,
+    local_curvature_op,
+    poisson_bracket_op,
+)
+from .fci_2_field_rhs import Fci2FieldRhsParameters, Fci2FieldRhsResult, Fci2FieldState, compute_2field_rhs
+from .fci_drb_EB_rhs import (
+    FciDrbEBBoundaryConditions,
+    FciDrbEBRhsParameters,
+    FciDrbEBRhsResult,
+    FciDrbEBState,
+    compute_fci_drb_eb_rhs,
+)
+from .fci_4_field_rhs import (
+    Fci4FieldBlobParameters,
+    Fci4FieldFreeDecayParameters,
+    Fci4FieldRhsParameters,
+    Fci4FieldRhsResult,
+    Fci4FieldState,
+    compute_4field_blob_rhs,
+    compute_4field_curvature,
+    compute_4field_diffusion,
+    compute_4field_free_decay_rhs,
+    compute_4field_rhs,
+    compute_4field_poisson_diffusion,
+)
+__all__ += [
+    "ConservativeStencil3D",
+    "FaceGradientStencil3D",
+    "Fci2FieldRhsParameters",
+    "Fci2FieldRhsResult",
+    "Fci2FieldState",
+    "Fci4FieldBlobParameters",
+    "Fci4FieldFreeDecayParameters",
+    "Fci4FieldRhsParameters",
+    "Fci4FieldRhsResult",
+    "Fci4FieldState",
+    "FciCutWallValueEvaluator",
+    "FciDrbEBBoundaryConditions",
+    "FciDrbEBRhsParameters",
+    "FciDrbEBRhsResult",
+    "FciDrbEBState",
+    "FciFieldBundle",
+    "FciModelState",
+    "FourFieldBoundaryConditions",
+    "GhostFillWeights1D",
+    "HaloExchange3D",
+    "LocalBoundaryConditionBuilder",
+    "LocalBoundaryData3D",
+    "LocalBoundaryFaceBC3D",
+    "LocalCoordinateNormalDerivativeConstructor3D",
+    "LocalCoordinateSideValues1D",
+    "LocalCoordinateSideValues3D",
+    "LocalCutWallBC3D",
+    "LocalCutWallGeometry3D",
+    "LocalCutWallNormalDerivativeConstructor3D",
+    "LocalCutWallValueReconstructor3D",
+    "LocalFciCutWallValueEvaluator",
+    "LocalPeriodicTopologyRule3D",
+    "LocalStateAndBoundaryPreparer3D",
+    "LocalStencil1D",
+    "LocalStencil3D",
+    "PerpLaplacianMgHierarchy",
+    "PerpLaplacianMgLevel",
+    "PhysicalGhostCellFiller3D",
+    "PolarAxisRegularScalarRule3D",
+    "PolarAxisRegularVectorRule3D",
+    "PreparedLocalState3D",
+    "RemoteFciDependencyExchange",
+    "Rk4StepResult",
+    "TopologyHaloFiller3D",
+    "assert_matching_field_names",
+    "build_conservative_stencil_from_field",
+    "build_local_conservative_stencil_from_field",
+    "build_local_perp_laplacian_face_projectors",
+    "build_local_perp_laplacian_stencil",
+    "build_local_projected_laplacian_flux_stencil",
+    "build_perp_laplacian_face_projectors",
+    "build_perp_laplacian_mg_hierarchy",
+    "build_perp_laplacian_stencil",
+    "compute_2field_rhs",
+    "compute_4field_blob_rhs",
+    "compute_4field_curvature",
+    "compute_4field_diffusion",
+    "compute_4field_free_decay_rhs",
+    "compute_4field_poisson_diffusion",
+    "compute_4field_rhs",
+    "compute_fci_drb_eb_rhs",
+    "curvature_op",
+    "divergence_conservative_op",
+    "extract_owned_field_from_halo",
+    "extract_owned_state_from_halo",
+    "grad_parallel_op_fci",
+    "grad_perp_op",
+    "inject_owned_field_to_halo",
+    "inject_owned_state_to_halo",
+    "inject_owned_vector_field_to_halo",
+    "local_curvature_op",
+    "local_divergence_conservative_op",
+    "local_grad_parallel_op_direct",
+    "local_grad_parallel_op_fci",
+    "local_grad_perp_op_direct",
+    "local_parallel_laplacian_direct_op",
+    "local_perp_laplacian_conservative_op",
+    "local_perp_laplacian_local_op",
+    "local_poisson_bracket_op",
+    "make_default_topology_halo_filler_3d",
+    "mg_apply_preconditioner",
+    "parallel_laplacian_conservative_op",
+    "parallel_laplacian_direct_op",
+    "perp_laplacian_conservative_op",
+    "perp_laplacian_local_op",
+    "poisson_bracket_op",
+    "rk4_step",
+    "Sharded2FieldStepInfo",
+    "ShardedFciGeometry3D",
+    "assemble_local_fci_geometry",
+    "build_local_fci_geometries",
+    "make_shard_mesh",
+    "make_sharded_2field_step",
+    "sum_stage_outputs",
+    "update_halo_owned_slice",
+    "update_state_halo_owned_slices",
+]

drbx/native/array_backend.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from numbers import Number
+from typing import Any
+
+import jax.numpy as jnp
+import numpy as np
+
+
+def is_jax_array(value: Any) -> bool:
+    """Return whether ``value`` is a JAX array or tracer.
+
+    Several hot-path kernels mix static NumPy metric arrays with dynamic JAX
+    state arrays during ``jax.jvp`` or ``jax.linearize``. Backend selection must
+    prefer JAX whenever any dynamic argument is a JAX value; otherwise a single
+    NumPy metric would force ``np.asarray`` on a tracer.
+    """
+
+    if value is None:
+        return False
+    if isinstance(value, (np.ndarray, np.generic, Number, bool)):
+        return False
+    module = type(value).__module__
+    if module == "numpy" or module.startswith("numpy."):
+        return False
+    return hasattr(value, "aval") or module.startswith("jax") or module.startswith("jaxlib")
+
+
+def use_jax_backend(*values: Any) -> bool:
+    """Select JAX when any argument is a JAX array/tracer."""
+
+    for value in values:
+        if value is not None and is_jax_array(value):
+            return True
+    return False
+
+
+def asarray(value: Any, *, use_jax: bool):
+    """Cast to the selected floating array backend."""
+
+    if use_jax:
+        return jnp.asarray(value, dtype=jnp.float64)
+    return np.asarray(value, dtype=np.float64)
+
+
+def zeros_like(value: Any, *, use_jax: bool):
+    """Create a float64 zeros array on the selected backend."""
+
+    if use_jax:
+        return jnp.zeros_like(jnp.asarray(value, dtype=jnp.float64), dtype=jnp.float64)
+    return np.zeros_like(np.asarray(value, dtype=np.float64), dtype=np.float64)
+
+
+def maximum_reduce(values: tuple[Any, ...], *, use_jax: bool):
+    """Backend-preserving equivalent of ``np.maximum.reduce``."""
+
+    if not values:
+        raise ValueError("maximum_reduce requires at least one value.")
+    if use_jax:
+        result = jnp.asarray(values[0], dtype=jnp.float64)
+        for value in values[1:]:
+            result = jnp.maximum(result, jnp.asarray(value, dtype=jnp.float64))
+        return result
+    return np.maximum.reduce(tuple(np.asarray(value, dtype=np.float64) for value in values))