Source code for xgc_analysis.distribution_moments

from __future__ import annotations

from dataclasses import dataclass
import warnings

import numpy as np
from .constants import echarge



@dataclass
class SplitMoment:
    """Linear moment split into adiabatic / non-adiabatic pieces."""

    adiabatic: np.ndarray | None = None
    nonadiabatic: np.ndarray | None = None

    @property
    def total(self):
        if self.adiabatic is None and self.nonadiabatic is None:
            return None
        if self.adiabatic is None:
            return self.nonadiabatic
        if self.nonadiabatic is None:
            return self.adiabatic
        return self.adiabatic + self.nonadiabatic





@dataclass
class FluxSplit:
    """Flux moment split by physical mechanism."""

    drift: SplitMoment
    exb: SplitMoment
    delta_b: SplitMoment | None = None





@dataclass
class F0MomentComponent:
    """
    Moment set for one component branch (``n0`` or ``turb``).

    Notes
    -----
    Basic moments (density, flow, temperatures) are stored as plain arrays for
    this branch and are *not* split into adiabatic/non-adiabatic parts.
    Flux-like moments keep the adiabatic/non-adiabatic split.
    """

    density: np.ndarray
    mean_parallel_flow: np.ndarray
    parallel_temperature: np.ndarray
    perpendicular_temperature: np.ndarray
    parallel_energy_flux: SplitMoment
    radial_particle_flux: FluxSplit
    radial_energy_flux: FluxSplit





@dataclass
class F0MomentResult:
    """
    Common output structure for axisymmetric and non-axisymmetric runs.

    ``n0`` is always present. ``turb`` is present only for non-axisymmetric
    simulations when ``split_n0_turb`` is enabled.
    """

    n0: F0MomentComponent
    turb: F0MomentComponent | None = None





