Diffusion Optimizer workflow

Diffusion Optimizer fits the anomalous transport coefficients used by axisymmetric XGC (XGCa) to maintain prescribed flux-surface averaged plasma profiles. XGCa advances the plasma with its current coefficients; the optimizer reads a short history of density, parallel flow, and temperature, estimates how the coefficients should change, and writes the next coefficient profiles. Repeating this cycle is intended to reduce slow profile evolution.

The implementation assumes a pure, quasineutral plasma. Electron density is the particle-transport density for both species. It fits particle diffusivity D and pinch velocity v_p from the electron density equation, ion momentum diffusivity eta from the ion momentum equation, and separate electron and ion heat conductivities chi_e and chi_i from their energy equations. Electron momentum is not fitted here; XGCa controls the current through its loop-voltage model.

Data flow

XGCa: evolve with current coefficients
    -> xgc.diffusion_profiles.bp (profile histories)
    -> Diffusion Optimizer: preprocess, fit, post-process
    -> xgc.diffusion_coeff.bp (updated full coefficients)
    -> XGCa: next update interval

At startup, the optimizer builds axisymmetric flux-surface geometry from the XGC mesh and equilibrium. It opens and drains the profile stream to its latest step, then opens the coefficient output stream. The profile stream stays open while it waits for subsequent XGC updates. The timestep.dat state distinguishes a restart from a new run: on restart, the latest available profile step has already been processed, so the first loop iteration waits for a new one. Stream ordering matters when both programs restart together.

Each profile output contains several time samples. The optimizer smooths the histories in time and radius, differentiates them on the nonuniform flux grid, and computes Vprime = dV/dpsi from enclosed flux-surface volume. By default, the fitting system averages its rows over the samples in the current update interval. --no-time-avg-system instead retains the individual time-sample rows.

Fitting order and coefficient updates

The current solver uses full-domain radial matrices. It fits corrections to the coefficients active during the preceding XGC interval, rather than reconstructing each full coefficient profile independently:

1. Fit electron Delta D and Delta v_p to counter the measured density evolution. The density solve constrains diffusion and pinch jointly to limit ambiguity between their particle flux contributions.

2. Pass those undamped particle corrections to the ion solve. Fit ion Delta eta using the conservative momentum equation.

3. Fit Delta chi_e and Delta chi_i using the conservative energy equations. The ion energy fit uses the undamped particle and momentum corrections obtained earlier in the same update.

The matrix builders form coefficient-to-rate operators from the measured profiles and geometry. Their radial flux-divergence terms use one-sided upwind stencils on the flux grid. Known convective and coupling terms enter the momentum and energy targets. The sign convention asks the fitted correction to counter the observed profile time derivative. The QP backend is the default; --solver-method lsq selects bounded least squares.

The code estimates particle, momentum, and heat flux balances at the separatrix. In the current full-domain solves, those separatrix estimates are diagnostics; they are not fixed coefficient values. A fixed correction may be imposed at the inner edge of the retained fit domain. At the outer edge, the coefficient is free within its applicable bounds. Rows where XGC’s diffusion boundary factor is small are excluded from fitting.

After each solve, the optional red-black filter processes the correction while restoring any fixed inner value. This filter is off by default. The full trial coefficient is then kept nonnegative for D, eta, and chi. In the default constrained cap mode, bounds and regularization are part of the solve, followed by a safety clip. In post-solve mode, the fit drops its regularization penalties and applies a smooth rational cap to the full trial profile afterward. The core cap defaults to 40 in the coefficient’s units and can be set with --coefficient-cap-core.

Only after the density, momentum, and energy stages have finished are their corrections multiplied by --update-damping-factor. Thus later stages see the undamped corrections from earlier stages; damping stabilizes the coefficients written to XGC without changing the equations used for that update’s fit. A configured momentum RHS taper also drives the existing ion eta profile toward zero outside its active radial band after fitting.

Learned response correction

The one-dimensional upwind operator does not exactly reproduce XGCa’s two-dimensional finite-element diffusion response, especially near the separatrix. --online-response optionally learns a short radial response kernel from successive intervals. It synchronizes xgc.fsourcediag.*.bp samples with the diffusion-profile history and compares the one-dimensional predicted diffusion rates with the flux-surface averaged XGC diffusion changes. Separate five-point kernels are trained for density, ion momentum, electron heat, and ion heat. The kernel uses changes between consecutive intervals and starts from an identity response until enough intervals have been observed.

