NEOPAX._ambipolarity¶

Attributes¶

AMBIPOLARITY_MODEL_REGISTRY

Functions¶

`_normalize_er_scan_batch_mode`(→ str)
`_map_er_scan_axis`(fn, values, *[, batch_mode, batch_size])
`find_ambipolar_Er_min_entropy_jit`(Gamma_func, entropy_func)	JIT-friendly ambipolar root search without Python loops.
`find_ambipolar_Er_multistart_clustered`(Gamma_func[, ...])	Multi-start root-finding for ambipolarity using least-squares and clustering.
`find_ambipolar_Er_min_entropy_jit_multires`(Gamma_func, ...)	Two-stage (coarse/fine) JIT-friendly ambipolar root search for a single radius.
`find_ambipolar_Er_min_entropy_jit_multires_continuation`(...)	Continuation upgrade of the scalar two-stage search.
`initialize_ambipolar_best_roots_fast`(state, config, ...)	Fast ambipolar best-root profile initializer.
`find_ambipolar_root_tracked_local`(Gamma_func, prev_root)	Cheap single-root tracker around a previous best root.
`_build_initializer_evaluators`(state, params, flux_model)
`initialize_ambipolar_best_roots_tracked`(state, config, ...)	Faster initialization-oriented ambipolar best-root tracker.
`initialize_ambipolar_best_roots_hybrid`(state, config, ...)	Hybrid initialization:
`_deduplicate_sorted_roots`(roots, merge_tol)
`initialize_ambipolar_best_roots_multibranch`(state, ...)	Multibranch continuation initializer.
`find_ambipolar_Er_min_entropy_jit_adaptive`(Gamma_func, ...)	JAX-compatible adaptive bracketing root-finder with static padding.
`pad_and_sort_roots_for_plotting`(roots_all, ...[, ...])	For each radius, sort roots and entropies, pad to max_roots (default 3), and return arrays for plotting.
`plot_roots`(→ pathlib.Path)
`write_ambipolarity_hdf5`(→ pathlib.Path)
`register_ambipolarity_model`(name, func)
`solve_ambipolarity_roots_radial`(state, config, params, ...)	JIT/differentiable radial root-finder with blocksize option (default: pure vmap).
`solve_ambipolarity_roots_from_config`(state, config, params)	Config-driven entrypoint for ambipolarity root finding. Reads config, builds models, and calls the JIT/diff radial solver.

Module Contents¶

NEOPAX._ambipolarity._normalize_er_scan_batch_mode(mode: str | None) → str¶

NEOPAX._ambipolarity._map_er_scan_axis(fn, values, *, batch_mode: str = 'vmap', batch_size: int | None = None)¶

NEOPAX._ambipolarity.find_ambipolar_Er_min_entropy_jit(Gamma_func, entropy_func, Er_range=(-20.0, 20.0), n_scan=96, tol=1e-06, x_tol=1e-06, maxiter=30, er_scan_batch_mode='vmap', er_scan_batch_size=None)¶

JIT-friendly ambipolar root search without Python loops.

Uses a batched safeguarded Newton-bisection update over all scan intervals, then selects the minimum-entropy valid root.

Returns:: scalar Er minimizing entropy among valid roots roots_all: candidate roots for each interval (shape n_scan-1) entropies_all: entropy values with inf for invalid intervals valid_mask: boolean mask for valid intervals
Return type:: best_root

NEOPAX._ambipolarity.find_ambipolar_Er_multistart_clustered(Gamma_func, entropy_func=None, Er_range=(-20.0, 20.0), n_starts=32, tol=1e-06, maxiter=30, cluster_tol=0.001, return_entropies=False)¶: Multi-start root-finding for ambipolarity using least-squares and clustering. Uses many initial guesses, runs a root finder (optimistix if available), and clusters roots. Returns unique roots (within cluster_tol).

NEOPAX._ambipolarity.find_ambipolar_Er_min_entropy_jit_multires(Gamma_func, entropy_func, Er_range=(-20.0, 20.0), n_coarse=24, n_refine=8, max_roots=3, tol=1e-06, x_tol=1e-06, maxiter=12, er_scan_batch_mode='vmap', er_scan_batch_size=None)¶: Two-stage (coarse/fine) JIT-friendly ambipolar root search for a single radius. 1. Coarse scan to bracket roots (n_coarse) 2. Refine each bracket with Newton-bisection (n_refine) 3. Pad to max_roots for static shape

NEOPAX._ambipolarity.find_ambipolar_Er_min_entropy_jit_multires_continuation(Gamma_func, entropy_func, prev_root, global_Er_range=(-20.0, 20.0), local_half_width=6.0, expand_factor=2.0, max_local_expands=2, fallback_global=True, n_coarse=24, n_refine=8, max_roots=3, tol=1e-06, x_tol=1e-06, maxiter=12)¶

