Profiles

NTX now exposes a first imported profile workflow in src/ntx/profiles.py. This layer sits above the monoenergetic solve and above the radial scan builders, and it is intended for ambipolar electric-field studies and reduced bootstrap-current response analysis.

Scope

This module does not replace a full multi-species transport code. It uses the monoenergetic transport coefficients already produced by NTX and builds a clean, differentiable profile-level closure around them.

The current closure uses:

  • reduced monoenergetic particle-flux responses

  • reduced monoenergetic parallel-current responses

  • a smooth ambipolar electric-field profile solve on a precomputed NTX scan

Main Objects

MonoenergeticSpeciesProfile

One species is described by:

  • charge

  • nu_v

  • A1

  • A3

  • particle_weight

  • current_weight

where A1(r) and A3(r) are the thermodynamic-force channels used in the monoenergetic closure.

AmbipolarProfileResult

The solver returns:

  • rho

  • er_profile

  • ambipolar_residual

  • bootstrap_current_response

  • species_particle_flux

  • species_current_response

  • loss_history

The read-only bootstrap_current_proxy attribute remains as an NTX 0.2.x compatibility alias. New code should use bootstrap_current_response.

Reduced Monoenergetic Model

For one species, NTX currently uses the monoenergetic closures

\[\Gamma_a(r) = -w^{(\Gamma)}_a(r)\left[D_{11,a}(r) A_{1,a}(r) + D_{13,a}(r) A_{3,a}(r)\right],\]
\[J_a(r) = -w^{(J)}_a(r)\left[D_{31,a}(r) A_{1,a}(r) + D_{33,a}(r) A_{3,a}(r)\right].\]

The ambipolar residual is then

\[R(r) = \sum_a Z_a \Gamma_a(r),\]

so a charge-symmetric pair with identical particle-flux response must cancel exactly. The fast physics-gate suite checks this local ambipolarity identity before the nonlinear radial-electric-field solve is trusted.

and the reduced current response is

\[J_{\mathrm{red}}(r) = \sum_a J_a(r).\]

The ambipolar electric-field profile is obtained by minimizing a smooth radial objective on the precomputed E_r scan stored in the NTX scan payload:

\[\mathcal L_E[\hat E_r] = \left\langle R(r;\hat E_r)^2 \right\rangle_r + \lambda_E \left\langle \left(\partial_r \hat E_r\right)^2 + \tfrac12\left(\partial_r^2 \hat E_r\right)^2 \right\rangle_r,\]

where \lambda_E is the user-controlled smoothing weight exposed through smoothing_strength. NTX then applies bounded backtracking updates to the full profile vector rather than solving each radius independently. This suppresses the checkerboard artifacts that appear when adjacent radii are updated without any radial regularization.

Primitive density and temperature updates use exponential relaxation and a strict positive floor. That preserves the physical state-space constraint n(r) > 0, T(r) > 0 even when a transport mismatch is large enough to underflow an unconstrained explicit update.

Main Helpers

  • evaluate_scan_channel(...)

  • evaluate_species_particle_flux(...)

  • evaluate_species_current_response(...)

  • ambipolar_residual_profile(...)

  • solve_ambipolar_er_profile(...)

  • solve_ambipolar_profile_family(...)

  • current_response_objective(...)

  • bootstrap_current_objective(...)

  • apply_profile_control(...)

  • optimize_profile_control(...)

  • ProfileBasisControlSpec

  • apply_profile_basis_control(...)

  • optimize_profile_basis_control(...)

  • ProfileTransportClosureSpec

  • profile_transport_loss(...)

  • advance_profile_transport(...)

  • solve_profile_transport_loop(...)

Typical Workflow

import jax.numpy as jnp
from ntx import (
    GridSpec,
    MonoenergeticSpeciesProfile,
    build_ntx_neopax_scan_from_surfaces,
    example_surface,
    solve_ambipolar_er_profile,
)

rho = jnp.linspace(0.2, 0.8, 6)
nu_v = jnp.asarray([3.0e-4, 1.0e-3, 3.0e-3, 1.0e-2])
er_axis = jnp.asarray([-3.0e-3, -1.0e-3, -3.0e-4, 0.0, 3.0e-4, 1.0e-3, 3.0e-3])
er_grid = jnp.tile(er_axis[None, :], (rho.size, 1))
surfaces = tuple(example_surface() for _ in range(rho.size))

scan = build_ntx_neopax_scan_from_surfaces(
    surfaces,
    rho=rho,
    nu_v=nu_v,
    Es=er_grid,
    Er=er_grid,
    drds=jnp.ones_like(rho),
    grid=GridSpec(7, 9, 6),
)

