drbx 2.0.0.dev0__py3-none-any.whl

drbx/data/atomic_rates/rec_AMJUEL_H.x_2.3.13a.json

@@ -0,0 +1,213 @@
+{
+    "info": {
+        "sigma_v_coeffs_ref": "Amjuel H.4 2.3.13a, page 181",
+        "sigma_v_E_coeffs_ref": "Amjuel H.10 2.3.13a, page 295",
+        "reaction_lbl": "e + he+ -> he",
+        "sigma_v_E_notes": "Helium multi-step model, here recombination: radiative + threebody + dielectronic. Fujimoto Formulation II (only ground level transported, no meta-stables kept explicit). The quantity given here happens to be the radiation loss. The loss of potential energy still needs to be subtracted to make this a total electron energy loss (or gain) rate.",
+        "sigma_v_notes": "Helium multi-step model, here recombination: radiative + threebody + dielectronic. Fujimoto Formulation II (only ground level transported, no metastables kept explicit).",
+        "includes_sigma_v_e": true
+    },
+    "electron_heating": 24.586,
+    "sigma_v_coeffs": [
+        [
+            -28.72754373123,
+            -0.006171082987797,
+            0.02414548639597,
+            -0.007188662067622,
+            0.0009481268604767,
+            -1.958887458637e-05,
+            -5.507786383328e-06,
+            4.35828868693e-07,
+            -9.50327209101e-09
+        ],
+        [
+            1.564233603544,
+            -0.03972220721457,
+            -0.04466712599181,
+            0.01247359158796,
+            -0.001660591942878,
+            6.019181402025e-05,
+            3.800156798817e-06,
+            -3.377807793756e-07,
+            6.828447501225e-09
+        ],
+        [
+            -6.182140631482,
+            0.1626641668186,
+            0.03366589582541,
+            -0.007413737965595,
+            0.001220189896183,
+            -9.50529572475e-05,
+            4.459492214068e-06,
+            -1.552772441333e-07,
+            2.866586118879e-09
+        ],
+        [
+            5.459428677778,
+            -0.1700323494998,
+            -0.01540106384088,
+            0.0009524545793262,
+            -8.734341535385e-05,
+            -2.796027477899e-06,
+            4.561981097438e-07,
+            4.940311502014e-09,
+            -6.52572501076e-10
+        ],
+        [
+            -2.128115924661,
+            0.07233939709414,
+            0.005819196258503,
+            -7.655935845761e-05,
+            -1.83794906705e-05,
+            4.72578983298e-06,
+            -3.99778241186e-07,
+            1.036731541123e-08,
+            -3.373845712183e-11
+        ],
+        [
+            0.4373730373037,
+            -0.01574917019835,
+            -0.001456253436544,
+            4.772491845078e-05,
+            -1.827059132463e-06,
+            -6.94116329271e-08,
+            2.716740135949e-08,
+            -1.143121626264e-09,
+            1.295139027087e-11
+        ],
+        [
+            -0.04972257208732,
+            0.001866175274689,
+            0.0002047337498511,
+            -1.004438052808e-05,
+            7.59073486585e-07,
+            -6.771179147667e-08,
+            1.218720257518e-09,
+            6.78702447954e-11,
+            -2.432253541918e-12
+        ],
+        [
+            0.002967287371427,
+            -0.0001147811325052,
+            -1.460813593905e-05,
+            7.422385993164e-07,
+            -3.281946488134e-08,
+            2.164459880579e-09,
+            1.113868237282e-10,
+            -1.513922678655e-11,
+            3.951084520871e-13
+        ],
+        [
+            -7.271204747116e-05,
+            2.874049670122e-06,
+            4.124421172202e-07,
+            -1.689203971933e-08,
+            -9.071172814458e-10,
+            1.844295219334e-10,
+            -2.055023511556e-11,
+            1.101902611511e-12,
+            -2.206082129473e-14
+        ]
+    ],
+    "sigma_v_E_coeffs": [
+        [
+            -25.38377692766,
+            -0.04826880987619,
+            0.0679657596731,
+            -0.0240171002139,
+            0.004130156138736,
+            -0.0003494803122018,
+            1.34502510054e-05,
+            -1.323917127568e-07,
+            -2.551716207606e-09
+        ],
+        [
+            2.472758419513,
+            0.1668058989207,
+            -0.1265192781981,
+            0.02938171777028,
+            -0.00221652505507,
+            -0.0001261523686946,
+            2.701401918133e-05,
+            -1.370446267883e-06,
+            2.313673787201e-08
+        ],
+        [
+            -8.864417999957,
+            -0.1882326730037,
+            0.119402867431,
+            -0.02382836629119,
+            0.001134820469638,
+            0.0001782113978272,
+            -2.305554399898e-05,
+            9.43029409318e-07,
+            -1.305188423829e-08
+        ],
+        [
+            8.394970578944,
+            0.08397993216045,
+            -0.0579697281374,
+            0.01158600348753,
+            -0.0007504743150582,
+            -2.602911694939e-05,
+            4.568209602293e-06,
+            -1.458110560501e-07,
+            9.826599911934e-10
+        ],
+        [
+            -3.465864794112,
+            -0.0157268418022,
+            0.01398192327776,
+            -0.002700181027443,
+            0.0001902304157269,
+            -1.534387905925e-06,
+            -1.03206007926e-07,
+            -1.355858638619e-08,
+            5.917279771473e-10
+        ],
+        [
+            0.7479071085372,
+            0.0005997666028811,
+            -0.001614053457119,
+            0.0002620866439317,
+            -1.103039382799e-05,
+            -4.215447554819e-07,
+            -9.926133276192e-09,
+            4.148813674084e-09,
+            -1.207867670158e-10
+        ],
+        [
+            -0.08863575102304,
+            0.0001901540166344,
+            6.941090299375e-05,
+            -3.042043168371e-06,
+            -1.677907209787e-06,
+            2.153652742395e-07,
+            -6.354480058307e-09,
+            -1.223103792568e-10,
+            5.54305794673e-12
+        ],
+        [
+            0.005484926807853,
+            -2.510359436743e-05,
+            1.123735445147e-06,
+            -9.198494797723e-07,
+            2.058851315121e-07,
+            -1.866416375894e-08,
+            7.564835556537e-10,
+            -1.622993472948e-11,
+            2.428986170198e-13
+        ],
+        [
+            -0.0001388441945179,
+            9.1419955967e-07,
+            -1.16891589033e-07,
+            3.485370731777e-08,
+            -5.086412415216e-09,
+            3.674153797642e-10,
+            -1.621809988343e-11,
+            6.737654534264e-13,
+            -1.678705755876e-14
+        ]
+    ]
+}

