Source code for ewokstomo.tasks.energycalculation

from __future__ import annotations
import logging
from typing import Any, TYPE_CHECKING
import numpy as np
import xraylib
from ewokscore import Task
from ewokscore.model import BaseInputModel
from ewokscore.model import BaseOutputModel
from pydantic import Field
from xoppylib.sources.xoppy_bm_wiggler import xoppy_calc_bm
from xoppylib.sources.xoppy_bm_wiggler import xoppy_calc_wiggler_on_aperture

logger = logging.getLogger(__name__)

_NIST_ALIASES = {
    "air": "Air, Dry (near sea level)",
    "water": "Water, Liquid",
    "kapton": "Kapton Polyimide Film",
}
try:
    for n in xraylib.GetCompoundDataNISTList():
        _NIST_ALIASES[n] = n
        _NIST_ALIASES[n.lower()] = n
        _NIST_ALIASES[n.replace(" ", "")] = n
        _NIST_ALIASES[n.lower().replace(" ", "")] = n
except Exception as e:
    logger.debug(
        "Could not load NIST compound list from xraylib; using base aliases only: %s",
        e,
    )


def _is_number(x) -> bool:
    try:
        float(x)
        return True
    except Exception:
        return False


if TYPE_CHECKING:  # pragma: no cover - typing-only import
    pass

try:
    from ewokscore.missing_data import MissingData as _MissingDataRuntime  # type: ignore[assignment]
except Exception:  # pragma: no cover - fallback for optional dependency
    _MissingDataRuntime = None  # type: ignore[assignment]


def _is_missing_data(value: Any) -> bool:
    if value is None:
        return False
    if _MissingDataRuntime is not None and isinstance(value, _MissingDataRuntime):
        return True
    return value.__class__.__name__ == "MissingData"


def _canonical_compound_name(name: str) -> str | None:
    if not name:
        return None
    key1 = name.strip()
    key2 = key1.lower()
    key3 = key2.replace(" ", "")
    return _NIST_ALIASES.get(key1) or _NIST_ALIASES.get(key2) or _NIST_ALIASES.get(key3)


def _mu_over_rho_element(Z: int, E_keV: np.ndarray) -> np.ndarray:
    # cm^2/g
    return np.array([xraylib.CS_Total(Z, e) for e in E_keV], dtype=float)


def _mu_over_rho_compound_from_nist(
    cname: str, E_keV: np.ndarray
) -> tuple[np.ndarray, float | None]:
    """
    Return (mu_over_rho [cm^2/g], density [g/cm^3 or None]) for a NIST compound name.
    mu/ρ is computed as the mass-fraction weighted sum of elemental mu/ρ.
    """
    data = xraylib.GetCompoundDataNISTByName(cname)
    Zs = data["Elements"]
    w = data["massFractions"]
    mu_over_rho = np.zeros_like(E_keV, dtype=float)
    for Z, wf in zip(Zs, w):
        mu_over_rho += wf * _mu_over_rho_element(int(Z), E_keV)
    return mu_over_rho, float(data["density"])


