NEOPAX._transport_equations¶

Attributes¶

`DENSITY_STATE_TO_PHYSICAL`
`PARTICLE_FLUX_PHYSICAL_TO_STATE`
`HEAT_FLUX_PHYSICAL_TO_STATE`
`__equation_registry`

Classes¶

`EquationBase`	Base class for transport equations. Subclasses must implement __call__.
`TransportLaggedResponse`
`DensityEquation`	Base class for transport equations. Subclasses must implement __call__.
`TemperatureEquation`	Base class for transport equations. Subclasses must implement __call__.
`ElectricFieldEquation`	Base class for transport equations. Subclasses must implement __call__.
`ComposedEquationSystem`

Functions¶

`_minmod_pair`(a, b)
`_minmod3`(a, b, c)
`_mc_limited_face_states`(profile_ghost)
`_temperature_face_states`(temperature_ghost, ...)
`_cell_centered_flux_faces`(flux, reconstruction_mode)
`enforce_quasi_neutrality`(state, species)	Reconstruct electron density from ion densities and species charges.
`project_fixed_temperature_species`(state[, ...])	Keep the closure temperature fixed for species whose temperature equation
`apply_er_dirichlet_boundary_state`(state, er_bc_model)	Clamp Er state values at configured Dirichlet boundaries.
`_expand_density_rhs_to_full_shape`(density_rhs, ...)	Expand a reduced density RHS back to full physical species ordering.
`_plasma_permitivity_from_prefactor`(state, ...)	Plasma permittivity on the transport grid using a precomputed geometry prefactor.
`register_equation`(name)	Decorator to register equation classes in the registry.
`get_equation`(→ Type)
`list_equations`()
`build_density_equation`(field, flux_model, ...[, ...])
`_build_species_faces_builder`(field, bc_model[, ...])
`_build_species_ghost_builder`(bc_model)
`build_temperature_equation`(field, flux_model, ...[, ...])
`build_electric_field_equation`(field, species_names, ...)
`_resolve_er_boundary_mode`(config, solver_cfg)
`build_equation_system`(config, species, field, flux_model)	Build the list of equation instances to evolve using prebuilt runtime
`build_equation_system_from_config`(config, species)	Backward-compatible wrapper that builds the required runtime objects from

Module Contents¶

NEOPAX._transport_equations.DENSITY_STATE_TO_PHYSICAL = 1e+20¶

NEOPAX._transport_equations.PARTICLE_FLUX_PHYSICAL_TO_STATE = 1e-20¶

NEOPAX._transport_equations.HEAT_FLUX_PHYSICAL_TO_STATE = 1e-23¶

NEOPAX._transport_equations._minmod_pair(a, b)¶

NEOPAX._transport_equations._minmod3(a, b, c)¶

NEOPAX._transport_equations._mc_limited_face_states(profile_ghost)¶

NEOPAX._transport_equations._temperature_face_states(temperature_ghost, reconstruction_mode)¶

NEOPAX._transport_equations._cell_centered_flux_faces(flux, reconstruction_mode)¶

NEOPAX._transport_equations.enforce_quasi_neutrality(state, species)¶: Reconstruct electron density from ion densities and species charges. Returns a new TransportState with quasi-neutral electron density.

NEOPAX._transport_equations.project_fixed_temperature_species(state, temperature_active_mask=None, fixed_temperature_profile=None, density_floor=DEFAULT_TRANSPORT_DENSITY_FLOOR)¶: Keep the closure temperature fixed for species whose temperature equation is disabled by projecting pressure = n * T_fixed on the working state.

NEOPAX._transport_equations.apply_er_dirichlet_boundary_state(state, er_bc_model)¶: Clamp Er state values at configured Dirichlet boundaries.

NEOPAX._transport_equations._expand_density_rhs_to_full_shape(density_rhs, template_density, species)¶: Expand a reduced density RHS back to full physical species ordering.