@dataclass
class F0MomentInputs:
    """
    Geometry / field inputs required by electrostatic f0 moments (f0_sum-style).

    Arrays are expected in axisymmetric-compatible shapes:
      - node arrays: ``(n_node,)`` or ``(1, n_node)``
      - vector arrays: ``(n_node, n_comp)`` or ``(1, n_node, n_comp)``

    Parameters
    ----------
    mass_kg, charge_C : float
        Species mass and charge.
    temp_node_ev : ndarray
        Reference local temperature used for velocity normalization (``temp_node``
        in the legacy/C++ diagnostics path, typically ``f0_fg_T_ev`` for the
        species).
    den_node_m3 : ndarray
        Reference local density (``f0_den`` for the species).
    bfield : ndarray
        Magnetic field components in cylindrical ordering ``(R, Z, phi)``.
    epsi, etheta : ndarray
        Electric-field components used by the electrostatic ExB velocity terms.
    gradpsi : ndarray
        Gradient of poloidal flux with shape ``(..., 2)`` (R,Z components).
    nb_curl_nb : ndarray
        Scalar ``b x curl(b)`` factor used in the denominator ``D``.
    v_gradb, v_curv : ndarray
        Drift vectors with the legacy XGC ordering; component 0 is the radial-like
        (flux) component, component 1 is Z, component 2 is toroidal.
    dpot : ndarray | None
        Electrostatic potential perturbation (V) for the Boltzmann term.
    flux_surface_avg_project : callable | None
        Optional callback mapping a node-array to its flux-surface-averaged
        projection back on the nodes (Fortran ``mat_transpose_mult`` +
        interpolation path analogue). Signature: ``arr_node -> arr_node``.
    f_3d_re, f_3d_im : ndarray | None
        Optional real/imaginary non-axisymmetric distribution components for
        delta-B radial flux moments. Same shape as the primary distribution.
    dbre, dbim : ndarray | None
        Optional real/imaginary magnetic perturbation vectors (R,Z,phi comps).
    species_index : int | None
        Species index used for species-dependent velocity-volume normalization
        (``f0_grid_vol_vonly``) when available.
    split_n0_turb : bool, default False
        Enable decomposition into toroidally averaged (``n0``) and
        non-axisymmetric (``turb``) branches.
    sim_is_axisymmetric : bool | None
        Simulation-level symmetry flag (distinct from mesh symmetry). Used to
        choose default branch behavior in :meth:`from_simulation`.

    Notes
    -----
    Moment definitions follow the standard velocity integrals used in XGC
    analysis:
    ``n = ∫ f d^3v``, ``u_parallel = (1/n) ∫ v_parallel f d^3v``,
    ``T_parallel = (m/en) ∫ (v_parallel-u_parallel)^2 f d^3v``,
    ``T_perp = (m/2en) ∫ v_perp^2 f d^3v``.
    """

    mass_kg: float
    charge_C: float
    temp_node_ev: np.ndarray
    den_node_m3: np.ndarray
    bfield: np.ndarray
    epsi: np.ndarray
    etheta: np.ndarray
    gradpsi: np.ndarray
    nb_curl_nb: np.ndarray
    v_gradb: np.ndarray
    v_curv: np.ndarray
    dpot: np.ndarray | None = None
    flux_surface_avg_project: callable | None = None
    f_3d_re: np.ndarray | None = None
    f_3d_im: np.ndarray | None = None
    dbre: np.ndarray | None = None
    dbim: np.ndarray | None = None
    species_index: int | None = None
    split_n0_turb: bool = False
    sim_is_axisymmetric: bool | None = None



    @staticmethod
    def make_flux_surface_avg_projector(geometry):
        """
        Build a ``flux_surface_avg_project`` callback from a ``Mesh``.

        The returned callable performs the same conceptual operation as the
        legacy ``mat_transpose_mult + interpolation`` path in ``f0_sum.F90``:
        flux-surface average followed by projection back to nodal values.

        For non-axisymmetric meshes, the mesh-level implementation performs
        per-plane FSA, toroidal averaging, and nodal back-projection. This is
        the required path for 3D XGC diagnostics and keeps axisymmetric/3D
        behavior unified.

        Supported input shapes to the callback
        --------------------------------------
        - ``(n_node,)``                 -> ``(n_node,)``
        - ``(1, n_node)``               -> ``(1, n_node)``
        - ``(nphi, n_node)``            -> ``(nphi, n_node)``
        """
        if not hasattr(geometry, "flux_avg_to_from_surf") or not hasattr(geometry, "nphi"):
            raise TypeError(
                "flux_surface_avg_geometry must be a Mesh-like object "
                "with 'nphi' and flux_avg_to_from_surf(...)."
            )

        def _project(arr):
            a = np.asarray(arr)
            if a.ndim == 1:
                return np.asarray(geometry.flux_avg_to_from_surf(a))
            if a.ndim == 2:
                # Axisymmetric path often carries a canonical singleton phi axis
                # (1, n_node). Mesh.flux_avg_to_from_surf accepts 1D nodal
                # arrays in this case, so strip and restore that axis.
                if a.shape[0] == 1:
                    out = np.asarray(geometry.flux_avg_to_from_surf(a[0]))
                    return out[np.newaxis, :] if out.ndim == 1 else out

                out = np.asarray(geometry.flux_avg_to_from_surf(a))
                # Mesh.flux_avg_to_from_surf currently returns a representative
                # nodal field (1D) after toroidal averaging. For moment
                # calculations on (nphi, nnode), restore canonical shape by
                # repeating the projected nodal profile across phi.
                if out.ndim == 1:
                    return np.broadcast_to(out[np.newaxis, :], a.shape)
                if out.shape == a.shape:
                    return out
                if out.shape == (1, a.shape[1]):
                    return np.broadcast_to(out, a.shape)
                raise ValueError(
                    "flux_surface_avg_project returned unsupported shape "
                    f"{out.shape} for input shape {a.shape}."
                )
            raise ValueError(
                "flux_surface_avg_project expects a node array of shape "
                "(n_node,), (1,n_node), or (nphi,n_node)."
            )

        return _project




    @classmethod
    def from_species(
        cls,
        *,
        species,
        species_index=None,
        temp_node_ev,
        den_node_m3,
        bfield,
        epsi,
        etheta,
        gradpsi,
        nb_curl_nb,
        v_gradb,
        v_curv,
        dpot=None,
        flux_surface_avg_project=None,
        f_3d_re=None,
        f_3d_im=None,
        dbre=None,
        dbim=None,
        split_n0_turb=False,
        sim_is_axisymmetric=None,
    ):
        return cls(
            mass_kg=float(species.mass_kg),
            charge_C=float(species.charge_C),
            temp_node_ev=np.asarray(temp_node_ev),
            den_node_m3=np.asarray(den_node_m3),
            bfield=np.asarray(bfield),
            epsi=np.asarray(epsi),
            etheta=np.asarray(etheta),
            gradpsi=np.asarray(gradpsi),
            nb_curl_nb=np.asarray(nb_curl_nb),
            v_gradb=np.asarray(v_gradb),
            v_curv=np.asarray(v_curv),
            dpot=None if dpot is None else np.asarray(dpot),
            flux_surface_avg_project=flux_surface_avg_project,
            f_3d_re=None if f_3d_re is None else np.asarray(f_3d_re),
            f_3d_im=None if f_3d_im is None else np.asarray(f_3d_im),
            dbre=None if dbre is None else np.asarray(dbre),
            dbim=None if dbim is None else np.asarray(dbim),
            species_index=None if species_index is None else int(species_index),
            split_n0_turb=bool(split_n0_turb),
            sim_is_axisymmetric=None if sim_is_axisymmetric is None else bool(sim_is_axisymmetric),
        )




    @classmethod
    def from_simulation(
        cls,
        *,
        simulation,
        species_index: int,
        epsi,
        etheta,
        dpot=None,
        flux_surface_avg_geometry=None,
        f_3d_re=None,
        f_3d_im=None,
        dbre=None,
        dbim=None,
        bfield=None,
        gradpsi=None,
        nb_curl_nb=None,
        v_gradb=None,
        v_curv=None,
        temp_node_ev=None,
        den_node_m3=None,
        split_n0_turb=None,
        sim_is_axisymmetric=None,
    ):
        """
        Convenience constructor using ``Simulation`` and ``VelocityGrid`` metadata.

        This helper auto-populates most inputs needed by the f0-moment calculator:
        species mass/charge from ``simulation.species[species_index]``, local
        reference profiles from ``simulation.velocity_grid``, and the equilibrium
        magnetic field from ``simulation.magnetic_field``.

        ``epsi`` and ``etheta`` are still required explicit inputs because they are
        time-dependent field quantities.

        Parameters
        ----------
        simulation : Simulation
            XGC-Analysis simulation object with ``species``, ``velocity_grid``,
            ``mesh``, and ``magnetic_field`` attributes.
        species_index : int
            Index into ``simulation.species`` and species-indexed arrays in
            ``simulation.velocity_grid``.
        flux_surface_avg_geometry : Mesh | None, optional
            Mesh object providing ``flux_avg_to_from_surf``. If omitted,
            ``simulation.mesh`` is used.
        split_n0_turb : bool | None, optional
            If ``None``, defaults to ``not simulation.sim_is_axisymmetric``.
        sim_is_axisymmetric : bool | None, optional
            Explicit simulation-level symmetry override. This is intentionally
            separate from ``mesh.is_axisymmetric`` because tokamak turbulence
            runs can have an axisymmetric mesh but non-axisymmetric dynamics.
        *overrides*
            Any of ``bfield``, ``gradpsi``, ``nb_curl_nb``, ``v_gradb``,
            ``v_curv``, ``temp_node_ev``, or ``den_node_m3`` may be supplied to
            override values taken from ``simulation``.
        """
        sp = simulation.species[species_index]
        vg = getattr(simulation, "velocity_grid", None)
        if vg is None:
            raise ValueError("simulation.velocity_grid is not available.")

        if temp_node_ev is None:
            if hasattr(vg, "fg_T_ev") and vg.fg_T_ev is not None:
                temp_node_ev = vg.fg_T_ev[species_index]
            else:
                temp_node_ev = vg.T_ev[species_index]
        if den_node_m3 is None:
            den_node_m3 = vg.den[species_index]

        if bfield is None:
            mf = getattr(simulation, "magnetic_field", None)
            if mf is None or not hasattr(mf, "bfield"):
                raise ValueError("simulation.magnetic_field.bfield is not available.")
            bfield = np.asarray(mf.bfield.get_data())

        if gradpsi is None:
            gradpsi = getattr(vg, "gradpsi", None)
        if nb_curl_nb is None:
            nb_curl_nb = getattr(vg, "nb_curl_nb", None)
        if v_gradb is None:
            v_gradb = getattr(vg, "v_gradb", None)
        if v_curv is None:
            v_curv = getattr(vg, "v_curv", None)

        missing = [
            name for name, val in (
                ("gradpsi", gradpsi),
                ("nb_curl_nb", nb_curl_nb),
                ("v_gradb", v_gradb),
                ("v_curv", v_curv),
            ) if val is None
        ]
        if missing:
            raise ValueError(
                "Missing required f0-mesh geometry arrays from simulation.velocity_grid "
                f"(or overrides): {', '.join(missing)}"
            )

        if flux_surface_avg_geometry is None:
            flux_surface_avg_geometry = getattr(simulation, "mesh", None)

        projector = None
        if flux_surface_avg_geometry is not None:
            if not hasattr(flux_surface_avg_geometry, "nphi") or not hasattr(
                flux_surface_avg_geometry, "flux_avg_to_from_surf"
            ):
                raise TypeError(
                    "flux_surface_avg_geometry must be a Mesh-like object "
                    "(pass simulation.mesh)."
                )
            projector = cls.make_flux_surface_avg_projector(flux_surface_avg_geometry)

        # Gate n0/turb split by simulation symmetry (not mesh geometry).
        if sim_is_axisymmetric is None:
            sim_is_axisymmetric = getattr(simulation, "sim_is_axisymmetric", None)
        if sim_is_axisymmetric is None:
            raise ValueError(
                "simulation.sim_is_axisymmetric is required to choose the "
                "axisymmetric vs non-axisymmetric moment split path."
            )
        if split_n0_turb is None:
            split_n0_turb = not bool(sim_is_axisymmetric)

        return cls.from_species(
            species=sp,
            species_index=species_index,
            temp_node_ev=temp_node_ev,
            den_node_m3=den_node_m3,
            bfield=bfield,
            epsi=epsi,
            etheta=etheta,
            gradpsi=gradpsi,
            nb_curl_nb=nb_curl_nb,
            v_gradb=v_gradb,
            v_curv=v_curv,
            dpot=dpot,
            flux_surface_avg_project=projector,
            f_3d_re=f_3d_re,
            f_3d_im=f_3d_im,
            dbre=dbre,
            dbim=dbim,
            split_n0_turb=split_n0_turb,
            sim_is_axisymmetric=sim_is_axisymmetric,
        )