def _resolve_mu_density(
    material: str, density_g_cm3: Any, E_keV: np.ndarray
) -> tuple[np.ndarray, float]:
    """
    Resolve mass attenuation (mu/ρ) and density for either an element or a NIST compound.
    - If density is '?', None or missing, we auto-fill:
        * elements: xraylib.ElementDensity(Z)
        * compounds: NIST density
    """
    try:
        Z = xraylib.SymbolToAtomicNumber(material)
    except Exception as e:
        logger.debug(
            "Material is not an element symbol (%r); trying compound paths: %s",
            material,
            e,
        )
    else:
        mu_over_rho = _mu_over_rho_element(Z, E_keV)
        if _is_number(density_g_cm3):
            rho = float(density_g_cm3)
        else:
            rho = float(xraylib.ElementDensity(Z))
        return mu_over_rho, rho

    cname = _canonical_compound_name(material)
    if cname is not None:
        mu_over_rho, rho_nist = _mu_over_rho_compound_from_nist(cname, E_keV)
        if _is_number(density_g_cm3):
            rho = float(density_g_cm3)
        elif rho_nist is not None:
            rho = float(rho_nist)
        else:
            raise ValueError(f"NIST density unavailable for compound '{cname}'")
        return mu_over_rho, rho
    try:
        comp = xraylib.CompoundParser(material)
        Zs = comp["Elements"]
        w = comp["massFractions"]
        mu_over_rho = np.zeros_like(E_keV, dtype=float)
        for Z, wf in zip(Zs, w):
            mu_over_rho += wf * _mu_over_rho_element(int(Z), E_keV)
        rho = float(density_g_cm3) if _is_number(density_g_cm3) else 1.0
        return mu_over_rho, rho
    except Exception as e:
        raise ValueError(f"Unknown material '{material}': {e}")


def _transmission(
    material: str, thickness_mm: float, density_g_cm3: Any, energy_eV: np.ndarray
) -> np.ndarray:
    if thickness_mm <= 0:
        return np.ones_like(energy_eV, dtype=float)
    E_keV = np.asarray(energy_eV, dtype=float) / 1e3
    mu_over_rho, rho = _resolve_mu_density(
        material, density_g_cm3, E_keV
    )  # cm^2/g, g/cm^3
    mu = mu_over_rho * rho  # cm^-1
    t_cm = float(thickness_mm) / 10.0
    return np.exp(-mu * t_cm)




class ComputeBMSpectrumInputModel(BaseInputModel):
    TYPE_CALC: int = Field(
        0, description="Calculation type; must be 0 (energy spectrum)."
    )
    VER_DIV: int = Field(
        0, description="0 integrates flux over full Psi, 2 over [PSI_MIN, PSI_MAX]."
    )
    MACHINE_NAME: str = Field(
        "ESRF bending magnet",
        description="Machine name label; not used in the calculation.",
    )
    RB_CHOICE: int = Field(
        0, description="Radius source: 0 uses MACHINE_R_M, 1 derives it from BFIELD_T."
    )
    MACHINE_R_M: float = Field(
        25.0, description="Bending magnet radius (m); used when RB_CHOICE=0."
    )
    BFIELD_T: float = Field(
        0.8,
        description="Magnetic field (T); used when RB_CHOICE=1 to derive the radius.",
    )
    BEAM_ENERGY_GEV: float = Field(6.0, description="Beam energy (GeV).")
    CURRENT_A: float = Field(0.2, description="Ring current (A).")
    HOR_DIV_MRAD: float = Field(1.0, description="Horizontal divergence (mrad).")
    PHOT_ENERGY_MIN: float = Field(100.0, description="Minimum photon energy (eV).")
    PHOT_ENERGY_MAX: float = Field(200000.0, description="Maximum photon energy (eV).")
    NPOINTS: int = Field(500, description="Number of energy points.")
    LOG_CHOICE: int = Field(1, description="Energy grid spacing: 0=linear, 1=log.")
    PSI_MRAD_PLOT: float = Field(
        1.0, description="Unused when TYPE_CALC=0 (only affects other calc types)."
    )
    PSI_MIN: float = Field(
        -1.0, description="Minimum psi integration angle (mrad); used when VER_DIV=2."
    )
    PSI_MAX: float = Field(
        1.0, description="Maximum psi integration angle (mrad); used when VER_DIV=2."
    )
    PSI_NPOINTS: int = Field(
        500, description="Number of psi integration points; used when VER_DIV=2."
    )
    FILE_DUMP: bool = Field(
        False, description="Whether XOPPY dumps results to bm.spec."
    )