electron = MonoenergeticSpeciesProfile(
    charge=-1.0,
    nu_v=jnp.linspace(4.0e-4, 1.0e-3, rho.size),
    A1=1.1 - 0.25 * rho,
    A3=0.55 - 0.12 * rho,
    current_weight=-1.0,
    name="electron",
)
ion = MonoenergeticSpeciesProfile(
    charge=1.0,
    nu_v=jnp.linspace(2.0e-3, 5.0e-3, rho.size),
    A1=0.7 + 0.35 * rho,
    A3=0.24 + 0.08 * rho,
    particle_weight=1.08,
    current_weight=1.0,
    name="ion",
)

result = solve_ambipolar_er_profile(scan, (electron, ion), steps=12, damping=0.7)

Example Script

The repository example

python examples/ambipolar_profile.py

writes:

docs/_static/ambipolar_profile.png
docs/_static/ambipolar_profile.pdf

It shows:

  • the ambipolar residual landscape over the scanned E_r axis

  • the reduced bootstrap-current response profile

  • species particle-flux responses and the charge-weighted residual

  • the integrated ambipolar landscape used by the smooth-profile solver

Ambipolar profile

Control-Parameter Families

NTX also exposes a small family-solve layer:

\[\mathcal J(c) = \int w(r) J_{\mathrm{red}}(r;c)^2\,dr,\]

where c is any explicit profile control and w(r) is an optional radial weight.

Use:

  • solve_ambipolar_profile_family(...) to solve several profile closures on the same NTX scan

  • current_response_objective(...) to reduce one solved current profile to a scalar optimization objective; bootstrap_current_objective(...) remains as a compatibility wrapper

The repository example

python examples/ambipolar_profile_family.py

writes:

docs/_static/ambipolar_profile_family.png
docs/_static/ambipolar_profile_family.pdf

It shows:

  • the integrated residual landscape across the control family

  • the resulting family of reduced bootstrap-current responses

  • a scalar objective landscape across the control parameter

  • the final ambipolar residual norm across that family

Ambipolar profile family

Differentiable Profile-Control Optimization

On top of the family solve, NTX now exposes a scalar control optimization:

\[\mathcal J(c) = \int w(r) J_{\mathrm{red}}(r;c)^2\,dr + \lambda \left\langle R(r;c)^2 \right\rangle,\]

where c is a scalar profile control, w(r) is an optional radial weight, and \lambda is a residual penalty.

The corresponding helpers are:

  • ProfileControlSpec

  • apply_profile_control(...)

  • optimize_profile_control(...)

The scalar-control implementation lives in src/ntx/_profiles_control_scalar.py; the compatibility facade src/ntx/_profiles_controls.py preserves the existing public import surface. The fast test suite gates the intended linear response: zero control leaves A1 and A3 unchanged, and finite control multiplies each species profile by the prescribed response factor.

The repository example

python examples/profile_control_optimization.py

writes:

docs/_static/profile_control_optimization.png
docs/_static/profile_control_optimization.pdf

It shows:

  • objective descent across optimization iterations

  • scalar control updates

  • the residual-profile reduction relative to the uncontrolled baseline

  • the best reduced bootstrap-current response profile

The implementation lives entirely in src/ntx/profiles.py, so the optimization stays in the imported JAX lane instead of leaving the NTX runtime.

Profile control optimization

Low-Dimensional Radial Basis Controls

For richer profile optimization studies, NTX also exposes a low-dimensional radial basis control:

\[\delta A_{1,a}(r) = \sum_k c_k R^{(1)}_{a,k}\phi_k(r), \qquad \delta A_{3,a}(r) = \sum_k c_k R^{(3)}_{a,k}\phi_k(r),\]

where:

  • c_k are the optimized basis amplitudes

  • \phi_k(r) are user-supplied radial basis functions

  • R^{(1)}_{a,k} and R^{(3)}_{a,k} map each basis function into each species

The profile objective then becomes

\[\mathcal J(\mathbf c) = \int w(r) J_{\mathrm{red}}(r;\mathbf c)^2\,dr + \lambda \left\langle R(r;\mathbf c)^2 \right\rangle + \mu \|\mathbf c\|_2^2.\]

In the current implementation, the optimization step is stabilized in two ways:

  • the control update uses a normalized gradient direction rather than the raw gradient magnitude

  • a small backtracking line search rejects steps that increase the objective

This keeps the public examples in a physically interpretable regime even when the reduced bootstrap-current response changes rapidly with control amplitude.

The corresponding helpers are:

  • ProfileBasisControlSpec

  • apply_profile_basis_control(...)

  • optimize_profile_basis_control(...)