NEOPAX._transport_equations._plasma_permitivity_from_prefactor(state, species_mass, permitivity_prefactor)¶: Plasma permittivity on the transport grid using a precomputed geometry prefactor.

NEOPAX._transport_equations.__equation_registry: Dict[str, Type]¶

NEOPAX._transport_equations.register_equation(name: str)¶: Decorator to register equation classes in the registry.

class NEOPAX._transport_equations.EquationBase¶

Base class for transport equations. Subclasses must implement __call__.

abstractmethod __call__(state, models, field, species, energy_grid, database, solver_parameters, bc=None, **kwargs)¶

NEOPAX._transport_equations.get_equation(name: str) → Type¶

NEOPAX._transport_equations.list_equations()¶

class NEOPAX._transport_equations.TransportLaggedResponse¶

flux_response: object = None¶

class NEOPAX._transport_equations.DensityEquation¶

Bases: EquationBase

Base class for transport equations. Subclasses must implement __call__.

dr_cells: jax.Array¶

Vprime: jax.Array¶

Vprime_half: jax.Array¶

flux_model: callable¶

flux_faces_builder: callable¶

active_species_mask: jax.Array¶

independent_density_mask: jax.Array¶

face_flux_builder: callable = None¶

density_bc_model: object = None¶

particle_flux_reconstruction: str = 'closure_face_flux'¶

particle_face_closure_mode: str = 'reconstructed'¶

source_model: callable = None¶

species: object = None¶

name: str = 'density'¶

_mode_requests_face_fluxes(mode_value)¶

_use_model_face_particle_fluxes()¶

enforce_dirichlet_boundary_rhs(state, density_rhs)¶

debug_components(state, fluxes=None, source_outputs=None)¶

__call__(state, fluxes=None, source_outputs=None)¶

NEOPAX._transport_equations.build_density_equation(field, flux_model, source_model, bc_density, species, bc_temperature=None, bc_er=None, reconstruction='linear', active_species_mask=None, particle_flux_reconstruction='closure_face_flux', particle_face_closure_mode='reconstructed', density_floor=DEFAULT_TRANSPORT_DENSITY_FLOOR, temperature_floor=DEFAULT_TRANSPORT_TEMPERATURE_FLOOR)¶

class NEOPAX._transport_equations.TemperatureEquation¶

Bases: EquationBase

Base class for transport equations. Subclasses must implement __call__.

dr_cells: jax.Array¶

Vprime: jax.Array¶

Vprime_half: jax.Array¶

flux_model: callable¶

flux_faces_builder: callable¶

temperature_ghost_builder: callable¶

charge_qp: jax.Array¶

active_species_mask: jax.Array¶

face_flux_builder: callable = None¶

temperature_bc_model: object = None¶

convection_reconstruction: str = 'tvd_mc'¶

heat_flux_reconstruction: str = 'tvd_mc'¶

include_neo_convection: bool = True¶

include_turbulent_convection: bool = True¶

include_classical_convection: bool = True¶

include_work_term: bool = True¶

source_model: callable = None¶

species: object = None¶

name: str = 'temperature'¶

_mode_requests_face_fluxes(mode_value)¶

_use_model_face_heat_fluxes()¶

_use_model_face_particle_fluxes()¶

enforce_dirichlet_boundary_rhs(state, density_rhs, pressure_rhs)¶

debug_components(state, fluxes=None, source_outputs=None)¶

__call__(state, fluxes=None, source_outputs=None)¶

NEOPAX._transport_equations._build_species_faces_builder(field, bc_model, reconstruction='linear')¶

NEOPAX._transport_equations._build_species_ghost_builder(bc_model)¶