class ComputeSpectrumOutputModel(BaseOutputModel):
    energy_eV: np.ndarray = Field(
        ..., description="Photon energy grid (eV), sorted ascending."
    )
    flux: np.ndarray = Field(
        ..., description="Flux (phot/s/0.1%bw) as returned by XOPPY."
    )
    spectral_power: np.ndarray = Field(..., description="Spectral power (W/eV).")
    cumulated_power: np.ndarray = Field(..., description="Cumulated power (W).")





class ComputeBMSpectrum(  # type: ignore[call-arg]
    Task,
    input_model=ComputeBMSpectrumInputModel,
    output_model=ComputeSpectrumOutputModel,
):
    """
    Compute a bending-magnet (BM) spectrum using XOPPY's ``xoppy_calc_bm``.
    """



    def run(self):
        TYPE_CALC = self.inputs.TYPE_CALC
        VER_DIV = self.inputs.VER_DIV
        MACHINE_NAME = self.inputs.MACHINE_NAME
        RB_CHOICE = self.inputs.RB_CHOICE
        MACHINE_R_M = self.inputs.MACHINE_R_M
        BFIELD_T = self.inputs.BFIELD_T
        BEAM_ENERGY_GEV = self.inputs.BEAM_ENERGY_GEV
        CURRENT_A = self.inputs.CURRENT_A
        HOR_DIV_MRAD = self.inputs.HOR_DIV_MRAD
        PHOT_ENERGY_MIN = self.inputs.PHOT_ENERGY_MIN
        PHOT_ENERGY_MAX = self.inputs.PHOT_ENERGY_MAX
        NPOINTS = self.inputs.NPOINTS
        LOG_CHOICE = self.inputs.LOG_CHOICE
        PSI_MRAD_PLOT = self.inputs.PSI_MRAD_PLOT
        PSI_MIN = self.inputs.PSI_MIN
        PSI_MAX = self.inputs.PSI_MAX
        PSI_NPOINTS = self.inputs.PSI_NPOINTS
        FILE_DUMP = self.inputs.FILE_DUMP

        a6_T, fm, a, energy_eV = xoppy_calc_bm(
            TYPE_CALC=TYPE_CALC,
            MACHINE_NAME=MACHINE_NAME,
            RB_CHOICE=RB_CHOICE,
            MACHINE_R_M=MACHINE_R_M,
            BFIELD_T=BFIELD_T,
            BEAM_ENERGY_GEV=BEAM_ENERGY_GEV,
            CURRENT_A=CURRENT_A,
            HOR_DIV_MRAD=HOR_DIV_MRAD,
            VER_DIV=VER_DIV,
            PHOT_ENERGY_MIN=PHOT_ENERGY_MIN,
            PHOT_ENERGY_MAX=PHOT_ENERGY_MAX,
            NPOINTS=NPOINTS,
            LOG_CHOICE=LOG_CHOICE,
            PSI_MRAD_PLOT=PSI_MRAD_PLOT,
            PSI_MIN=PSI_MIN,
            PSI_MAX=PSI_MAX,
            PSI_NPOINTS=PSI_NPOINTS,
            FILE_DUMP=FILE_DUMP,
        )
        if TYPE_CALC != 0 or VER_DIV not in (0, 2):
            raise ValueError(
                "ComputeBMSpectrum expects TYPE_CALC=0 and VER_DIV in {0,2}"
            )

        flux = a6_T[:, 5]
        spectral_power = a6_T[:, 6]
        cum_power = a6_T[:, 7]

        order = np.argsort(energy_eV)
        self.outputs.energy_eV = energy_eV[order]
        self.outputs.flux = flux[order]
        self.outputs.spectral_power = spectral_power[order]
        self.outputs.cumulated_power = cum_power[order]