The radial-basis implementation lives in src/ntx/_profiles_control_basis.py. The matching gate checks that the basis-control modifier is exactly the contracted response-basis map, with zero control again preserving the original thermodynamic-force profiles.

The repository example

python examples/profile_basis_optimization.py

writes:

docs/_static/profile_basis_optimization.png
docs/_static/profile_basis_optimization.pdf
docs/_static/profile_basis_optimization.json

It shows:

  • the basis-coefficient history

  • the basis functions and the final optimized modifier

  • the residual-profile reduction relative to the uncontrolled baseline

  • the optimized reduced bootstrap-current response profile

The JSON sidecar records the objective improvement, residual norm ratio, and optimized basis coefficients so this profile-basis workflow can be tracked as a stress artifact rather than only as a picture.

Profile basis optimization

Profile Transport Relaxation Loop

The next step beyond explicit control families is a simple self-consistent transport-relaxation loop. NTX now exposes

  • ProfileTransportClosureSpec

  • profile_transport_loss(...)

  • advance_profile_transport(...)

  • solve_profile_transport_loop(...)

The closure updates the profile-force channels explicitly after each ambipolar solve, using prescribed source and target channels:

\[A_{1,a}^{(n+1)}(r) = A_{1,a}^{(n)}(r) - \alpha^{(\Gamma)}_a(r)\, \widetilde{\Delta\Gamma}_a^{(n)}(r),\]
\[A_{3,a}^{(n+1)}(r) = A_{3,a}^{(n)}(r) - \alpha^{(J)}_a(r)\, \widetilde{\Delta J}_a^{(n)}(r),\]

where the normalized mismatches are

\[\Delta\Gamma_a^{(n)}(r) = \Gamma_a^{(n)}(r) - \Gamma_{a,\mathrm{target}}(r) - \Gamma_{a,\mathrm{source}}(r),\]
\[\Delta J_a^{(n)}(r) = J_a^{(n)}(r) - J_{a,\mathrm{target}}(r) - J_{a,\mathrm{source}}(r),\]
\[\widetilde{\Delta\Gamma}_a^{(n)}(r) = \mathrm{clip}\!\left( \frac{\Delta\Gamma_a^{(n)}(r)} {\max\!\bigl(\sqrt{\langle (\Delta\Gamma_a^{(n)})^2\rangle_r},\,\epsilon_a^{(\Gamma)}(r)\bigr)}, -u_{\max,a}^{(\Gamma)}(r), u_{\max,a}^{(\Gamma)}(r) \right),\]
\[\widetilde{\Delta J}_a^{(n)}(r) = \mathrm{clip}\!\left( \frac{\Delta J_a^{(n)}(r)} {\max\!\bigl(\sqrt{\langle (\Delta J_a^{(n)})^2\rangle_r},\,\epsilon_a^{(J)}(r)\bigr)}, -u_{\max,a}^{(J)}(r), u_{\max,a}^{(J)}(r) \right).\]

The loop records a quadratic transport mismatch loss,

\[\mathcal L_{\mathrm{transport}}^{(n)} = \left\langle \sum_a \left(\Delta\Gamma_a^{(n)}\right)^2 + \left(\Delta J_a^{(n)}\right)^2 \right\rangle_r,\]

Each explicit update is then checked with a short backtracking acceptance rule: the next state is kept only if the recomputed transport loss does not increase. This makes the shipped workflow materially stronger than the older raw-response relaxation loop and removes the large runaway profiles that were easy to trigger in coarse public examples.

The closure-spec fields are therefore:

  • particle_relaxation

  • current_relaxation

  • particle_target

  • current_target

  • particle_source

  • current_source

  • normalization_floor

  • max_normalized_update

  • radial_smoothing_strength

The simple default update is still explicit, but it is now much more robust for publication-facing profile studies.

The corresponding helpers are:

  • ProfileTransportClosureSpec

  • profile_transport_loss(...)

  • advance_profile_transport(...)

  • solve_profile_transport_loop(...)

  • PrimitiveSpeciesProfile

  • build_species_profile_from_primitives(...)

  • build_species_profiles_from_primitives(...)

  • primitive_profile_transport_loss(...)

  • advance_primitive_profile_transport(...)

  • solve_primitive_profile_transport_loop(...)

The old raw-update form,

\[A_{1,a}^{(n+1)}(r) = A_{1,a}^{(n)}(r) - \alpha^{(\Gamma)}_a(r)\left[\Gamma_a^{(n)}(r)-\Gamma_{a,\mathrm{target}}(r)\right],\]
\[A_{3,a}^{(n+1)}(r) = A_{3,a}^{(n)}(r) - \alpha^{(J)}_a(r)\left[J_a^{(n)}(r)-J_{a,\mathrm{target}}(r)\right].\]

