NEOPAX._transport_flux_models¶

Attributes¶

`DENSITY_STATE_TO_PHYSICAL`
`TEMPERATURE_STATE_TO_PHYSICAL`
`PRESSURE_SOURCE_STATE_TO_MW_M3`
`TRANSPORT_FLUX_MODEL_REGISTRY`
`TRANSPORT_FLUX_MODEL_CAPABILITIES`

Classes¶

`TransportFluxModelBase`	Abstract base class for transport flux models.
`JVPTransportFluxResponse`
`NTXPreparedCoefficientResponse`
`NTXInterpolatedMomentResponse`
`NTXExactLijLaggedResponse`
`FaceTransportState`
`CombinedTransportLaggedResponse`
`CombinedTransportFluxModel`	Abstract base class for transport flux models.
`NTXDatabaseTransportModel`	Abstract base class for transport flux models.
`NTXRuntimeScanChannels`
`NTXExactLijRuntimeSupport`
`NTXRuntimeScanTransportModel`	Abstract base class for transport flux models.
`NTXExactLijRuntimeTransportModel`	Abstract base class for transport flux models.
`ZeroTransportModel`	Abstract base class for transport flux models.
`FluxesRFileTransportModel`	Abstract base class for transport flux models.
`AnalyticalTurbulentTransportModel`	Abstract base class for transport flux models.
`PowerAnalyticalTurbulentTransportModel`	Abstract base class for transport flux models.

Functions¶

