process_module_runner.py

#!/usr/bin/env python3
"""Run selected PROCESS model calculations for cross-code module comparison."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any

import numpy as np

from process import constants
from process.current_drive import NeutralBeam
from process.data_structure import current_drive_variables as cdv
from process.data_structure import physics_variables as pv
from process.fusion_reactions import (
    REACTION_CONSTANTS_DT,
    BoschHaleConstants,
    beam_fusion,
    bosch_hale_reactivity,
)
from process.l_h_transition import calculate_martin08_nominal
from process.plasma_geometry import fvol, perim, surfa
from process.plasma_profiles import PlasmaProfile


def apply_case(case: dict[str, Any]) -> None:
    rmajor = case["rmajor"]
    rminor = case.get("rminor", case["rmajor"] / case["aspect"])
    aspect = case.get("aspect", rmajor / rminor)
    eps = rminor / rmajor
    f_t_beam = case["f_density_t_beam"]
    f_d_fuel = case["f_density_d_fuel"]
    f_t_fuel = case["f_density_t_fuel"]

    pv.rmajor = rmajor
    pv.rminor = rminor
    pv.eps = eps
    pv.aspect = aspect
    pv.kappa = case.get("kappa", 1.7)
    pv.delta = case.get("delta", 0.33)
    pv.vol_plasma = case.get("vol_plasma", 800.0)
    pv.alphan = case.get("alpha_ndensity", 1.0)
    pv.alphat = case.get("alpha_temp", 1.0)
    pv.dlamie = case["dlamie"]
    pv.m_beam_amu = case["m_beam_amu"]
    pv.nd_plasma_electrons_vol_avg = case["ndensity_e_vol_avg"]
    pv.nd_plasma_electron_line = case["ndensity_e_line_avg"]
    pv.nd_plasma_fuel_ions_vol_avg = case["ndensity_fuel_vol_avg"]
    pv.temp_plasma_electron_vol_avg_kev = case["temp_e_vol_avg"]
    pv.temp_plasma_electron_density_weighted_kev = case["temp_e_avg_density_weighted"]
    pv.temp_plasma_ion_vol_avg_kev = case["temp_i_vol_avg"]
    pv.temp_plasma_ion_density_weighted_kev = case["temp_i_avg_density_weighted"]
    pv.n_charge_plasma_effective_vol_avg = case["charge_eff_total_ion"]
    pv.n_charge_plasma_effective_mass_weighted_vol_avg = case["charge_eff_mass_weighted"]
    pv.f_nd_alpha_electron = case["f_density_alpha_electron"]
    pv.f_nd_plasma_carbon_electron = case["f_density_imp_c_electron"]
    pv.f_nd_plasma_oxygen_electron = case["f_density_imp_o_electron"]
    pv.f_nd_plasma_iron_argon_electron = case["f_density_imp_fe_ar_electron"]
    pv.f_plasma_fuel_deuterium = f_d_fuel
    pv.f_plasma_fuel_tritium = f_t_fuel
    pv.b_plasma_poloidal_average = case["b_poloidal_avg"]
    pv.b_plasma_toroidal_on_axis = case["b_toroidal_rmajor"]
    pv.i_plasma_ignited = case.get("i_plasma_ignited", 0)
    pv.beamfus0 = case["beamfus0"]
    pv.betbm0 = case["betbm0"]
    pv.sigmav_dt_average = case["sigmav_dt_average"]

    cdv.e_beam_kev = case["e_beam_kev"]
    cdv.f_radius_beam_tangency_rmajor = case["f_r_beam_tangency_rmajor"]
    cdv.f_beam_tritium = f_t_beam
    cdv.c_beam_total = case["current_beam_total"]


def run_efficiency(kind: str) -> dict[str, float]:
    nb = NeutralBeam(plasma_profile=PlasmaProfile())
    if kind == "iter":
        eta_cd, f_ion, f_shine = nb.iternb()
        n_decay = cdv.n_beam_decay_lengths_core
    elif kind == "cul":
        eta_cd, f_ion, f_shine = nb.culnbi()
        n_decay = cdv.n_beam_decay_lengths_core
    else:
        raise ValueError(f"unknown kind: {kind}")
    return {
        "eta_cd": float(eta_cd),
        "f_power_beam_injected_to_ion": float(f_ion),
        "f_shine_through": float(f_shine),
        "n_beam_decay_lengths_core": float(n_decay),
    }


def run_beam_fusion() -> dict[str, float]:
    beta_beam, nd_beam, p_alpha_mw = beam_fusion(
        pv.beamfus0,
        pv.betbm0,
        pv.b_plasma_poloidal_average,
        pv.b_plasma_toroidal_on_axis,
        cdv.c_beam_total,
        pv.nd_plasma_electrons_vol_avg,
        pv.nd_plasma_fuel_ions_vol_avg,
        pv.dlamie,
        cdv.e_beam_kev,
        pv.f_plasma_fuel_deuterium,
        pv.f_plasma_fuel_tritium,
        cdv.f_beam_tritium,
        pv.sigmav_dt_average,
        pv.temp_plasma_electron_density_weighted_kev,
        pv.temp_plasma_ion_density_weighted_kev,
        pv.vol_plasma,
        pv.n_charge_plasma_effective_mass_weighted_vol_avg,
    )
    return {
        "beta_beam": float(beta_beam),
        "ndensity_beam": float(nd_beam),
        "power_alpha_beam_mw": float(p_alpha_mw),
    }


def run_fusion_reactivity(case: dict[str, Any]) -> dict[str, float]:
    temps = np.asarray(case.get("fusion_reactivity_temps_kev", [10.0]), dtype=float)
    dt = BoschHaleConstants(**REACTION_CONSTANTS_DT)
    sigmv = bosch_hale_reactivity(temps, dt)
    return {f"sigmv_dt_t{temp:g}kev": float(value) for temp, value in zip(temps, sigmv)}


def run_plasma_geometry(case: dict[str, Any]) -> dict[str, float]:
    rmajor = case["rmajor"]
    rminor = case.get("rminor", case["rmajor"] / case["aspect"])
    kappa = case.get("kappa", 1.7)
    delta = case.get("delta", 0.33)
    surface_total, surface_outer = surfa(rminor, rmajor, kappa, delta)
    return {
        "vol_plasma": float(fvol(rmajor, rminor, kappa, delta)),
        "surface_area_plasma": float(surface_total),
        "surface_area_plasma_outer": float(surface_outer),
        "length_plasma": float(perim(rminor, kappa, delta)),
    }


def run_b_poloidal(case: dict[str, Any]) -> dict[str, float]:
    geom = run_plasma_geometry(case)
    i_plasma_current = case.get("i_plasma_current", 1)
    if i_plasma_current != 2:
        b_pol = constants.RMU0 * case["current_plasma_total_ampere"] / geom["length_plasma"]
    else:
        raise NotImplementedError("TART b_poloidal model not wired in module runner")
    return {"b_poloidal_avg": float(b_pol)}


def run_lh_martin08_nominal(case: dict[str, Any]) -> dict[str, float]:
    geom = run_plasma_geometry(case)
    dnla20 = case["ndensity_e_line_avg"] / 1.0e20
    p_lh = calculate_martin08_nominal(
        dnla20,
        case["b_toroidal_rmajor"],
        geom["surface_area_plasma"],
        case["m_ions_total_amu"],
    )
    return {"p_lh_martin08_nominal_mw": float(p_lh)}


TASKS = {
    "efficiency": lambda case, kind: run_efficiency(kind),
    "beam_fusion": lambda case, kind: run_beam_fusion(),
    "fusion_reactivity": lambda case, kind: run_fusion_reactivity(case),
    "plasma_geometry": lambda case, kind: run_plasma_geometry(case),
    "b_poloidal": lambda case, kind: run_b_poloidal(case),
    "lh_martin08_nominal": lambda case, kind: run_lh_martin08_nominal(case),
}


def main() -> int:
    parser = argparse.ArgumentParser(description="Run PROCESS calculations for module comparison.")
    parser.add_argument("--task", choices=tuple(TASKS), required=True)
    parser.add_argument("--kind", choices=("iter", "cul"), default="iter")
    parser.add_argument("--case-file", type=Path, required=True)
    args = parser.parse_args()

    case = json.loads(args.case_file.read_text(encoding="utf-8"))
    apply_case(case)
    result = TASKS[args.task](case, args.kind)
    json.dump(result, sys.stdout)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())