is retained only as intuition. The shipped implementation uses the normalized source/target mismatch form above.

The repository example

python examples/profile_transport_loop.py

writes:

docs/_static/profile_transport_loop.png
docs/_static/profile_transport_loop.pdf

It shows:

  • the ambipolar residual evolution across accepted transport iterations

  • the corresponding reduced bootstrap-current response history

  • the transport-loss and ambipolar-residual histories

  • the final A1(r) and A3(r) profiles for each species

Profile transport loop

Primitive Density And Temperature Transport

NTX also now supports a stronger imported workflow in which the thermodynamic forces are reconstructed from primitive density and temperature profiles rather than updated directly. For one species,

\[A_{3,a}(r) = \frac{d \ln T_a}{dr},\]
\[A_{1,a}(r) = \frac{d \ln n_a}{dr} - \frac{3}{2}\frac{d \ln T_a}{dr} + C_{E,a}(r) Z_a E_r(r),\]

where C_{E,a}(r) is the user-supplied electrostatic prefactor. This mapping is covered by a fast analytical physics gate and implemented by:

  • PrimitiveSpeciesProfile

  • build_species_profile_from_primitives(...)

  • build_species_profiles_from_primitives(...)

The primitive closure uses the same normalized transport mismatches as the explicit A1/A3 loop, augments them with explicit density and temperature source-target channels, and updates the primitive fields multiplicatively:

\[n_a^{(n+1)}(r) = n_a^{(n)}(r) \exp\!\left[ -\alpha_a^{(\Gamma)}(r)\widetilde{\Delta\Gamma}_a^{(n)}(r) -\alpha_a^{(n)}(r)\widetilde{\Delta n}_a^{(n)}(r) \right],\]
\[T_a^{(n+1)}(r) = T_a^{(n)}(r) \exp\!\left[ -\alpha_a^{(J)}(r)\widetilde{\Delta J}_a^{(n)}(r) -\alpha_a^{(T)}(r)\widetilde{\Delta T}_a^{(n)}(r) \right].\]

Here

\[\Delta n_a^{(n)}(r) = n_a^{(n)}(r) - n_{a,\mathrm{target}}(r) - n_{a,\mathrm{source}}(r), \qquad \Delta T_a^{(n)}(r) = T_a^{(n)}(r) - T_{a,\mathrm{target}}(r) - T_{a,\mathrm{source}}(r),\]

with the same normalized-and-clipped structure used for the monoenergetic transport channels. NTX then applies radial smoothing to the updated primitive profiles before reconstructing A1(r) and A3(r) for the next ambipolar solve. This keeps density and temperature positive, suppresses coarse-grid spikes, and still feeds their gradients back into the ambipolar closure through A1(r) and A3(r).

The additional primitive-closure fields on ProfileTransportClosureSpec are:

  • density_relaxation

  • temperature_relaxation

  • density_target

  • temperature_target

  • density_source

  • temperature_source

  • primitive_normalization_floor

  • max_primitive_normalized_update

  • radial_smoothing_strength

The repository example

python examples/primitive_profile_transport.py

writes:

docs/_static/primitive_profile_transport.png
docs/_static/primitive_profile_transport.pdf

and shows:

  • the initial-versus-final primitive ambipolar residual and current profiles

  • the derived monoenergetic force profiles reconstructed from the primitive state

  • final density and temperature profiles for each species

Primitive profile transport

Literature-Anchored Primitive-To-Force Audit

The repository also includes a benchmark-family audit for the primitive profile reconstruction itself:

python examples/profile_force_reconstruction_audit.py

This writes:

docs/_static/profile_force_reconstruction_audit.png
docs/_static/profile_force_reconstruction_audit.pdf
docs/_static/profile_force_reconstruction_audit.json

It compares the reconstructed

\[A_3 = \partial_\rho \ln T, \qquad A_1 = \partial_\rho \ln n - \tfrac32 \partial_\rho \ln T + Z \alpha \hat E_r\]

profiles against the exact derivatives implied by the archived precise-QS QA/QH profile family. This is the first literature-anchored validation surface for the imported primitive-profile workflow itself, but it should be read as a coarse benchmark-family stress test for the current reconstruction scheme, not as a parity gate. It is still the right figure to carry into the paper when discussing how NTX reconstructs monoenergetic force profiles from archived benchmark inputs.

Primitive-to-force reconstruction audit

Source-Code Map