Continuation upgrade of the scalar two-stage search.

Start from a narrow range around the previous radial best root and only widen/fallback if the local search fails. This preserves the JAX-friendly scalar solver while avoiding a broad global scan at every radius.

NEOPAX._ambipolarity.initialize_ambipolar_best_roots_fast(state, config, params, flux_model, amb_cfg)¶

Fast ambipolar best-root profile initializer.

Optimized for initialization only: - robust global solve on the first radius - cheap local continuation search on later radii - optional fallback to global search only when local search fails

Returns:: (n_radial,) ndarray
Return type:: best_roots

NEOPAX._ambipolarity.find_ambipolar_root_tracked_local(Gamma_func, prev_root, global_Er_range=(-20.0, 20.0), local_half_width=4.0, n_scan=12, tol=1e-06, x_tol=1e-06, maxiter=10)¶

Cheap single-root tracker around a previous best root.

This is intended for initialization only: - scan a narrow local window - pick the bracket closest to the previous root - refine a single root with safeguarded Newton/bisection

NEOPAX._ambipolarity._build_initializer_evaluators(state, params, flux_model)¶

NEOPAX._ambipolarity.initialize_ambipolar_best_roots_tracked(state, config, params, flux_model, amb_cfg)¶

Faster initialization-oriented ambipolar best-root tracker.

Strategy: - first radius: robust global two-stage + entropy selection - later radii: track the selected branch locally using a cheap sign-scan and

single-root refinement

fallback to global two-stage only when local tracking fails

NEOPAX._ambipolarity.initialize_ambipolar_best_roots_hybrid(state, config, params, flux_model, amb_cfg)¶

Hybrid initialization: - use cheap tracked local continuation most radii - periodically or when suspicious, revalidate with full two_stage

This keeps the baseline two_stage logic in the loop, but avoids paying for it at every radius.

NEOPAX._ambipolarity._deduplicate_sorted_roots(roots, merge_tol)¶

NEOPAX._ambipolarity.initialize_ambipolar_best_roots_multibranch(state, config, params, flux_model, amb_cfg)¶

Multibranch continuation initializer.

first radius: full two_stage to get all candidate roots
later radii: track each previously found branch locally
choose minimum-entropy among tracked branches
periodically or on failure, fall back to full two_stage

NEOPAX._ambipolarity.find_ambipolar_Er_min_entropy_jit_adaptive(Gamma_func, entropy_func, Er_range=(-20.0, 20.0), n_init=16, n_subdiv=2, n_rounds=2, max_brackets=24, n_refine=8, max_roots=3, tol=1e-06, x_tol=1e-06, maxiter=12, er_scan_batch_mode='vmap', er_scan_batch_size=None)¶: JAX-compatible adaptive bracketing root-finder with static padding. - n_init: initial number of coarse intervals - n_subdiv: number of subdivisions per round (e.g., 2 for bisection) - n_rounds: number of adaptive subdivision rounds - max_brackets: maximum number of brackets to pad for JIT - n_refine: Newton-bisection steps per bracket - max_roots: maximum roots to pad for output

NEOPAX._ambipolarity.pad_and_sort_roots_for_plotting(roots_all, entropies_all, n_roots_all, best_roots=None, max_roots=3)¶

For each radius, sort roots and entropies, pad to max_roots (default 3), and return arrays for plotting. roots_all: (n_radial, n_found) array entropies_all: (n_radial, n_found) array n_roots_all: (n_radial,) array (number of valid roots per radius) best_roots: (n_radial,) array (optional) :returns: (max_roots, n_radial)

entropies_3: (max_roots, n_radial) best_root: (n_radial,) if best_roots is not None else None

Return type:: roots_3

NEOPAX._ambipolarity.plot_roots(rho, roots_3, entropies_3, best_root, output_dir: pathlib.Path, overlay_reference_er: bool = True, reference_er_file: str | pathlib.Path | None = None, reference_er_label: str | None = None, debug: bool = False) → pathlib.Path¶

NEOPAX._ambipolarity.write_ambipolarity_hdf5(rho, roots_3, entropies_3, best_root, output_dir: pathlib.Path) → pathlib.Path¶

NEOPAX._ambipolarity.AMBIPOLARITY_MODEL_REGISTRY¶

NEOPAX._ambipolarity.register_ambipolarity_model(name: str, func)¶

NEOPAX._ambipolarity.solve_ambipolarity_roots_radial(state, config, params, model_name, flux_model, entropy_model, amb_cfg)¶: JIT/differentiable radial root-finder with blocksize option (default: pure vmap).

NEOPAX._ambipolarity.solve_ambipolarity_roots_from_config(state, config, params, flux_model=None, entropy_model=None)¶: Config-driven entrypoint for ambipolarity root finding. Reads config, builds models, and calls the JIT/diff radial solver.