class ComputeWigglerSpectrumInputModel(BaseInputModel):
    PHOT_ENERGY_MIN: float = Field(100.0, description="Minimum photon energy (eV).")
    PHOT_ENERGY_MAX: float = Field(4e5, description="Maximum photon energy (eV).")
    NPOINTS: int = Field(2000, description="Number of energy points.")
    ENERGY: float = Field(6.0, description="Beam energy (GeV).")
    CURRENT: float = Field(200.0, description="Ring current (mA).")
    FIELD: int = Field(
        1,
        description="Trajectory source: 0=sinusoidal, 1=field map (FILE), 2=harmonics.",
    )
    NPERIODS: int = Field(
        1, description="Number of periods; used when FIELD is 0 or 2."
    )
    ULAMBDA: float = Field(
        0.15, description="Period length (m); used when FIELD is 0 or 2."
    )
    K: float = Field(22.591, description="Deflection parameter; used when FIELD=0.")
    NTRAJPOINTS: int = Field(101, description="Number of trajectory points.")
    FILE: str = Field("", description="Magnetic field map path; required when FIELD=1.")
    SLIT_FLAG: int = Field(
        1, description="0=full-space spectrum, nonzero=integrate over the slit below."
    )
    SLIT_D: float = Field(
        56.5, description="Distance to the slit (m); used when SLIT_FLAG!=0."
    )
    SLIT_NY: int = Field(
        101, description="Number of vertical angular points; used when SLIT_FLAG!=0."
    )
    SLIT_WIDTH_H_MM: float = Field(
        10.0, description="Slit horizontal width (mm); used when SLIT_FLAG!=0."
    )
    SLIT_HEIGHT_V_MM: float = Field(
        5.0, description="Slit vertical height (mm); used when SLIT_FLAG!=0."
    )
    SLIT_CENTER_H_MM: float = Field(
        0.0, description="Slit horizontal center (mm); used when SLIT_FLAG!=0."
    )
    SLIT_CENTER_V_MM: float = Field(
        0.0, description="Slit vertical center (mm); used when SLIT_FLAG!=0."
    )
    SHIFT_X_FLAG: int = Field(
        1, description="Horizontal shift flag; used when FIELD=1."
    )
    SHIFT_X_VALUE: float = Field(
        -0.002385, description="Horizontal shift value (m); used when FIELD=1."
    )
    SHIFT_BETAX_FLAG: int = Field(
        5, description="Beta_x shift flag; used when FIELD=1."
    )
    SHIFT_BETAX_VALUE: float = Field(
        0.005, description="Beta_x shift value (rad); used when FIELD=1."
    )
    TRAJ_RESAMPLING_FACTOR: float = Field(
        10000.0, description="Trajectory resampling factor; used when SLIT_FLAG!=0."
    )
    SLIT_POINTS_FACTOR: float = Field(
        3.0, description="Slit points factor; used when SLIT_FLAG!=0."
    )
    LOG_CHOICE: int = Field(
        1,
        description="Unused: xoppy_calc_wiggler_on_aperture has no log-sampling option.",
    )