`compute_total_power_mw`(state, species, ...[, fallback_mw])
`compute_total_power_breakdown_mw`(state, ...)
`register_transport_flux_model`(→ None)
`get_transport_flux_model`(→ Callable[Ellipsis, ...)
`get_transport_flux_model_capabilities`(...)
`transport_flux_model`(name, **register_kwargs)
`_flatten_flux_dict`(→ tuple[jax.Array, tuple[str, ...)
`_unflatten_flux_dict`(→ dict)
`_extract_right_constraints`(→ tuple[jax.Array, jax.Array])
`_extract_face_constraints`(→ tuple[jax.Array \| None, ...)
`_face_profile`(profile, face_centers[, bc_model, ...])
`_face_profile_gradient`(profile, face_centers[, bc_model])
`build_face_transport_state`(→ FaceTransportState)
`_ntss_like_face_profile`(profile, face_centers[, ...])
`_ntss_like_face_gradient`(profile, face_centers[, bc_model])
`build_ntss_like_face_transport_state`(→ FaceTransportState)
`_as_float_array`(→ jax.Array)
`_import_ntx`()
`_load_ntx_vmec_boozer_channels`(→ dict[str, ...)
`_build_ntx_field_channels`(→ tuple[jax.Array, ...)
`build_ntx_runtime_scan_channels`(→ NTXRuntimeScanChannels)
`_build_ntx_surface_loader`(vmec_file, boozer_file[, ...])
`build_ntx_runtime_surfaces`(→ tuple[Any, Ellipsis])
`build_ntx_exact_lij_runtime_support`(...)
`build_ntx_exact_lij_runtime_transport_model`(species, ...)
`build_ntx_runtime_scan_transport_model`(species, ...[, ...])
`_normalize_flux_dataset`(arr, n_species)
`read_flux_profile_file`(path, n_species)
`_flux_profile_debug_summary`(name, arr)
`build_fluxes_r_file_transport_model`(species, geometry, *)
`build_transport_flux_model`(→ CombinedTransportFluxModel)	Build the composed transport model from explicit model instances.

Module Contents¶

NEOPAX._transport_flux_models.DENSITY_STATE_TO_PHYSICAL = 1e+20¶

NEOPAX._transport_flux_models.TEMPERATURE_STATE_TO_PHYSICAL = 1000.0¶

NEOPAX._transport_flux_models.PRESSURE_SOURCE_STATE_TO_MW_M3 = 0.016019992951203103¶

NEOPAX._transport_flux_models.compute_total_power_mw(state, species, pressure_source_model, geometry, fallback_mw=3.0)¶

NEOPAX._transport_flux_models.compute_total_power_breakdown_mw(state, pressure_source_model, geometry)¶

NEOPAX._transport_flux_models.TRANSPORT_FLUX_MODEL_REGISTRY: dict[str, Callable[[], TransportFluxModelBase]]¶

NEOPAX._transport_flux_models.TRANSPORT_FLUX_MODEL_CAPABILITIES: dict[str, NEOPAX._model_api.ModelCapabilities]¶

NEOPAX._transport_flux_models.register_transport_flux_model(name: str, builder: Callable[Ellipsis, TransportFluxModelBase], *, capabilities: NEOPAX._model_api.ModelCapabilities | None = None, validate: bool = False, validation_context: NEOPAX._model_api.ModelValidationContext | None = None) → None¶

NEOPAX._transport_flux_models.get_transport_flux_model(name: str) → Callable[Ellipsis, TransportFluxModelBase]¶

NEOPAX._transport_flux_models.get_transport_flux_model_capabilities(name: str) → NEOPAX._model_api.ModelCapabilities¶

NEOPAX._transport_flux_models.transport_flux_model(name: str, **register_kwargs)¶

class NEOPAX._transport_flux_models.TransportFluxModelBase¶

Bases: abc.ABC

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

abstractmethod __call__(state, geometry=None, params=None) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

build_lagged_response(state, **kwargs)¶

evaluate_with_lagged_response(state, lagged_response, **kwargs)¶

class NEOPAX._transport_flux_models.JVPTransportFluxResponse¶

reference_state: Any¶

reference_flux: dict¶

class NEOPAX._transport_flux_models.NTXPreparedCoefficientResponse¶

reference_transport_moments: jax.Array¶

reference_nu_hat: jax.Array¶

reference_epsi_hat: jax.Array¶

class NEOPAX._transport_flux_models.NTXInterpolatedMomentResponse¶

reference_er: jax.Array¶

reference_log_nu_star: jax.Array¶

reference_transport_moments: jax.Array¶

dtransport_moments_d_er: jax.Array¶

dtransport_moments_d_log_nu_star: jax.Array¶

class NEOPAX._transport_flux_models.NTXExactLijLaggedResponse¶

center_response: Any¶

class NEOPAX._transport_flux_models.FaceTransportState¶

density: jax.Array¶

pressure: jax.Array¶

Er: jax.Array¶

property temperature¶

NEOPAX._transport_flux_models._flatten_flux_dict(fluxes: dict) → tuple[jax.Array, tuple[str, Ellipsis]]¶

NEOPAX._transport_flux_models._unflatten_flux_dict(flat_flux: jax.Array, reference_flux: dict) → dict¶

NEOPAX._transport_flux_models._extract_right_constraints(bc_model: Any, state_arr: jax.Array) → tuple[jax.Array, jax.Array]¶

NEOPAX._transport_flux_models._extract_face_constraints(bc_model: Any, state_arr: jax.Array, face_centers: jax.Array) → tuple[jax.Array | None, jax.Array | None, jax.Array | None, jax.Array | None]¶

NEOPAX._transport_flux_models._face_profile(profile, face_centers, bc_model=None, reconstruction='linear')¶

NEOPAX._transport_flux_models._face_profile_gradient(profile, face_centers, bc_model=None)¶

NEOPAX._transport_flux_models.build_face_transport_state(state: NEOPAX._state.TransportState, geometry: Any, *, bc_density: Any = None, bc_temperature: Any = None, bc_er: Any = None, reconstruction: str = 'linear', density_floor: Any = DEFAULT_TRANSPORT_DENSITY_FLOOR, temperature_floor: Any = DEFAULT_TRANSPORT_TEMPERATURE_FLOOR) → FaceTransportState¶

NEOPAX._transport_flux_models._ntss_like_face_profile(profile, face_centers, bc_model=None, density_floor=None)¶

NEOPAX._transport_flux_models._ntss_like_face_gradient(profile, face_centers, bc_model=None)¶

NEOPAX._transport_flux_models.build_ntss_like_face_transport_state(state: NEOPAX._state.TransportState, geometry: Any, *, bc_density: Any = None, bc_temperature: Any = None, bc_er: Any = None, density_floor: Any = DEFAULT_TRANSPORT_DENSITY_FLOOR, temperature_floor: Any = DEFAULT_TRANSPORT_TEMPERATURE_FLOOR) → FaceTransportState¶

class NEOPAX._transport_flux_models.CombinedTransportLaggedResponse¶

neoclassical_response: object¶

turbulent_response: object¶

classical_response: object¶

class NEOPAX._transport_flux_models.CombinedTransportFluxModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

neoclassical_model: TransportFluxModelBase¶

turbulent_model: TransportFluxModelBase¶

classical_model: TransportFluxModelBase¶

include_turbulent_particle_flux: bool = True¶

static _zero_like_flux(reference, fallback=0)¶

__call__(state, *args, **kwargs) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

build_lagged_response(state, **kwargs)¶

evaluate_with_lagged_response(state, lagged_response, **kwargs)¶

class NEOPAX._transport_flux_models.NTXDatabaseTransportModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

species: Any¶

energy_grid: Any¶

geometry: Any¶

database: Any¶

collisionality_model: str = 'default'¶

bc_density: Any = None¶

bc_temperature: Any = None¶

__call__(state) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

NEOPAX._transport_flux_models._as_float_array(value, *, name: str, positive: bool = False) → jax.Array¶

NEOPAX._transport_flux_models._import_ntx()¶

NEOPAX._transport_flux_models._load_ntx_vmec_boozer_channels(wout_path: pathlib.Path, boozmn_path: pathlib.Path, rho: jax.Array) → dict[str, jax.Array | float]¶

NEOPAX._transport_flux_models._build_ntx_field_channels(rho: jax.Array, er_tilde: jax.Array, channels: dict[str, jax.Array | float]) → tuple[jax.Array, jax.Array, jax.Array]¶

class NEOPAX._transport_flux_models.NTXRuntimeScanChannels¶

rho: Any¶

a_b: float¶

psia: float¶

b00: Any¶

r00: Any¶

boozer_i: Any¶

boozer_g: Any¶

iota: Any¶

drds: Any¶

dr_tildedr: Any¶

dr_tildeds: Any¶

fac_reference_to_sfincs_11: Any¶

fac_reference_to_sfincs_31: Any¶

fac_reference_to_sfincs_33: Any¶

fac_sfincs_to_dkes_11: Any¶

fac_sfincs_to_dkes_31: Any¶

fac_sfincs_to_dkes_33: Any¶

fac_dkes_to_d11star: Any¶

fac_dkes_to_d31star: Any¶

fac_dkes_to_d33star: Any¶

classmethod from_mapping(rho, channels: dict[str, jax.Array | float]) → NTXRuntimeScanChannels¶

as_mapping() → dict[str, jax.Array | float]¶

NEOPAX._transport_flux_models.build_ntx_runtime_scan_channels(vmec_file, boozer_file, rho_scan) → NTXRuntimeScanChannels¶

NEOPAX._transport_flux_models._build_ntx_surface_loader(vmec_file, boozer_file, surface_backend='auto')¶

NEOPAX._transport_flux_models.build_ntx_runtime_surfaces(vmec_file, boozer_file, rho_values, *, surface_backend='auto') → tuple[Any, Ellipsis]¶

class NEOPAX._transport_flux_models.NTXExactLijRuntimeSupport¶

center_channels: NTXRuntimeScanChannels¶

face_channels: NTXRuntimeScanChannels¶

center_prepared: Any¶

face_prepared: Any¶

grid: Any¶

class NEOPAX._transport_flux_models.NTXRuntimeScanTransportModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

species: Any¶

energy_grid: Any¶

geometry: Any¶

vmec_file: str | None¶

boozer_file: str | None¶

rho_scan: Any¶

nu_v_scan: Any¶

er_tilde_scan: Any¶

n_theta: int = 25¶

n_zeta: int = 25¶

n_xi: int = 64¶

surface_backend: str = 'vmec'¶

source_name: str = 'ntx_scan_runtime'¶

collisionality_model: str = 'default'¶

bc_density: Any = None¶

bc_temperature: Any = None¶

channels: NTXRuntimeScanChannels | None = None¶

database: Any = None¶

_scan_axes() → tuple[jax.Array, jax.Array, jax.Array]¶

_static_channels() → NTXRuntimeScanChannels¶

_surface_loader(ntx)¶

_build_runtime_database()¶

with_static_channels() → NTXRuntimeScanTransportModel¶

with_scan_inputs(*, rho_scan=None, nu_v_scan=None, er_tilde_scan=None, clear_database: bool = True) → NTXRuntimeScanTransportModel¶

_database_model() → NTXDatabaseTransportModel¶

__call__(state) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

with_runtime_database() → NTXRuntimeScanTransportModel¶

NEOPAX._transport_flux_models.build_ntx_exact_lij_runtime_support(vmec_file, boozer_file, rho_center, rho_face, *, surface_backend='auto', n_theta=25, n_zeta=25, n_xi=64) → NTXExactLijRuntimeSupport¶

class NEOPAX._transport_flux_models.NTXExactLijRuntimeTransportModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

species: Any¶

energy_grid: Any¶

geometry: Any¶

vmec_file: str | None¶

boozer_file: str | None¶

n_theta: int = 25¶

n_zeta: int = 25¶

n_xi: int = 64¶

surface_backend: str = 'vmec'¶

face_response_mode: str = 'face_local_response'¶

radial_batch_size: int | None = None¶

radial_batch_mode: str = 'simple'¶

scan_batch_size: int | None = None¶

response_anchor_count: int | None = None¶

use_remat: bool = False¶

er_v_floor: float | None = None¶

collisionality_model: str = 'default'¶

bc_density: Any = None¶

bc_temperature: Any = None¶

support: NTXExactLijRuntimeSupport | None = None¶

_rho_center_face()¶

_static_support() → NTXExactLijRuntimeSupport¶

with_static_support() → NTXExactLijRuntimeTransportModel¶

with_transport_resolution(*, n_theta=None, n_zeta=None, n_xi=None) → NTXExactLijRuntimeTransportModel¶

with_face_response_mode(face_response_mode: str) → NTXExactLijRuntimeTransportModel¶

with_radial_batch_size(radial_batch_size: int | None) → NTXExactLijRuntimeTransportModel¶

with_radial_batch_mode(radial_batch_mode: str | None) → NTXExactLijRuntimeTransportModel¶

with_scan_batch_size(scan_batch_size: int | None) → NTXExactLijRuntimeTransportModel¶

with_response_anchor_count(response_anchor_count: int | None) → NTXExactLijRuntimeTransportModel¶

with_use_remat(use_remat: bool) → NTXExactLijRuntimeTransportModel¶

with_er_v_floor(er_v_floor: float | None) → NTXExactLijRuntimeTransportModel¶

static _normalize_radial_batch_mode(radial_batch_mode: str | None) → str¶

_map_radius_axis_hybrid(fn, radius_indices)¶

_map_radius_axis(fn, radius_indices)¶

_map_radius_axis_unbatched(fn, radius_indices)¶

_response_anchor_indices(n_radius: int) → jax.Array¶

_interpolate_anchor_values(anchor_indices, anchor_values, target_rho)¶

_regularize_axis_radius0(values_by_radius, radius_coordinates)¶

_map_radius_axis_regularized_at_axis0(fn, radius_indices, radius_coordinates, *, unbatched: bool = False)¶

_regularize_center_fluxes_axis0(gamma, q, upar)¶

_log_nu_star_from_nu_hat(nu_hat_a)¶

_local_scan_inputs(*, drds_value, species_index: int, er_value, temperature_local, density_local, vthermal_local, collisionality_kind)¶

_lij_from_coefficient_scan(coeff_scan, *, drds_value, species_index: int, vth_a)¶

_transport_moments_from_coefficient_scan(coeff_scan, *, drds_value)¶

_transport_moments_from_inputs_impl(prepared, nu_hat_a, epsi_hat_a, *, drds_value)¶

_transport_moments_from_inputs(prepared, nu_hat_a, epsi_hat_a, *, drds_value)¶

_lij_from_transport_moments(transport_moments, *, species_index: int, vth_a)¶

_batched_lij_from_transport_moments(transport_moments, v_thermal)¶

_solve_coefficient_scan_prepared_impl(prepared, nu_hat_a, epsi_hat_a)¶

_solve_coefficient_scan_prepared(prepared, nu_hat_a, epsi_hat_a)¶

_coefficient_scan_from_inputs(prepared, nu_hat_a, epsi_hat_a)¶

_solve_lij_prepared_local_impl(prepared, *, drds_value, species_index: int, er_value, temperature_local, density_local, vthermal_local, collisionality_kind)¶

_solve_lij_prepared_local(prepared, *, drds_value, species_index: int, er_value, temperature_local, density_local, vthermal_local, collisionality_kind)¶

_build_coefficient_response_local(prepared, *, drds_value, species_index: int, er_value, temperature_local, density_local, vthermal_local, collisionality_kind)¶

_build_interpolated_moment_response_local(prepared, *, drds_value, species_index: int, er_value, temperature_local, density_local, vthermal_local, collisionality_kind)¶

_lij_center(Er, temperature, density)¶

_lij_faces(Er_faces, temperature_faces, density_faces)¶

_assemble_center_fluxes(Er, temperature, density, lij, n_right, n_right_grad, t_right, t_right_grad)¶

_cell_centered_flux_to_faces_centered(flux)¶

__call__(state) → dict¶

build_lagged_response(state, **kwargs)¶

evaluate_with_lagged_response(state, lagged_response, **kwargs)¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

NEOPAX._transport_flux_models.build_ntx_exact_lij_runtime_transport_model(species, energy_grid, geometry, *, vmec_file, boozer_file, ntx_exact_n_theta=25, ntx_exact_n_zeta=25, ntx_exact_n_xi=64, ntx_exact_surface_backend='vmec', ntx_exact_face_response_mode='face_local_response', ntx_exact_radial_batch_size=None, ntx_exact_radial_batch_mode='simple', ntx_exact_scan_batch_size=None, ntx_exact_response_anchor_count=None, ntx_exact_use_remat=False, ntx_exact_er_v_floor=None, ntx_exact_lij_support=None, preload_support=False, collisionality_model='default', bc_density=None, bc_temperature=None, **kwargs)¶

NEOPAX._transport_flux_models.build_ntx_runtime_scan_transport_model(species, energy_grid, geometry, *, vmec_file, boozer_file, ntx_scan_rho, ntx_scan_nu_v, ntx_scan_er_tilde, ntx_scan_n_theta=25, ntx_scan_n_zeta=25, ntx_scan_n_xi=64, ntx_scan_surface_backend='auto', ntx_scan_source_name='ntx_scan_runtime', collisionality_model='default', bc_density=None, bc_temperature=None, ntx_scan_channels=None, preload_channels=False, prebuild_database=True, **kwargs)¶

class NEOPAX._transport_flux_models.ZeroTransportModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

shape: Any = None¶

__call__(state) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

build_lagged_response(state, **kwargs)¶

NEOPAX._transport_flux_models._normalize_flux_dataset(arr, n_species)¶

NEOPAX._transport_flux_models.read_flux_profile_file(path, n_species)¶

NEOPAX._transport_flux_models._flux_profile_debug_summary(name, arr)¶

NEOPAX._transport_flux_models.build_fluxes_r_file_transport_model(species, geometry, *, fluxes_file=None, file=None, flux_file=None, neoclassical_file=None, turbulence_file=None, classical_file=None, grid_location='cell_centered', profile_location=None, **kwargs)¶

class NEOPAX._transport_flux_models.FluxesRFileTransportModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

species: Any¶

geometry: Any¶

r_data: Any¶

gamma_data: Any = None¶

q_data: Any = None¶

upar_data: Any = None¶

profile_location: str = 'cell_centered'¶

q_scale: float = 1.0¶

with_q_scale(q_scale: float) → FluxesRFileTransportModel¶

_interp_species_profile(data, target_r)¶

_normalize_profile_location()¶

_data_on_cell_grid(data)¶

_data_on_face_grid(data)¶

__call__(state) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

class NEOPAX._transport_flux_models.AnalyticalTurbulentTransportModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

species: Any¶

grid: Any¶

chi_t: Any¶

chi_n: Any¶

field: Any¶

with_transport_coeffs(*, chi_t=None, chi_n=None) → AnalyticalTurbulentTransportModel¶

__call__(state) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

build_lagged_response(state, **kwargs)¶

evaluate_with_lagged_response(state, lagged_response, **kwargs)¶

class NEOPAX._transport_flux_models.PowerAnalyticalTurbulentTransportModel¶

Bases: TransportFluxModelBase

Abstract base class for transport flux models. Output dict keys:

Gamma: particle flux

Q: heat flux

Upar: parallel flow

species: Any¶

field: Any¶

chi_t: Any¶

chi_n: Any¶

pressure_source_model: Any = None¶

total_power_mw: Any = None¶

with_transport_coeffs(*, chi_t=None, chi_n=None, pressure_source_model=None, total_power_mw=None) → PowerAnalyticalTurbulentTransportModel¶

_effective_total_power_mw(state)¶

__call__(state) → dict¶

build_local_particle_flux_evaluator(state)¶

evaluate_face_fluxes(state, face_state, **kwargs)¶

build_lagged_response(state, **kwargs)¶

evaluate_with_lagged_response(state, lagged_response, **kwargs)¶

NEOPAX._transport_flux_models.build_transport_flux_model(neo_model: TransportFluxModelBase, turb_model: TransportFluxModelBase, classical_model: TransportFluxModelBase = None, *, include_turbulent_particle_flux: bool = True) → CombinedTransportFluxModel¶: Build the composed transport model from explicit model instances. All models must be constructed up front by the orchestrator.