drbx/geometry/__init__.py

@@ -0,0 +1,207 @@
+from .essos_import import (
+    EssosFieldLineBundle,
+    EssosImportedFciGeometry,
+    build_essos_vmec_scaled_qa_coordinates,
+    build_essos_imported_fci_geometry,
+    essos_runtime_available,
+    load_essos_field_line_bundle_npz,
+    load_essos_coil_field_axis,
+    load_essos_vmec_field_axis,
+    resolve_essos_landreman_qa_json,
+    resolve_essos_landreman_qa_wout,
+    save_essos_field_line_bundle_npz,
+    trace_essos_coil_field_lines,
+    trace_essos_coil_initial_conditions,
+    trace_essos_vmec_initial_conditions,
+)
+from .fci_maps import FciMaps, identity_fci_maps, load_fci_maps_netcdf
+from .metric_tensor import MetricTensor3D, build_metric_report, metric_inverse_residual
+from .stellarator import SyntheticStellaratorGeometry, build_synthetic_stellarator_geometry
+from .island_divertor import (
+    IslandDivertorField,
+    build_island_divertor_geometry,
+    island_divertor_connection_length,
+    island_divertor_field_line_rhs,
+)
+from .open_slab import build_open_slab_geometry
+from .rotating_ellipse import build_rotating_ellipse_geometry, rotating_ellipse_position
+from .shifted_torus import build_shifted_torus_geometry
+from .vmec_jax_import import (
+    evaluate_vmec_jax_surface_field,
+    load_vmec_jax_wout,
+    trace_vmec_jax_field_lines,
+    traced_rotational_transform,
+    vmec_jax_boundary_rz,
+    vmec_jax_half_mesh_s,
+    vmec_jax_runtime_available,
+    vmec_jax_surface_rz,
+    vmec_jax_wout_summary,
+)
+from .vmec_extender_import import (
+    VmecExtenderGrid,
+    build_vmec_extender_fci_maps,
+    interpolate_vmec_extender_B_cyl,
+    load_vmec_extender_grid_netcdf,
+    validate_vmec_extender_points_in_bounds,
+    vmec_extender_absB,
+    vmec_extender_fieldline_rhs_RZ_phi,
+    vmec_extender_points_in_bounds,
+)
+
+__all__ = [
+    "EssosFieldLineBundle",
+    "EssosImportedFciGeometry",
+    "build_essos_vmec_scaled_qa_coordinates",
+    "build_essos_imported_fci_geometry",
+    "essos_runtime_available",
+    "FciMaps",
+    "MetricTensor3D",
+    "SyntheticStellaratorGeometry",
+    "VmecExtenderGrid",
+    "build_metric_report",
+    "IslandDivertorField",
+    "build_island_divertor_geometry",
+    "island_divertor_connection_length",
+    "island_divertor_field_line_rhs",
+    "build_open_slab_geometry",
+    "build_rotating_ellipse_geometry",
+    "rotating_ellipse_position",
+    "build_shifted_torus_geometry",
+    "build_synthetic_stellarator_geometry",
+    "build_vmec_extender_fci_maps",
+    "identity_fci_maps",
+    "interpolate_vmec_extender_B_cyl",
+    "load_essos_field_line_bundle_npz",
+    "load_essos_coil_field_axis",
+    "load_essos_vmec_field_axis",
+    "load_fci_maps_netcdf",
+    "load_vmec_extender_grid_netcdf",
+    "metric_inverse_residual",
+    "resolve_essos_landreman_qa_json",
+    "resolve_essos_landreman_qa_wout",
+    "save_essos_field_line_bundle_npz",
+    "trace_essos_coil_field_lines",
+    "trace_essos_coil_initial_conditions",
+    "trace_essos_vmec_initial_conditions",
+    "validate_vmec_extender_points_in_bounds",
+    "evaluate_vmec_jax_surface_field",
+    "load_vmec_jax_wout",
+    "trace_vmec_jax_field_lines",
+    "traced_rotational_transform",
+    "vmec_jax_boundary_rz",
+    "vmec_jax_half_mesh_s",
+    "vmec_jax_runtime_available",
+    "vmec_jax_surface_rz",
+    "vmec_jax_wout_summary",
+    "vmec_extender_absB",
+    "vmec_extender_fieldline_rhs_RZ_phi",
+    "vmec_extender_points_in_bounds",
+]
+
+# --- FCI stack incorporated from PR #3 (Aiken Xie, branch 3D_fci) ---
+from .fci_geometry import (
+    BFieldGeometry,
+    CellCenteredGrid3D,
+    CellVolumeGeometry3D,
+    LocalFciDirectionMap,
+    LocalCellVolumeGeometry3D,
+    FaceBFieldGeometry,
+    FaceMetricGeometry,
+    FciGeometry3D,
+    FciMaps3D,
+    LocalFciGeometry3D,
+    ConservativeStencilBuilder,
+    ConservativeStencilBuilderContext,
+    LocalConservativeStencilBuilder,
+    LocalStencilBuilder,
+    LocalStencilBuilderContext,
+    StencilBuilderContext,
+    Grid1D,
+    HaloLayout3D,
+    LocalCellCenteredGrid3D,
+    LocalGrid1D,
+    LocalBFieldGeometry,
+    LocalFciLocalDependencyTable,
+    LocalFciMaps3D,
+    LocalFciRemoteDependencyTable,
+    FCI_DEP_INVALID,
+    FCI_DEP_FIELD_INTERIOR,
+    FCI_DEP_PHYSICAL_BOUNDARY,
+    FCI_DEP_CUT_WALL,
+    LocalDomain3D,
+    LocalFaceBFieldGeometry,
+    LocalRegularFaceGeometry3D,
+    LocalFaceMetricGeometry,
+    MetricGeometry,
+    LocalMetricGeometry,
+    NeighborMap3D,
+    ShardSpec3D,
+    SIDE_PHYSICAL,
+    SIDE_SIMPLE_PERIODIC,
+    SIDE_AXIS_REGULAR,
+    SIDE_TOPOLOGY_MAPPED,
+    SIDE_UNUSED,
+    RegularFaceGeometry3D,
+    LocalSpacing3D,
+    Spacing3D,
+    build_fci_maps_from_b_contravariant,
+    build_curvature_coefficients,
+    build_conservative_stencil_from_field,
+    build_local_conservative_stencil_from_field,
+    build_local_direct_stencil_one_sided_physical_from_halo,
+    build_local_stencil_from_field,
+    logical_grid_from_axis_vectors,
+)
+__all__ += [
+    "BFieldGeometry",
+    "CellCenteredGrid3D",
+    "CellVolumeGeometry3D",
+    "ConservativeStencilBuilder",
+    "ConservativeStencilBuilderContext",
+    "FCI_DEP_CUT_WALL",
+    "FCI_DEP_FIELD_INTERIOR",
+    "FCI_DEP_INVALID",
+    "FCI_DEP_PHYSICAL_BOUNDARY",
+    "FaceBFieldGeometry",
+    "FaceMetricGeometry",
+    "FciGeometry3D",
+    "FciMaps3D",
+    "Grid1D",
+    "HaloLayout3D",
+    "LocalBFieldGeometry",
+    "LocalCellCenteredGrid3D",
+    "LocalCellVolumeGeometry3D",
+    "LocalConservativeStencilBuilder",
+    "LocalDomain3D",
+    "LocalFaceBFieldGeometry",
+    "LocalFaceMetricGeometry",
+    "LocalFciDirectionMap",
+    "LocalFciGeometry3D",
+    "LocalFciLocalDependencyTable",
+    "LocalFciMaps3D",
+    "LocalFciRemoteDependencyTable",
+    "LocalGrid1D",
+    "LocalMetricGeometry",
+    "LocalRegularFaceGeometry3D",
+    "LocalSpacing3D",
+    "LocalStencilBuilder",
+    "LocalStencilBuilderContext",
+    "MetricGeometry",
+    "NeighborMap3D",
+    "RegularFaceGeometry3D",
+    "SIDE_AXIS_REGULAR",
+    "SIDE_PHYSICAL",
+    "SIDE_SIMPLE_PERIODIC",
+    "SIDE_TOPOLOGY_MAPPED",
+    "SIDE_UNUSED",
+    "ShardSpec3D",
+    "Spacing3D",
+    "StencilBuilderContext",
+    "build_conservative_stencil_from_field",
+    "build_curvature_coefficients",
+    "build_fci_maps_from_b_contravariant",
+    "build_local_conservative_stencil_from_field",
+    "build_local_direct_stencil_one_sided_physical_from_halo",
+    "build_local_stencil_from_field",
+    "logical_grid_from_axis_vectors",
+]