NEOPAX._transport_equations.build_temperature_equation(field, flux_model, source_model, species, bc_temperature, bc_density=None, bc_gamma=None, bc_er=None, active_species_mask=None, charge_qp=None, include_neo_convection=True, include_turbulent_convection=True, include_classical_convection=True, include_work_term=True, convection_reconstruction='tvd_mc', heat_flux_reconstruction='tvd_mc', reconstruction='linear', density_floor=DEFAULT_TRANSPORT_DENSITY_FLOOR, temperature_floor=DEFAULT_TRANSPORT_TEMPERATURE_FLOOR)¶

class NEOPAX._transport_equations.ElectricFieldEquation¶

Bases: EquationBase

Base class for transport equations. Subclasses must implement __call__.

dr_cells: jax.Array¶

Vprime: jax.Array¶

Vprime_half: jax.Array¶

flux_model: callable¶

species_mass: jax.Array¶

charge_qp: jax.Array¶

permitivity_prefactor: jax.Array¶

gamma_faces_builder: callable¶

er_diffusive_flux_builder: callable¶

er_bc_model: object = None¶

source_mode: str = 'ambipolar_local'¶

permitivity_mode: str = 'neopax_local'¶

Er_relax: float = 1.0¶

DEr: float = 1.0¶

boundary_mode: str = 'standard'¶

ntss_B0_mid: float = 0.0¶

ntss_psfactor_mid: float = 1.0¶

ntss_density_indices: jax.Array = None¶

name: str = 'Er'¶

_charge_flux_from_gamma(Gamma)¶

_er_diffusion(Er)¶

_charge_flux_and_ambi_term(state, Gamma, plasma_permitivity)¶

debug_components(state, fluxes=None)¶

__call__(state, fluxes=None)¶

enforce_dirichlet_boundary_rhs(state, er_rhs)¶

ap_linear_split(state)¶: Return diagonal linearization and explicit source for optional AP preconditioning. Uses only attributes set at construction and the current state.

NEOPAX._transport_equations.build_electric_field_equation(field, species_names, flux_model, species_mass, charge_qp, bc_gamma, bc_er, Er_relax=1.0, DEr=1.0, source_mode='ambipolar_local', permitivity_mode='neopax_local', reconstruction='linear', boundary_mode='standard')¶

NEOPAX._transport_equations._resolve_er_boundary_mode(config, solver_cfg)¶

NEOPAX._transport_equations.build_equation_system(config, species, field, flux_model, source_models=None, solver_cfg=None, boundary_models=None)¶: Build the list of equation instances to evolve using prebuilt runtime objects. This avoids rebuilding geometry, databases, and flux models inside the equation builder and keeps compile closures smaller.

NEOPAX._transport_equations.build_equation_system_from_config(config, species)¶: Backward-compatible wrapper that builds the required runtime objects from config before delegating to build_equation_system.

class NEOPAX._transport_equations.ComposedEquationSystem¶

equations: tuple¶

density_equation: object | None = None¶

temperature_equation: object | None = None¶

er_equation: object | None = None¶

species: object | None = None¶

shared_flux_model: object | None = None¶

density_floor: object¶

temperature_floor: object¶

temperature_active_mask: object | None = None¶

fixed_temperature_profile: object | None = None¶

er_bc_model: object | None = None¶

_prepare_working_state(state)¶

_resolve_equations()¶

build_lagged_response(state)¶

evaluate_with_lagged_response(t, state, runtime, lagged_response)¶

_evaluate_state(state, lagged_response=None)¶

__call__(t, state, runtime)¶: Call all equations with state, return a TransportState matching the state structure. Always output all three fields, setting missing ones to zero arrays of the correct shape. When electrons are present, evaluate the RHS on a quasi-neutral working state, but keep electron density out of the solved density subsystem. This matches the NTSS-style pattern: evolve independent ion/impurity density rows, reconstruct electron density algebraically for the working state and accepted/output states.

vector_field(t, y, args)¶: Torax-style vector field for JAX ODE solvers: (t, y, args) -> dy/dt y is the state, args[0] is the runtime dict.