class ComputeWigglerSpectrum(  # type: ignore[call-arg]
    Task,
    input_model=ComputeWigglerSpectrumInputModel,
    output_model=ComputeSpectrumOutputModel,
):
    """
    Compute a wiggler spectrum on an aperture using
    ``xoppy_calc_wiggler_on_aperture``.
    """



    def run(self):
        energy, flux, sp, cum, *_ = xoppy_calc_wiggler_on_aperture(
            FIELD=int(self.inputs.FIELD),
            NPERIODS=int(self.inputs.NPERIODS),
            ULAMBDA=float(self.inputs.ULAMBDA),  # m
            K=float(self.inputs.K),
            ENERGY=float(self.inputs.ENERGY),  # GeV
            PHOT_ENERGY_MIN=float(self.inputs.PHOT_ENERGY_MIN),
            PHOT_ENERGY_MAX=float(self.inputs.PHOT_ENERGY_MAX),
            NPOINTS=int(self.inputs.NPOINTS),
            NTRAJPOINTS=int(self.inputs.NTRAJPOINTS),
            CURRENT=float(self.inputs.CURRENT),  # mA
            FILE=str(self.inputs.FILE),
            SLIT_FLAG=int(self.inputs.SLIT_FLAG),
            SLIT_D=float(self.inputs.SLIT_D),  # m
            SLIT_NY=int(self.inputs.SLIT_NY),
            SLIT_WIDTH_H_MM=float(self.inputs.SLIT_WIDTH_H_MM),
            SLIT_HEIGHT_V_MM=float(self.inputs.SLIT_HEIGHT_V_MM),
            SLIT_CENTER_H_MM=float(self.inputs.SLIT_CENTER_H_MM),
            SLIT_CENTER_V_MM=float(self.inputs.SLIT_CENTER_V_MM),
            SHIFT_X_FLAG=int(self.inputs.SHIFT_X_FLAG),
            SHIFT_X_VALUE=float(self.inputs.SHIFT_X_VALUE),  # m
            SHIFT_BETAX_FLAG=int(self.inputs.SHIFT_BETAX_FLAG),
            SHIFT_BETAX_VALUE=float(self.inputs.SHIFT_BETAX_VALUE),
            TRAJ_RESAMPLING_FACTOR=float(self.inputs.TRAJ_RESAMPLING_FACTOR),
            SLIT_POINTS_FACTOR=float(self.inputs.SLIT_POINTS_FACTOR),
        )
        energy = np.asarray(energy, dtype=float)
        order = np.argsort(energy)
        self.outputs.energy_eV = energy[order]
        self.outputs.flux = np.asarray(flux, dtype=float)[order]
        self.outputs.spectral_power = np.asarray(sp, dtype=float)[order]
        self.outputs.cumulated_power = np.asarray(cum, dtype=float)[order]






class ApplyAttenuatorsInputModel(BaseInputModel):
    energy_eV: np.ndarray = Field(
        ..., description="Photon energy grid in electron-volts."
    )
    spectral_power: np.ndarray = Field(..., description="Power spectrum (W/eV).")
    attenuators: dict[str, dict[str, Any]] = Field(
        ...,
        description=(
            "Mapping where each value contains ``material``, ``thickness_mm`` and "
            "optional ``density_g_cm3``. ``material`` accepts element symbols "
            '(e.g. "Al"), NIST aliases (e.g. "kapton") or chemical formulae '
            "parsable by xraylib."
        ),
    )
    order: list[str] | None = Field(
        None,
        description=(
            "Explicit stacking order of the attenuator keys. Defaults to the "
            "dictionary insertion order."
        ),
    )
    flux: np.ndarray | None = Field(
        None,
        description="Source flux array (phot/s/0.1%bw) matching ``energy_eV``.",
    )





class ApplyAttenuatorsOutputModel(BaseOutputModel):
    energy_eV: np.ndarray = Field(..., description="Same energy grid passed through.")
    attenuated_spectral_power: np.ndarray = Field(
        ..., description="``spectral_power`` (W/eV) multiplied by ``transmission``."
    )
    transmission: np.ndarray = Field(
        ..., description="Cumulative transmission of the attenuator stack (0-1)."
    )
    attenuated_flux: np.ndarray | None = Field(
        None,
        description=(
            "``flux`` (phot/s/0.1%bw) multiplied by ``transmission`` when "
            "provided, else None."
        ),
    )