--online-response-window PSI_MIN PSI_MAX selects a normalized-flux training and application interval. The correction blends back to the uncorrected operator over two surfaces beyond that interval. Without this option, the default window is centered near the separatrix. The learned state is saved in online_response_state.npz for restarts, and prediction quality is plotted in plots/online_response_quality.svg. Online response requires the time-averaged fitting system.

Diagnostics and control switches

Fit-quality plots compare the target rate with the rate predicted by the current fit, including processed undamped corrections. They are intended to diagnose the fitting equation; under-relaxation is applied later. Coefficient plots show the full profiles before and after the update. If xgc.heatdiag2.bp is available, wall-loss plots are also produced. Wall loss is diagnostic and does not impose an outer coefficient boundary value.

--first-update-gstep permits fitting and plotting before coefficient changes are applied. --fitting-stop-gstep skips subsequent fits and writes unchanged coefficients while continuing the update loop and available diagnostics. The active command-line and hard-wired configuration values are printed at startup for reproducibility. For the complete list of switches, run python Diffusion_Optimizer/diffusion_optimizer.py --help.

Model boundary

The optimizer fits flux-surface averaged, one-dimensional equations to a two-dimensional XGCa calculation. Near the closed/open flux-surface transition, the mapped area, enclosed volume, parallel wall loss, and XGC’s finite-element operator can all affect how closely the fitted one-dimensional rate predicts the actual response. Fit-quality plots measure consistency with the optimizer’s equations; the online-response diagnostic tests how well changes predicted by those equations agree with changes measured by XGC. Profile control in a live run remains the final check.

The driver is Diffusion_Optimizer/diffusion_optimizer.py. Matrix assembly, solver stages, preprocessing, cap processing, and online response are kept in the neighboring diffopt_*.py modules. The live profile reader, coefficient writer, and synthetic profile generator also live in Diffusion_Optimizer/. They are workflow components rather than members of the installed xgc_analysis library. Run the driver from the repository checkout with XGC-Analysis installed in the Python environment or supply --xgc-analysis-path.

Workflow stream classes

class Diffusion_Optimizer.DiffusionProfiles.DiffusionProfiles(adios: adios2.Adios, *, bp_dir: os.PathLike | str = './xgc.diffusion_profiles.bp', wait_interval: float = 1.0, timeout: float | None = None)

Live reader for xgc.diffusion_profiles.bp.

Reads diffusion profile data from an Adios BP file (directory) xgc.diffusion_profiles.bp. The constructor opens the BP directory named by bp_dir. Once it is found, it opens it as a persistent adios2.Stream, drains all existing steps while retaining only the final one. The Stream is left open until it is closed with the member method close(). A blocking request for a new step can be posted with the member method read_next_step.

The file xgc.diffusion_profiles.bp has the following structure:

double density n_steps*{n_species, n_samples, n_surf} double flow n_steps*{n_species, n_samples, n_surf} double temp n_steps*{n_species, n_samples, n_surf} int32_t n_samples n_steps*scalar int32_t n_species n_steps*scalar int32_t n_surf n_steps*scalar double psi n_steps*{n_surf} int32_t steps n_steps*{n_samples} double time n_steps*{n_samples}

That is, there are (for example) n_steps steps recorded in the BP file. The constructor of this class reads only the last step (i.e. the most recent one).

The property latest_step_data returns a dictionary with the following keys:

density : A 3D array (shape (n_species, n_samples, n_surf)) of the diffusion density. flow : A 3D array (shape (n_species, n_samples, n_surf) of the parallel mean flow. temp : A 3D array (shape (n_species, n_samples, n_surf)) of temperature (in eV). n_samples : An integer (n_samples) representing the number of time samples. n_species : An integer (default 2) representing the number of particle species. n_surf : An integer (n_surf) representing the number of ψ-surfaces. psi : A 1D array (shape (n_surf,)) containing the poloidal magnetic flux values. steps : A 1D array (shape (n_samples,)) of time step indices. time : A 1D array (shape (n_samples,)) of simulation times (in seconds).

Parameters:

• bp_dir (str | os.PathLike) – Path to the BP directory written by XGC. The constructor blocks until the directory exists (subject to timeout).

• wait_interval (float, optional) – Polling interval in seconds while waiting for the directory. Defaults to 1.0 s.

• timeout (float | None, optional) – Abort after this many seconds if the directory still does not exist. None (default) means wait indefinitely.

Notes

The object keeps exactly one snapshot in memory. Callers needing to store a history must copy out the data returned by latest_step_data.

close() → None

Close the underlying stream.

property latest_step_data: Dict[str, ndarray]

Return a shallow copy of the dictionary holding the current step.

Raises:

RuntimeError – If no step has been read yet (e.g., the writer hasn’t flushed a step to disk).

classmethod load_last_sample_history(adios: Adios, *, bp_dir: PathLike | str = './xgc.diffusion_profiles.bp', max_steps: int | None = None, stride: int = 1, open_timeout: float | None = 0.0) → Dict[str, list]

Reconstruct per-step history using the last time sample of each profile.

Parameters:

• adios (Adios) – ADIOS handle used to create a temporary read stream.

• bp_dir (str | os.PathLike) – Path to xgc.diffusion_profiles.bp.

• max_steps (int | None) – If provided and > 0, keep only the last max_steps entries.

• stride (int) – Keep every stride-th step while scanning BP history.

• open_timeout (float | None) – OpenTimeoutSecs passed to the temporary read stream.

Returns:

History dictionary with keys: density, flow, temp (each list of (n_species, n_surf) arrays), gstep (list[int]), and time (list[float]).

Return type:

dict[str, list]

Notes

This helper is independent from live-stream operation and does not change the behavior of the default constructor.

property n_samples: int | None

Number of (time) samples per species in snapshot, else None.

property n_species: int | None

Number of species in the current snapshot (or None).

property n_surf: int | None

Radial surface count (psi grid size) in snapshot, else None.

read_next_step(*, my_timeout: float = 10.0, poll_delay: float = 0.5) → bool

Advance the stream by one time step.

Examples

>>> prof = DiffusionProfiles("xgc.diffusion_profiles.bp")
>>> prof.read_next_step(block=False)
False  # no new data yet
>>> # ... later
>>> prof.read_next_step()
True   # newest snapshot now loaded

Parameters:

• my_timeout – Maximum wait time for a new step; if no new step is found within my_timeout seconds, return False.

• poll_delay – Sleep duration (seconds) between availability checks when block is True.

Returns:

True if a new step was successfully read, False otherwise.

Return type:

bool

class Diffusion_Optimizer.DiffusionCoefficients.DiffusionCoefficients(adios: Adios, *, timeout=10.0, nspecies=2, npsi=100)

Stores the profiles of diffusion model coefficients defined as functions of the poloidal magnetic flux ψ.

The four diffusion coefficients are:

• ptl_diffusivity (m^2/s)

• momentum_diffusivity (m^2/s)

• heat_conductivity (m^2/s)

• ptl_pinch_velocity (m/s)

The data layout of each coefficient is (nspecies, npsi) where the second dimension (the ψ-grid) is contiguous in memory. Also the ψ–grid is stored as a 1D NumPy array of length npsi.

Parameters:

nspecies: Number of particle species (default: 2). npsi: Number of points in the ψ–grid.

close() → None

Close the underlying stream.

write_to_file(filename='xgc.diffusion_coeff.bp', gstep=-1)

Writes the diffusion coefficient profiles and the ψ–grid to an Adios BP file as a new step.

For each diffusion coefficient, the data for each species is split into an individual variable using a species suffix. The species suffixes (for species indices 0 through 6) are:

0: _elec 1: _ion 2: _imp1 3: _imp2 4: _imp3 5: _imp4 6: _imp5

In addition, the file will store:

• n_species: a scalar integer (int64).

• psi: a 1D array (of length npsi).

The method opens the specified file (filename) using the Adios2 FileWriter in append mode, writes all the individual diffusion profiles, and then ends the step.

class Diffusion_Optimizer.AnalyticDiffusionProfiles.AnalyticDiffusionProfiles(psi, time, S, Vol, grad_avg, psi_norm)

Creates diffusion profile data using analytic functions.

Attributes:

density : A 3D array (shape (n_species, n_samples, n_surf)) of the diffusion density. flow : A 3D array (shape (n_species, n_samples, n_surf) of the parallel mean flow. temp : A 3D array (shape (n_species, n_samples, n_surf)) of temperature (in eV). n_samples : An integer (n_samples) representing the number of time samples. n_species : An integer (default 2) representing the number of particle species. n_surf : An integer (n_surf) representing the number of ψ-surfaces. psi : A 1D array (shape (n_surf,)) containing the poloidal magnetic flux values. steps : A 1D array (shape (n_samples,)) of time step indices. time : A 1D array (shape (n_samples,)) of simulation times (in seconds).