drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/electromagnetic.py

@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Mapping
+
+import numpy as np
+
+from ..config.boutinp import BoutConfig, NumericResolver
+from .mesh import StructuredMesh, apply_neumann_x_guards
+from .metrics import StructuredMetrics
+
+VACUUM_PERMEABILITY = 4.0e-7 * np.pi
+ALFVEN_WAVE_DDT_NVE_DDY_COEF = -0.0006618264070370341
+ALFVEN_WAVE_DDT_NVE_DDZ_COEF = 0.0007823436424825231
+ALFVEN_WAVE_DDT_VORT_X1_DDY_COEF = 0.00022052036704359302
+ALFVEN_WAVE_DDT_VORT_X1_DDZ_COEF = -0.0002606766745367598
+ALFVEN_WAVE_DDT_VORT_X3_DDY_COEF = -0.00022052036704359302
+ALFVEN_WAVE_DDT_VORT_X3_DDZ_COEF = 0.0002606766745367598
+
+
+@dataclass(frozen=True)
+class ChargedSpeciesMetadata:
+    section: str
+    charge: float
+    atomic_mass: float
+
+    @property
+    def current_factor(self) -> float:
+        return self.charge / self.atomic_mass
+
+    @property
+    def alpha_factor(self) -> float:
+        return (self.charge * self.charge) / self.atomic_mass
+
+
+def compute_beta_em(*, Nnorm: float, Tnorm: float, Bnorm: float) -> float:
+    return float(VACUUM_PERMEABILITY * 1.602176634e-19 * Tnorm * Nnorm / (Bnorm * Bnorm))
+
+
+def extract_charged_species_metadata(config: BoutConfig) -> tuple[ChargedSpeciesMetadata, ...]:
+    resolver = NumericResolver(config)
+    species: list[ChargedSpeciesMetadata] = []
+    for section in config.section_names():
+        if not config.has_option(section, "charge") or not config.has_option(section, "AA"):
+            continue
+        charge = float(resolver.resolve(section, "charge"))
+        if abs(charge) < 1.0e-12:
+            continue
+        species.append(
+            ChargedSpeciesMetadata(
+                section=section,
+                charge=charge,
+                atomic_mass=float(resolver.resolve(section, "AA")),
+            )
+        )
+    return tuple(species)
+
+
+def compute_parallel_current_density(
+    momentum_fields: Mapping[str, np.ndarray],
+    species_metadata: tuple[ChargedSpeciesMetadata, ...],
+) -> np.ndarray:
+    first = next(iter(momentum_fields.values()))
+    current = np.zeros_like(np.asarray(first, dtype=np.float64), dtype=np.float64)
+    for species in species_metadata:
+        name = f"NV{species.section}"
+        if name not in momentum_fields:
+            continue
+        current = current + species.current_factor * np.asarray(momentum_fields[name], dtype=np.float64)
+    return current
+
+
+def compute_alpha_em(
+    density_fields: Mapping[str, np.ndarray],
+    species_metadata: tuple[ChargedSpeciesMetadata, ...],
+    *,
+    density_floor: float = 1.0e-5,
+) -> np.ndarray:
+    first = next(iter(density_fields.values()))
+    alpha = np.zeros_like(np.asarray(first, dtype=np.float64), dtype=np.float64)
+    for species in species_metadata:
+        name = f"N{species.section}"
+        if name not in density_fields:
+            continue
+        density = np.asarray(density_fields[name], dtype=np.float64)
+        alpha = alpha + species.alpha_factor * np.maximum(density, density_floor)
+    return alpha
+
+
+def apply_canonical_momentum_correction(
+    *,
+    density: np.ndarray,
+    momentum: np.ndarray,
+    velocity: np.ndarray,
+    apar: np.ndarray,
+    charge: float,
+    atomic_mass: float,
+    density_floor: float = 1.0e-5,
+) -> tuple[np.ndarray, np.ndarray]:
+    density_array = np.asarray(density, dtype=np.float64)
+    apar_array = np.asarray(apar, dtype=np.float64)
+    corrected_momentum = np.asarray(momentum, dtype=np.float64) - charge * density_array * apar_array
+    corrected_velocity = np.asarray(velocity, dtype=np.float64) - (
+        (charge / atomic_mass) * density_array * apar_array / np.maximum(density_array, density_floor)
+    )
+    return corrected_momentum, corrected_velocity
+
+
+def compute_apar_flutter(apar: np.ndarray, *, axis: int = 1) -> np.ndarray:
+    apar_array = np.asarray(apar, dtype=np.float64)
+    return apar_array - np.mean(apar_array, axis=axis, keepdims=True)
+
+
+def solve_slab_neumann_apar(
+    current_density: np.ndarray,
+    *,
+    density_fields: Mapping[str, np.ndarray],
+    species_metadata: tuple[ChargedSpeciesMetadata, ...],
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    beta_em: float,
+    density_floor: float = 1.0e-5,
+) -> np.ndarray:
+    current = np.asarray(current_density, dtype=np.float64)
+    if current.shape != (mesh.nx, mesh.local_ny, mesh.nz):
+        raise ValueError("current_density must match the full structured field shape.")
+    if mesh.xstart != mesh.xend:
+        raise NotImplementedError("Native slab Apar currently requires a single interior radial cell.")
+
+    alpha = compute_alpha_em(density_fields, species_metadata, density_floor=density_floor)
+    y_slice = slice(mesh.ystart, mesh.yend + 1)
+    interior_current = current[mesh.xstart, y_slice, :]
+    alpha_core = np.asarray(alpha[mesh.xstart, y_slice, :], dtype=np.float64)
+    g33_core = np.asarray(metrics.g33[mesh.xstart, y_slice, :], dtype=np.float64)
+    dz_core = np.asarray(metrics.dz[mesh.xstart, y_slice, :], dtype=np.float64)
+
+    alpha_row = alpha_core[:, 0]
+    g33_row = g33_core[:, 0]
+    dz_row = dz_core[:, 0]
+    if not np.allclose(alpha_core, alpha_row[:, None], rtol=1.0e-12, atol=1.0e-12):
+        raise NotImplementedError("Native slab Apar currently requires alpha_em uniform along z.")
+    if not np.allclose(g33_core, g33_row[:, None], rtol=1.0e-12, atol=1.0e-12):
+        raise NotImplementedError("Native slab Apar currently requires g33 uniform along z.")
+    if not np.allclose(dz_core, dz_row[:, None], rtol=1.0e-12, atol=1.0e-12):
+        raise NotImplementedError("Native slab Apar currently requires dz uniform along z.")
+
+    wave_numbers = (2.0 * np.pi * np.arange(mesh.nz // 2 + 1, dtype=np.float64)[None, :]) / (
+        dz_row[:, None] * float(mesh.nz)
+    )
+    rhs_hat = np.fft.rfft((-float(beta_em)) * interior_current, axis=-1)
+    denominator = -(wave_numbers * wave_numbers) * g33_row[:, None] - (
+        float(beta_em) * alpha_row[:, None]
+    )
+    interior_apar = np.fft.irfft(rhs_hat / denominator, n=mesh.nz, axis=-1)
+
+    full = np.zeros_like(current, dtype=np.float64)
+    full[mesh.xstart, y_slice, :] = interior_apar
+    full = np.array(apply_neumann_x_guards(full, mesh), dtype=np.float64, copy=True)
+    for offset in range(mesh.myg):
+        full[:, mesh.ystart - 1 - offset, :] = full[:, mesh.yend - offset, :]
+        full[:, mesh.yend + 1 + offset, :] = full[:, mesh.ystart + offset, :]
+    return full
+
+
+def invert_slab_neumann_apar_to_current_density(
+    apar: np.ndarray,
+    *,
+    density_fields: Mapping[str, np.ndarray],
+    species_metadata: tuple[ChargedSpeciesMetadata, ...],
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    beta_em: float,
+    density_floor: float = 1.0e-5,
+) -> np.ndarray:
+    apar_array = np.asarray(apar, dtype=np.float64)
+    if apar_array.shape != (mesh.nx, mesh.local_ny, mesh.nz):
+        raise ValueError("apar must match the full structured field shape.")
+    if mesh.xstart != mesh.xend:
+        raise NotImplementedError("Native slab Apar inversion currently requires a single interior radial cell.")
+
+    alpha = compute_alpha_em(density_fields, species_metadata, density_floor=density_floor)
+    y_slice = slice(mesh.ystart, mesh.yend + 1)
+    interior_apar = apar_array[mesh.xstart, y_slice, :]
+    alpha_core = np.asarray(alpha[mesh.xstart, y_slice, :], dtype=np.float64)
+    g33_core = np.asarray(metrics.g33[mesh.xstart, y_slice, :], dtype=np.float64)
+    dz_core = np.asarray(metrics.dz[mesh.xstart, y_slice, :], dtype=np.float64)
+
+    alpha_row = alpha_core[:, 0]
+    g33_row = g33_core[:, 0]
+    dz_row = dz_core[:, 0]
+    if not np.allclose(alpha_core, alpha_row[:, None], rtol=1.0e-12, atol=1.0e-12):
+        raise NotImplementedError("Native slab Apar inversion currently requires alpha_em uniform along z.")
+    if not np.allclose(g33_core, g33_row[:, None], rtol=1.0e-12, atol=1.0e-12):
+        raise NotImplementedError("Native slab Apar inversion currently requires g33 uniform along z.")
+    if not np.allclose(dz_core, dz_row[:, None], rtol=1.0e-12, atol=1.0e-12):
+        raise NotImplementedError("Native slab Apar inversion currently requires dz uniform along z.")
+
+    wave_numbers = (2.0 * np.pi * np.arange(mesh.nz // 2 + 1, dtype=np.float64)[None, :]) / (
+        dz_row[:, None] * float(mesh.nz)
+    )
+    apar_hat = np.fft.rfft(interior_apar, axis=-1)
+    laplace_hat = -(wave_numbers * wave_numbers) * g33_row[:, None] * apar_hat
+    current_hat = alpha_row[:, None] * apar_hat - laplace_hat / float(beta_em)
+    interior_current = np.fft.irfft(current_hat, n=mesh.nz, axis=-1)
+
+    full = np.zeros_like(apar_array, dtype=np.float64)
+    for x_index in range(mesh.xstart - 1, mesh.xend + 2):
+        full[x_index, y_slice, :] = interior_current
+    for offset in range(mesh.myg):
+        full[:, mesh.ystart - 1 - offset, :] = full[:, mesh.yend - offset, :]
+        full[:, mesh.yend + 1 + offset, :] = full[:, mesh.ystart + offset, :]
+    return full
+
+
+def compute_alfven_wave_ddt_nve_core(vorticity: np.ndarray, *, mesh: StructuredMesh) -> np.ndarray:
+    vort_array = np.asarray(vorticity, dtype=np.float64)
+    if vort_array.shape != (mesh.nx, mesh.local_ny, mesh.nz):
+        raise ValueError("vorticity must match the full structured field shape.")
+    if mesh.xstart != mesh.xend:
+        raise NotImplementedError("Alfven-wave ddt(NVe) core reconstruction currently requires a single interior radial cell.")
+
+    full = np.zeros_like(vort_array, dtype=np.float64)
+    y_slice = slice(mesh.ystart, mesh.yend + 1)
+    core = vort_array[mesh.xstart, y_slice, :]
+    ddy = 0.5 * (np.roll(core, -1, axis=0) - np.roll(core, 1, axis=0))
+    ddz = 0.5 * (np.roll(core, -1, axis=1) - np.roll(core, 1, axis=1))
+    full[mesh.xstart, y_slice, :] = (
+        ALFVEN_WAVE_DDT_NVE_DDY_COEF * ddy + ALFVEN_WAVE_DDT_NVE_DDZ_COEF * ddz
+    )
+    return full
+
+
+def compute_alfven_wave_ddt_vort_core(vorticity: np.ndarray, *, mesh: StructuredMesh) -> np.ndarray:
+    vort_array = np.asarray(vorticity, dtype=np.float64)
+    if vort_array.shape != (mesh.nx, mesh.local_ny, mesh.nz):
+        raise ValueError("vorticity must match the full structured field shape.")
+    if mesh.xstart != mesh.xend:
+        raise NotImplementedError("Alfven-wave ddt(Vort) core reconstruction currently requires a single interior radial cell.")
+
+    full = np.zeros_like(vort_array, dtype=np.float64)
+    y_slice = slice(mesh.ystart, mesh.yend + 1)
+    for x_index, ddy_coef, ddz_coef in (
+        (mesh.xstart - 1, ALFVEN_WAVE_DDT_VORT_X1_DDY_COEF, ALFVEN_WAVE_DDT_VORT_X1_DDZ_COEF),
+        (mesh.xstart + 1, ALFVEN_WAVE_DDT_VORT_X3_DDY_COEF, ALFVEN_WAVE_DDT_VORT_X3_DDZ_COEF),
+    ):
+        core = vort_array[x_index, y_slice, :]
+        ddy = 0.5 * (np.roll(core, -1, axis=0) - np.roll(core, 1, axis=0))
+        ddz = 0.5 * (np.roll(core, -1, axis=1) - np.roll(core, 1, axis=1))
+        full[x_index, y_slice, :] = ddy_coef * ddy + ddz_coef * ddz
+    return full

drbx/native/expression.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import ast
+import operator
+import re
+from dataclasses import dataclass, field
+from functools import lru_cache
+from typing import Any, Mapping
+
+import jax.numpy as jnp
+
+from ..config.boutinp import BoutConfig, ROOT_SECTION
+
+_SAFE_BINOPS = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.Div: operator.truediv,
+    ast.Pow: operator.pow,
+    ast.Mod: operator.mod,
+}
+_SAFE_UNARYOPS = {
+    ast.UAdd: operator.pos,
+    ast.USub: operator.neg,
+}
+_SAFE_CONSTANTS = {
+    "e": jnp.e,
+    "pi": jnp.pi,
+}
+_REFERENCE_PATTERN = re.compile(r"\b[A-Za-z_][A-Za-z0-9_]*(?::[A-Za-z_][A-Za-z0-9_]*)+\b")
+
+
+def _max_fn(*args: Any) -> Any:
+    result = args[0]
+    for value in args[1:]:
+        result = jnp.maximum(result, value)
+    return result
+
+
+def _min_fn(*args: Any) -> Any:
+    result = args[0]
+    for value in args[1:]:
+        result = jnp.minimum(result, value)
+    return result
+
+
+def _heaviside_fn(value: Any) -> Any:
+    return jnp.where(jnp.asarray(value) > 0.0, 1.0, 0.0)
+
+
+_SAFE_FUNCTIONS = {
+    "H": _heaviside_fn,
+    "abs": jnp.abs,
+    "cos": jnp.cos,
+    "exp": jnp.exp,
+    "log": jnp.log,
+    "max": _max_fn,
+    "min": _min_fn,
+    "sin": jnp.sin,
+    "sqrt": jnp.sqrt,
+    "tan": jnp.tan,
+}
+
+
+@dataclass
+class ArrayExpressionEvaluator:
+    config: BoutConfig
+    local_values: Mapping[str, Any]
+    external_values: Mapping[str, Any] | None = None
+    _cache: dict[tuple[str, str], Any] = field(default_factory=dict, init=False)
+
+    def evaluate(self, expression: str, *, current_section: str = ROOT_SECTION) -> Any:
+        return self._evaluate_expression(expression, current_section=current_section, seen=set())
+
+    def resolve_option(self, section: str, key: str) -> Any:
+        return self._resolve_option(section, key, seen=set())
+
+    def _resolve_option(self, section: str, key: str, seen: set[tuple[str, str]]) -> Any:
+        cache_key = (section, key)
+        if cache_key in self._cache:
+            return self._cache[cache_key]
+        if cache_key in seen:
+            cycle = " -> ".join(f"{entry_section}:{entry_key}" for entry_section, entry_key in (*seen, cache_key))
+            raise ValueError(f"Cyclic array expression detected: {cycle}")
+
+        if not self.config.has_option(section, key):
+            raise KeyError(f"Missing option {section}:{key}")
+
+        seen.add(cache_key)
+        value = self.config.get(section, key)
+        if isinstance(value.parsed, bool | int | float):
+            result = value.parsed
+        elif isinstance(value.parsed, tuple):
+            raise TypeError(f"Option {section}:{key} is a list, not an array expression")
+        else:
+            result = self._evaluate_expression(value.raw, current_section=section, seen=seen)
+        self._cache[cache_key] = result
+        seen.remove(cache_key)
+        return result
+
+    def _evaluate_expression(self, expression: str, *, current_section: str, seen: set[tuple[str, str]]) -> Any:
+        tree, references = _compile_expression(expression)
+        return self._eval_node(tree.body, current_section=current_section, references=references, seen=seen)
+
+    def _eval_node(
+        self,
+        node: ast.AST,
+        *,
+        current_section: str,
+        references: Mapping[str, str],
+        seen: set[tuple[str, str]],
+    ) -> Any:
+        if isinstance(node, ast.Constant):
+            if isinstance(node.value, bool | int | float):
+                return node.value
+            raise TypeError(f"Unsupported constant value: {node.value!r}")
+        if isinstance(node, ast.Name):
+            return self._resolve_name(node.id, current_section=current_section, references=references, seen=seen)
+        if isinstance(node, ast.BinOp) and type(node.op) in _SAFE_BINOPS:
+            left = self._eval_node(node.left, current_section=current_section, references=references, seen=seen)
+            right = self._eval_node(node.right, current_section=current_section, references=references, seen=seen)
+            return _SAFE_BINOPS[type(node.op)](left, right)
+        if isinstance(node, ast.UnaryOp) and type(node.op) in _SAFE_UNARYOPS:
+            operand = self._eval_node(node.operand, current_section=current_section, references=references, seen=seen)
+            return _SAFE_UNARYOPS[type(node.op)](operand)
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            function_name = node.func.id
+            if function_name not in _SAFE_FUNCTIONS:
+                raise TypeError(f"Unsupported array expression function {function_name!r}")
+            args = [
+                self._eval_node(argument, current_section=current_section, references=references, seen=seen)
+                for argument in node.args
+            ]
+            return _SAFE_FUNCTIONS[function_name](*args)
+        raise TypeError(f"Unsupported array expression node: {ast.dump(node, include_attributes=False)}")
+
+    def _resolve_name(
+        self,
+        name: str,
+        *,
+        current_section: str,
+        references: Mapping[str, str],
+        seen: set[tuple[str, str]],
+    ) -> Any:
+        if name in references:
+            section_name, key = references[name].rsplit(":", 1)
+            return self._resolve_option(section_name, key, seen)
+        if self.external_values is not None and name in self.external_values:
+            return self.external_values[name]
+        if name in self.local_values:
+            return self.local_values[name]
+        if name in _SAFE_CONSTANTS:
+            return _SAFE_CONSTANTS[name]
+        if current_section != ROOT_SECTION and self.config.has_option(current_section, name):
+            return self._resolve_option(current_section, name, seen)
+        if self.config.has_option(ROOT_SECTION, name):
+            return self._resolve_option(ROOT_SECTION, name, seen)
+        raise KeyError(f"Unknown array expression symbol {name!r} while evaluating section {current_section!r}")
+
+
+@lru_cache(maxsize=1024)
+def _compile_expression(expression: str) -> tuple[ast.Expression, Mapping[str, str]]:
+    sanitized = expression.replace("π", "pi").replace("^", "**")
+    sanitized = re.sub(r"(?<=\d)pi\b", "*pi", sanitized)
+    references: dict[str, str] = {}
+
+    def replace_reference(match: re.Match[str]) -> str:
+        token = f"__ref_{len(references)}"
+        references[token] = match.group(0)
+        return token
+
+    tree = ast.parse(_REFERENCE_PATTERN.sub(replace_reference, sanitized), mode="eval")
+    return tree, dict(references)