class F0MomentCalculator:
    """
    Electrostatic f0 moment calculator (first-pass Python port of ``f0_sum.F90``).

    Notes
    -----
    - Assumes distribution data are stored in canonical order
      ``(phi, node, vpar, vperp)``.
    - Adiabatic/non-adiabatic split follows the XGC diagnostics strategy:
      compute local density/temperature from ``f``, apply mesh-level
      ``flux_avg_to_from_surf`` to obtain background fields, build adiabatic
      Maxwellian + Boltzmann response, then split moments.
    - For non-axisymmetric simulations, an additional branch decomposition is
      available: ``n0`` (toroidally averaged) and ``turb`` (deviation from
      ``n0``). This is controlled by ``F0MomentInputs.split_n0_turb``.
    """

    def __init__(self, velocity_grid):
        self.vgrid = velocity_grid

    # ------------------------------------------------------------------
    # Normalization helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _ensure_node_shape(arr, nphi, nnode):
        a = np.asarray(arr)
        if a.shape == (nnode,):
            return a[np.newaxis, :]
        if a.shape == (1, nnode):
            return a
        if a.shape == (nphi, nnode):
            return a
        raise ValueError(f"Expected node array shape {(nnode,)}, (1,{nnode}), or {(nphi,nnode)}; got {a.shape}")

    @staticmethod
    def _ensure_vec3_shape(arr, nphi, nnode):
        a = np.asarray(arr)
        if a.shape == (nnode, 3):
            return a[np.newaxis, :, :]
        if a.shape == (1, nnode, 3):
            return a
        if a.shape == (nphi, nnode, 3):
            return a
        raise ValueError(f"Expected vector-3 shape (nnode,3), (1,nnode,3), or (nphi,nnode,3); got {a.shape}")

    @staticmethod
    def _ensure_vec2_shape(arr, nphi, nnode):
        a = np.asarray(arr)
        if a.shape == (nnode, 2):
            return a[np.newaxis, :, :]
        if a.shape == (1, nnode, 2):
            return a
        if a.shape == (nphi, nnode, 2):
            return a
        raise ValueError(f"Expected vector-2 shape (nnode,2), (1,nnode,2), or (nphi,nnode,2); got {a.shape}")

    def _integrate(self, f, kernel=None, *, species_index=None):
        arr = np.asarray(f)
        if kernel is not None:
            arr = arr * np.asarray(kernel)
        return self.vgrid.integrate_over_velocity(
            arr,
            axis_mu=3,
            axis_vp=2,
            node_axis=1,
            data_includes_vperp=True,
            include_gyroangle=True,
            apply_grid_vol_vonly=True,
            species_index=species_index,
        )

    def _velocity_mesh(self, nphi, nnode):
        # Canonical distribution order is (phi, node, vpar, vperp)
        vpar = self.vgrid.vpara_norm[np.newaxis, np.newaxis, :, np.newaxis]
        vperp = self.vgrid.vperp_norm[np.newaxis, np.newaxis, np.newaxis, :]
        # Broadcast shape is implicit; returned arrays are lightweight.
        return vpar, vperp

    def _thermal_speed(self, temp_ev, mass_kg, charge_C):
        return np.sqrt(np.abs(charge_C) * temp_ev / mass_kg)

    # ------------------------------------------------------------------
    # Basic moments used to build the adiabatic split
    # ------------------------------------------------------------------
    def _density(self, f, *, species_index=None):
        return self._integrate(f, species_index=species_index)

    def _mean_parallel_flow(self, f, density, vth, *, species_index=None):
        vpar, _ = self._velocity_mesh(*f.shape[:2])
        mom = self._integrate(f, kernel=vpar, species_index=species_index)
        with np.errstate(invalid="ignore", divide="ignore"):
            upar_norm = np.where(density != 0.0, mom / density, 0.0)
        return upar_norm * vth

    def _temperature_components_ev(self, f, density, upar_ms, vth, mass_kg, charge_C, *, species_index=None):
        vpar, vperp = self._velocity_mesh(*f.shape[:2])
        upar_norm = np.where(vth != 0.0, upar_ms / vth, 0.0)[:, :, np.newaxis, np.newaxis]

        kinetic_fac_ev = mass_kg * (vth[:, :, np.newaxis, np.newaxis] ** 2) / (2.0 * np.abs(charge_C))
        # temp_par_out and temp_perp_out follow the legacy Fortran normalization.
        t_par_num = self._integrate(f, kernel=kinetic_fac_ev * (vpar - upar_norm) ** 2, species_index=species_index)
        t_perp_num = self._integrate(f, kernel=kinetic_fac_ev * (vperp ** 2), species_index=species_index)
        t_tot_num = self._integrate(
            f,
            kernel=kinetic_fac_ev * ((vpar - upar_norm) ** 2 + vperp ** 2),
            species_index=species_index,
        )

        with np.errstate(invalid="ignore", divide="ignore"):
            t_par = np.where(density != 0.0, 2.0 * t_par_num / density, 0.0)
            t_perp = np.where(density != 0.0, t_perp_num / density, 0.0)
            t = np.where(density != 0.0, 2.0 * t_tot_num / (3.0 * density), 0.0)
        return t, t_par, t_perp

    def _fsavg_project(self, arr, inputs: F0MomentInputs):
        if inputs.flux_surface_avg_project is None:
            warnings.warn(
                "No flux_surface_avg_project provided; using local values instead of "
                "f0_sum-style flux-surface-averaged background for the adiabatic split.",
                RuntimeWarning,
                stacklevel=2,
            )
            return np.asarray(arr)
        return np.asarray(inputs.flux_surface_avg_project(np.asarray(arr)))

    def _build_adiabatic_distribution(self, f, inputs: F0MomentInputs):
        nphi, nnode, _, _ = f.shape
        # In XGC diagnostics, this is the reference temperature used for
        # velocity normalization (fg_temp_ev), not the local kinetic temperature.
        fg_temp_node = self._ensure_node_shape(inputs.temp_node_ev, nphi, nnode)
        den_node = self._ensure_node_shape(inputs.den_node_m3, nphi, nnode)
        dpot = np.zeros_like(fg_temp_node) if inputs.dpot is None else self._ensure_node_shape(inputs.dpot, nphi, nnode)

        mass = float(inputs.mass_kg)
        charge = float(inputs.charge_C)  # Coulomb

        vth = self._thermal_speed(fg_temp_node, mass, charge)

        sp_index = inputs.species_index
        den = self._density(f, species_index=sp_index)
        upar = self._mean_parallel_flow(f, den, vth, species_index=sp_index)
        temp, _, _ = self._temperature_components_ev(
            f, den, upar, vth, mass, charge, species_index=sp_index
        )

        den_avg = self._fsavg_project(den, inputs)
        temp_avg = self._fsavg_project(temp, inputs)

        vpar, vperp = self._velocity_mesh(nphi, nnode)
        # Match diag_3d_f0_f.cpp adiabatic prefactor (axisymmetric n0 component):
        # f_adia_raw ~ sqrt(Tnorm^-1) * (n0/T0) * exp(-Tnorm^-1*en) * smu * boltz_fac
        # with Tnorm^-1 = fg_temp_ev / T0.
        tnorm_inv = np.where(temp_avg > 0.0, fg_temp_node / temp_avg, 0.0)[:, :, np.newaxis, np.newaxis]
        eps_norm = 0.5 * (vpar**2 + vperp**2)
        # boltz_fac = exp(- q[e] * dphi[V] / T0[eV]) in C++ diagnostics.
        charge_eu = charge / echarge
        boltz_arg = np.where(temp_avg > 0.0, -charge_eu * dpot / temp_avg, 0.0)[:, :, np.newaxis, np.newaxis]
        boltz_fac = np.exp(boltz_arg)

        prefac = (
            np.sqrt(np.where(temp_avg > 0.0, fg_temp_node / temp_avg, 0.0))
            * np.where(temp_avg > 0.0, den_avg / temp_avg, 0.0)
        )[:, :, np.newaxis, np.newaxis]
        # vperp is smu on this grid, including the mu0 correction at index 0.
        f_adi = prefac * np.exp(-tnorm_inv * eps_norm) * vperp * boltz_fac
        return f_adi, den, upar, temp, den_avg, temp_avg, vth

    # ------------------------------------------------------------------
    # Drift / ExB / deltaB kernels (f0_sum-style)
    # ------------------------------------------------------------------
    def _radial_kernels(self, inputs: F0MomentInputs, vth, den_avg, temp_avg):
        # den_avg and temp_avg are currently unused directly here, but retained to
        # mirror the f0_sum flow and make it easy to extend exact parity later.
        nphi, nnode = vth.shape
        bfield = self._ensure_vec3_shape(inputs.bfield, nphi, nnode)
        gradpsi = self._ensure_vec2_shape(inputs.gradpsi, nphi, nnode)
        v_gradb = self._ensure_vec3_shape(inputs.v_gradb, nphi, nnode)
        v_curv = self._ensure_vec3_shape(inputs.v_curv, nphi, nnode)
        etheta = self._ensure_node_shape(inputs.etheta, nphi, nnode)
        epsi = self._ensure_node_shape(inputs.epsi, nphi, nnode)
        nb_curl_nb = self._ensure_node_shape(inputs.nb_curl_nb, nphi, nnode)

        btot = np.linalg.norm(bfield, axis=-1)
        bphi = bfield[..., 2]
        gpsi = np.linalg.norm(gradpsi, axis=-1)
        with np.errstate(invalid="ignore", divide="ignore"):
            btheta = (-bfield[..., 0] * gradpsi[..., 1] + bfield[..., 1] * gradpsi[..., 0]) / gpsi

        # Flux-surface averaged radial drift coefficients (Fortran averages
        # v_gradb(1)/B^3 and v_curv(1)/B^2, then restores powers of B).
        vg_rad_norm = np.where(btot != 0.0, v_gradb[..., 0] / btot**3, 0.0)
        vc_rad_norm = np.where(btot != 0.0, v_curv[..., 0] / btot**2, 0.0)
        vg_rad_avg = self._fsavg_project(vg_rad_norm, inputs) * btot**3
        vc_rad_avg = self._fsavg_project(vc_rad_norm, inputs) * btot**2

        mass = float(inputs.mass_kg)
        charge = float(inputs.charge_C)
        vpar, vperp = self._velocity_mesh(nphi, nnode)

        btot4 = btot[:, :, np.newaxis, np.newaxis]
        vth4 = vth[:, :, np.newaxis, np.newaxis]
        rho = np.where(btot4 != 0.0, vpar * vth4 * (mass / charge) / btot4, 0.0)
        cmrho = (charge / mass) * rho
        cmrho2 = (charge / mass) * rho**2
        mu = np.where(btot4 != 0.0, mass * (vperp**2) * (vth4**2) / (2.0 * btot4), 0.0)

        D = 1.0 / (1.0 + rho * nb_curl_nb[:, :, np.newaxis, np.newaxis])

        v_exb_r = np.where(
            btot != 0.0,
            -etheta * bphi / (btot**2) * gpsi,
            0.0,
        )[:, :, np.newaxis, np.newaxis]

        murho2b_c = (mu + charge * cmrho2 * btot4) / charge
        drift_kernel = D * (
            ((v_gradb[..., 0] - vg_rad_avg)[:, :, np.newaxis, np.newaxis] * murho2b_c / (btot4**2))
            + cmrho2 * (v_curv[..., 0] - vc_rad_avg)[:, :, np.newaxis, np.newaxis]
        )
        exb_kernel = D * v_exb_r

        delta_b_re_kernel = None
        delta_b_im_kernel = None
        if inputs.dbre is not None and inputs.dbim is not None:
            dbre = self._ensure_vec3_shape(inputs.dbre, nphi, nnode)
            dbim = self._ensure_vec3_shape(inputs.dbim, nphi, nnode)
            gradpsi_r = gradpsi[..., 0][:, :, np.newaxis, np.newaxis]
            gradpsi_z = gradpsi[..., 1][:, :, np.newaxis, np.newaxis]
            delta_b_re_kernel = D * cmrho * (
                dbre[..., 0][:, :, np.newaxis, np.newaxis] * gradpsi_r
                + dbre[..., 1][:, :, np.newaxis, np.newaxis] * gradpsi_z
            )
            delta_b_im_kernel = D * cmrho * (
                dbim[..., 0][:, :, np.newaxis, np.newaxis] * gradpsi_r
                + dbim[..., 1][:, :, np.newaxis, np.newaxis] * gradpsi_z
            )

        return drift_kernel, exb_kernel, delta_b_re_kernel, delta_b_im_kernel

    @staticmethod
    def _toroidal_n0(arr):
        """Return toroidal-average (n0) component broadcast back to full phi shape."""
        a = np.asarray(arr)
        if a.ndim < 1:
            return a
        if a.shape[0] <= 1:
            return a
        mean0 = np.mean(a, axis=0, keepdims=True)
        return np.broadcast_to(mean0, a.shape)

    def _split_kernel_n0_turb(self, kernel):
        """
        Split a 4D velocity kernel into toroidally averaged (n0) and turbulent parts.

        Kernel shape is expected to follow the canonical distribution layout:
        ``(phi, node, vpar, vperp)``.
        """
        k = np.asarray(kernel)
        if k.ndim != 4:
            raise ValueError(f"Expected 4D kernel (phi,node,vpar,vperp), got shape {k.shape}.")
        k_n0 = self._toroidal_n0(k)
        return k_n0, (k - k_n0)



    def compute(self, distribution_field, inputs: F0MomentInputs) -> F0MomentResult:
        """
        Compute a first set of electrostatic f0 moments (f0_sum-style split).

        Implemented moments (semantic names)
        ------------------------------------
        - density
        - mean parallel flow
        - parallel energy flux (new; linear moment)
        - radial particle flux split into ``drift``, ``exb``, optional ``delta_b``
        - radial energy flux split into ``drift``, ``exb``, optional ``delta_b``

        Returns
        -------
        F0MomentResult
            ``result.n0`` always populated. ``result.turb`` populated only when
            ``inputs.split_n0_turb`` is true and the input distribution has more
            than one toroidal plane.

        Notes
        -----
        - Basic moments (density, flow, temperatures) are reported as branch
          totals (``n0`` and optional ``turb``), not adiabatic/non-adiabatic
          sub-splits.
        - Flux-like moments retain adiabatic/non-adiabatic decomposition.
        """
        f = np.asarray(distribution_field.get_data())
        if f.ndim != 4:
            raise ValueError("Distribution data must have canonical shape (phi,node,vpar,vperp).")

        f_adi, _, _, _, den_avg, temp_avg, vth = self._build_adiabatic_distribution(f, inputs)
        f_nonad = f - f_adi

        nphi, nnode, _, _ = f.shape
        mass = float(inputs.mass_kg)
        sp_index = inputs.species_index

        # Parallel energy flux (new): ∫ (m/2) v^2 v_par f dV
        vpar, vperp = self._velocity_mesh(nphi, nnode)
        energy_prefac = 0.5 * mass * (vth[:, :, np.newaxis, np.newaxis] ** 2)
        kernel_parallel_energy_flux = energy_prefac * (vpar**2 + vperp**2) * (vpar * vth[:, :, np.newaxis, np.newaxis])
        drift_k, exb_k, dbr_k_re, dbr_k_im = self._radial_kernels(inputs, vth, den_avg, temp_avg)
        energy_weight = energy_prefac * (vpar**2 + vperp**2)
        use_n0_turb = bool(inputs.split_n0_turb) and (f.shape[0] > 1)

        if use_n0_turb:
            # Match the C++ n0/turb decomposition at the distribution level.
            f_adi_n0 = self._toroidal_n0(f_adi)
            f_non_n0 = self._toroidal_n0(f_nonad)
            f_adi_turb = f_adi - f_adi_n0
            f_non_turb = f_nonad - f_non_n0
            # Split ExB kernel itself to reduce cancellation sensitivity in
            # near-zero toroidally averaged signals.
            exb_k_n0, exb_k_turb = self._split_kernel_n0_turb(exb_k)
        else:
            f_adi_n0 = f_adi
            f_non_n0 = f_nonad
            f_adi_turb = None
            f_non_turb = None
            exb_k_n0 = exb_k
            exb_k_turb = None

        pflux_delta_b = None
        eflux_delta_b = None
        has_delta_b = (
            inputs.f_3d_re is not None and inputs.f_3d_im is not None and
            dbr_k_re is not None and dbr_k_im is not None
        )
        if has_delta_b:
            f3r = np.asarray(inputs.f_3d_re)
            f3i = np.asarray(inputs.f_3d_im)
            # Match f0_sum factor 1/pi in the combined real+imaginary contribution.
            pflux_delta_b_total = (
                self._integrate(f3r, dbr_k_re, species_index=sp_index)
                + self._integrate(f3i, dbr_k_im, species_index=sp_index)
            ) / np.pi
            # delta-B path is not decomposed into adiabatic/non-adiabatic in
            # this implementation; store it on the adiabatic slot by convention.
            pflux_delta_b = SplitMoment(adiabatic=pflux_delta_b_total, nonadiabatic=None)

            eflux_delta_b_total = (
                self._integrate(f3r, dbr_k_re * energy_weight, species_index=sp_index) +
                self._integrate(f3i, dbr_k_im * energy_weight, species_index=sp_index)
            ) / np.pi
            eflux_delta_b = SplitMoment(adiabatic=eflux_delta_b_total, nonadiabatic=None)

        def _build_component(f_adia_comp, f_nonadia_comp, exb_kernel_comp, *, include_delta_b):
            f_total_comp = f_adia_comp + f_nonadia_comp
            den_adia = self._density(f_adia_comp, species_index=sp_index)
            den_non = self._density(f_nonadia_comp, species_index=sp_index)
            den_tot = den_adia + den_non

            upar_tot = self._mean_parallel_flow(f_total_comp, den_tot, vth, species_index=sp_index)

            _, tpar_adia, tperp_adia = self._temperature_components_ev(
                f_adia_comp, den_adia, upar_tot, vth, mass, float(inputs.charge_C), species_index=sp_index
            )
            _, tpar_non, tperp_non = self._temperature_components_ev(
                f_nonadia_comp, den_non, upar_tot, vth, mass, float(inputs.charge_C), species_index=sp_index
            )

            qpar_adia = self._integrate(f_adia_comp, kernel_parallel_energy_flux, species_index=sp_index)
            qpar_non = self._integrate(f_nonadia_comp, kernel_parallel_energy_flux, species_index=sp_index)

            pflux_drift_adia = self._integrate(f_adia_comp, drift_k, species_index=sp_index)
            pflux_drift_non = self._integrate(f_nonadia_comp, drift_k, species_index=sp_index)
            pflux_exb_adia = self._integrate(f_adia_comp, exb_kernel_comp, species_index=sp_index)
            pflux_exb_non = self._integrate(f_nonadia_comp, exb_kernel_comp, species_index=sp_index)

            eflux_drift_adia = self._integrate(f_adia_comp, drift_k * energy_weight, species_index=sp_index)
            eflux_drift_non = self._integrate(f_nonadia_comp, drift_k * energy_weight, species_index=sp_index)
            eflux_exb_adia = self._integrate(f_adia_comp, exb_kernel_comp * energy_weight, species_index=sp_index)
            eflux_exb_non = self._integrate(f_nonadia_comp, exb_kernel_comp * energy_weight, species_index=sp_index)

            return F0MomentComponent(
                density=den_tot,
                mean_parallel_flow=upar_tot,
                # Basic moments are branch totals by design (no adia/nonadia split).
                parallel_temperature=tpar_adia + tpar_non,
                perpendicular_temperature=tperp_adia + tperp_non,
                parallel_energy_flux=SplitMoment(adiabatic=qpar_adia, nonadiabatic=qpar_non),
                radial_particle_flux=FluxSplit(
                    drift=SplitMoment(adiabatic=pflux_drift_adia, nonadiabatic=pflux_drift_non),
                    exb=SplitMoment(adiabatic=pflux_exb_adia, nonadiabatic=pflux_exb_non),
                    delta_b=pflux_delta_b if include_delta_b else None,
                ),
                radial_energy_flux=FluxSplit(
                    drift=SplitMoment(adiabatic=eflux_drift_adia, nonadiabatic=eflux_drift_non),
                    exb=SplitMoment(adiabatic=eflux_exb_adia, nonadiabatic=eflux_exb_non),
                    delta_b=eflux_delta_b if include_delta_b else None,
                ),
            )

        n0_component = _build_component(f_adi_n0, f_non_n0, exb_k_n0, include_delta_b=True)
        turb_component = None
        if use_n0_turb and f_adi_turb is not None and f_non_turb is not None and exb_k_turb is not None:
            turb_component = _build_component(f_adi_turb, f_non_turb, exb_k_turb, include_delta_b=False)

        return F0MomentResult(n0=n0_component, turb=turb_component)