drbx/geometry/embedding.py

@@ -0,0 +1,56 @@
+"""Metric tensors from an analytic embedding, by automatic differentiation.
+
+Given a map from logical coordinates ``u = (x, theta, zeta)`` to Cartesian
+space, the covariant metric is ``g_ij = d_i X . d_j X`` — computed here exactly
+with ``jax.jacfwd`` instead of by hand. Every analytic geometry in the package
+(rotating ellipse, island divertor) builds its :class:`MetricGeometry` through
+this one function, so a new geometry only has to supply its embedding.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import jax
+import jax.numpy as jnp
+
+from .fci_geometry import MetricGeometry
+
+__all__ = ["metric_from_position_fn"]
+
+
+def metric_from_position_fn(
+    position_fn: Callable[[jax.Array], jax.Array],
+    logical_grid: jax.Array,
+) -> MetricGeometry:
+    """Build the metric on ``logical_grid`` from ``position_fn(u) -> (X, Y, Z)``.
+
+    The covariant metric is the Gram matrix of the embedding Jacobian, the
+    contravariant metric its inverse, and the Jacobian ``sqrt(det g_cov)``.
+    """
+
+    jacobian_of_position = jax.jacfwd(position_fn)
+    points = logical_grid.reshape(-1, 3)
+    cartesian_jacobian = jax.vmap(jacobian_of_position)(points)
+    g_cov = jnp.einsum("pki,pkj->pij", cartesian_jacobian, cartesian_jacobian)
+    g_contra = jnp.linalg.inv(g_cov)
+    det_g_cov = jnp.linalg.det(g_cov)
+    location_shape = logical_grid.shape[:-1]
+    g_cov = g_cov.reshape(location_shape + (3, 3))
+    g_contra = g_contra.reshape(location_shape + (3, 3))
+    jacobian = jnp.sqrt(jnp.abs(det_g_cov)).reshape(location_shape)
+    return MetricGeometry(
+        J=jacobian,
+        g11=g_contra[..., 0, 0],
+        g22=g_contra[..., 1, 1],
+        g33=g_contra[..., 2, 2],
+        g12=g_contra[..., 0, 1],
+        g13=g_contra[..., 0, 2],
+        g23=g_contra[..., 1, 2],
+        g_11=g_cov[..., 0, 0],
+        g_22=g_cov[..., 1, 1],
+        g_33=g_cov[..., 2, 2],
+        g_12=g_cov[..., 0, 1],
+        g_13=g_cov[..., 0, 2],
+        g_23=g_cov[..., 1, 2],
+    )