class ApplyAttenuators(  # type: ignore[call-arg]
    Task,
    input_model=ApplyAttenuatorsInputModel,
    output_model=ApplyAttenuatorsOutputModel,
):
    """
    Apply a stack of attenuators to the source spectrum (and optionally flux).
    """



    def run(self):
        energy_eV = self.inputs.energy_eV
        sp_in = self.inputs.spectral_power.copy()
        flux_in = self.inputs.flux
        if _is_missing_data(flux_in):
            flux_in = None

        attenuators: dict[str, dict[str, Any]] = dict(self.inputs.attenuators)

        # order of stacking
        if self.inputs.order:
            keys: list[str] = list(self.inputs.order)
        else:
            keys = list(attenuators.keys())

        transmission = np.ones_like(energy_eV, dtype=float)
        for key in keys:
            a = attenuators[key]
            T = _transmission(
                str(a["material"]),
                float(a["thickness_mm"]),
                float(a["density_g_cm3"]),
                energy_eV,
            )
            transmission *= T

        sp_out = sp_in * transmission
        flux_out = flux_in * transmission if flux_in is not None else None

        self.outputs.energy_eV = energy_eV
        self.outputs.attenuated_spectral_power = sp_out
        self.outputs.transmission = transmission
        self.outputs.attenuated_flux = flux_out






class SpectrumStatsInputModel(BaseInputModel):
    energy_eV: np.ndarray = Field(..., description="Photon energy grid (eV).")
    attenuated_flux: np.ndarray = Field(
        ..., description="Attenuated flux (ph/s/0.1%bw)."
    )





class SpectrumStatsOutputModel(BaseOutputModel):
    mean_energy_eV: float = Field(
        ...,
        description=(
            "Flux-weighted mean using bin weights: weights = flux * ΔE / (0.001 * E)."
        ),
    )
    mean_idx: int = Field(
        ..., description="Index of energy_eV closest to mean_energy_eV (−1 if N/A)."
    )
    pic_energy_eV: float = Field(
        ...,
        description=(
            "Energy at which flux * ΔE / (0.001 * E) is maximal (NaN if N/A)."
        ),
    )
    pic_idx: int = Field(..., description="Index of that maximum (−1 if N/A).")





class SpectrumStats(  # type: ignore[call-arg]
    Task,
    input_model=SpectrumStatsInputModel,
    output_model=SpectrumStatsOutputModel,
):
    """
    Stats on an attenuated spectrum.
    """



    def run(self):
        energy = np.asarray(self.inputs.energy_eV, dtype=float)
        flux = np.asarray(self.inputs.attenuated_flux, dtype=float)

        if not (energy.shape == flux.shape):
            raise ValueError("energy_eV and attenuated_flux must have identical shapes")

        mask = np.isfinite(energy) & np.isfinite(flux) & (energy > 0)
        if not mask.any():
            self.outputs.mean_energy_eV = float("nan")
            self.outputs.mean_idx = -1
            self.outputs.pic_energy_eV = float("nan")
            self.outputs.pic_idx = -1
            return

        idx_masked = np.flatnonzero(mask)
        e = energy[mask]
        f = flux[mask]

        order = np.argsort(e)
        e_sorted = e[order]
        f_sorted = f[order]
        idx_sorted = idx_masked[order]

        if e_sorted.size == 1:
            dE = np.array([1.0], dtype=float)
        else:
            dE = np.empty_like(e_sorted)
            dE[1:-1] = 0.5 * (e_sorted[2:] - e_sorted[:-2])
            dE[0] = e_sorted[1] - e_sorted[0]
            dE[-1] = e_sorted[-1] - e_sorted[-2]
            dE = np.clip(dE, 0.0, None)

        weights = f_sorted * dE / (1e-3 * e_sorted)
        wsum = float(np.sum(weights))

        if wsum > 0.0:
            mean_energy_eV = float(np.sum(e_sorted * weights) / wsum)
            mean_idx = int(np.abs(energy - mean_energy_eV).argmin())
        else:
            mean_energy_eV = float("nan")
            mean_idx = -1

        if weights.size > 0:
            rel = int(np.argmax(weights))
            pic_idx = int(idx_sorted[rel])
            pic_energy_eV = float(energy[pic_idx])
        else:
            pic_idx = -1
            pic_energy_eV = float("nan")

        self.outputs.mean_energy_eV = mean_energy_eV
        self.outputs.mean_idx = mean_idx
        self.outputs.pic_energy_eV = pic_energy_eV
        self.outputs.pic_idx = pic_idx