def compare_f0_moments_to_fmoment(
    result: F0MomentResult,
    fmoment_data,
    *,
    species_prefix: str,
    step_index: int = 0,
    atol: float = 0.0,
    rtol: float = 1e-3,
    rel_floor_frac: float = 1e-3,
):
    """
    Compare selected computed moments against ``FMomentData`` reference arrays.

    Parameters
    ----------
    result : F0MomentResult
        Output from :meth:`F0MomentCalculator.compute`.
    fmoment_data : FMomentData
        Reader containing XGC-written moment diagnostics.
    species_prefix : {"e", "i", "i2", ...}
        Species prefix used by the distribution variable. This helper currently
        supports the primary XGC fmoment prefixes ``e`` and ``i``.
    step_index : int, default 0
        Index into the ``FMomentData`` nested-dict storage.
    atol, rtol : float
        Absolute/relative tolerances for ``np.allclose`` checks.
    rel_floor_frac : float, default 1e-3
        Relative-denominator floor as a fraction of a robust reference scale
        (95th percentile of ``abs(ref)``). This prevents near-zero reference
        values from producing uninformative, arbitrarily large relative errors.

    Returns
    -------
    dict[str, dict]
        Per-moment comparison summary containing ``ok``, ``max_abs``, and
        ``max_rel``.

    Notes
    -----
    This helper is intentionally simple and targets quick validation. It checks
    the moments most likely to be present in ``xgc.fmoment`` diagnostics. The
    comparison currently targets the ``n0`` branch of :class:`F0MomentResult`.
    """
    # FMoment uses e*/i* names, not generalized i2/i3... in all campaigns.
    if species_prefix not in ("e", "i"):
        raise ValueError("compare_f0_moments_to_fmoment currently supports species_prefix 'e' or 'i'.")

    # Map semantic outputs to FMomentData variable names provided by the user.
    ref_map = {
        "density": f"{species_prefix}_den",
        "mean_parallel_flow": f"{species_prefix}_u_para",
        "parallel_temperature": f"{species_prefix}_T_para",
        "perpendicular_temperature": f"{species_prefix}_T_perp",
        "radial_particle_flux_drift_adiabatic": f"{species_prefix}_radial_flux_mag_f0",
        "radial_particle_flux_drift_nonadiabatic": f"{species_prefix}_radial_flux_mag_df",
        "radial_particle_flux_exb_adiabatic": f"{species_prefix}_radial_flux_ExB_f0",
        "radial_particle_flux_exb_nonadiabatic": f"{species_prefix}_radial_flux_ExB_df",
        "radial_energy_flux_drift_adiabatic": f"{species_prefix}_radial_en_flux_mag_f0",
        "radial_energy_flux_drift_nonadiabatic": f"{species_prefix}_radial_en_flux_mag_df",
        "radial_energy_flux_exb_adiabatic": f"{species_prefix}_radial_en_flux_ExB_f0",
        "radial_energy_flux_exb_nonadiabatic": f"{species_prefix}_radial_en_flux_ExB_df",
    }

    comp = result.n0
    computed = {
        "density": comp.density,
        "mean_parallel_flow": comp.mean_parallel_flow,
        "parallel_temperature": comp.parallel_temperature,
        "perpendicular_temperature": comp.perpendicular_temperature,
        "radial_particle_flux_drift_adiabatic": comp.radial_particle_flux.drift.adiabatic,
        "radial_particle_flux_drift_nonadiabatic": comp.radial_particle_flux.drift.nonadiabatic,
        "radial_particle_flux_exb_adiabatic": comp.radial_particle_flux.exb.adiabatic,
        "radial_particle_flux_exb_nonadiabatic": comp.radial_particle_flux.exb.nonadiabatic,
        "radial_energy_flux_drift_adiabatic": comp.radial_energy_flux.drift.adiabatic,
        "radial_energy_flux_drift_nonadiabatic": comp.radial_energy_flux.drift.nonadiabatic,
        "radial_energy_flux_exb_adiabatic": comp.radial_energy_flux.exb.adiabatic,
        "radial_energy_flux_exb_nonadiabatic": comp.radial_energy_flux.exb.nonadiabatic,
    }

    out = {}
    for key, ref_var in ref_map.items():
        if not fmoment_data.has_var(ref_var):
            out[key] = {"ok": False, "reason": f"reference variable '{ref_var}' missing"}
            continue

        ref_obj = fmoment_data.get_item(ref_var, step_index)
        if hasattr(ref_obj, "get_data"):
            ref = np.asarray(ref_obj.get_data())
        else:
            ref = np.asarray(ref_obj)

        calc = np.asarray(computed[key])

        # If calc is stored with phi singleton and ref is 1D, compare plane 0.
        if calc.ndim == 2 and calc.shape[0] == 1 and ref.ndim == 1:
            calc_cmp = calc[0]
        else:
            calc_cmp = calc

        if ref.ndim == 2 and ref.shape[0] == 1 and calc_cmp.ndim == 1:
            ref_cmp = ref[0]
        else:
            ref_cmp = ref

        if calc_cmp.shape != ref_cmp.shape:
            out[key] = {
                "ok": False,
                "reason": f"shape mismatch calc={calc_cmp.shape} ref={ref_cmp.shape}",
            }
            continue

        diff = calc_cmp - ref_cmp
        max_abs = float(np.nanmax(np.abs(diff)))
        ref_scale = float(np.nanpercentile(np.abs(ref_cmp), 95))
        rel_floor = max(float(atol), rel_floor_frac * ref_scale)
        denom = np.maximum(np.abs(ref_cmp), rel_floor)
        max_rel = float(np.nanmax(np.abs(diff) / np.where(denom == 0.0, 1.0, denom)))

        # Use the same denominator floor in the pass/fail condition.
        abs_err = np.abs(diff)
        tol = float(atol) + float(rtol) * np.maximum(np.abs(ref_cmp), rel_floor)
        ok = bool(np.all(abs_err <= tol))
        out[key] = {
            "ok": ok,
            "max_abs": max_abs,
            "max_rel": max_rel,
            "ref_var": ref_var,
            "rel_floor": rel_floor,
        }

    return out