Add allocdata Python callback by dpgrote · Pull Request #6966 · BLAST-WarpX/warpx

dpgrote · 2026-06-23T18:27:41Z

This adds a new call back to the Python interface. This call back, allocdata, is called during initialization, both from scratch and from restart. New MultiFabs should be allocated at this point.

This is associated with the PR #6954 which allows writing of any MultiFabs to the diagnostics and to checkpoint files. The new MultiFabs would need to be added in the allocdata call back to be included in the diagnostics.

This also implements similar capability from PR #6889.

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

ax3l

LGTM 👍

n01r

Thanks, @dpgrote!

This adds a new call back to the Python interface. This call back, `allocdata`, is called during initialization, both from scratch and from restart. New MultiFabs should be allocated at this point. This is associated with the PR BLAST-WarpX#6954 which allows writing of any MultiFabs to the diagnostics and to checkpoint files. The new MultiFabs would need to be added in the `allocdata` call back to be included in the diagnostics. --------- Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

This allows new MultiFabs and user defined MultiFabs to be added to the checkpoint/restart. Comments: - All of the specific writes in `FlushFormatCheckpoint::WriteToFile` and read in `WarpX::InitFromCheckpoint` could be removed by setting the new flag on those MultiFabs. However, this would break any existing restart dumps since the names of the MultiFabs would change. - The change was not implemented in `Python/pywarpx/fields.py` since that usage has been deprecated. - This needs a CI test This builds on PR #6966, allowing Python defined MultiFabs to be added to checkpoint/restart.

This allows new MultiFabs and user defined MultiFabs to be added to the checkpoint/restart. Comments: - All of the specific writes in `FlushFormatCheckpoint::WriteToFile` and read in `WarpX::InitFromCheckpoint` could be removed by setting the new flag on those MultiFabs. However, this would break any existing restart dumps since the names of the MultiFabs would change. - The change was not implemented in `Python/pywarpx/fields.py` since that usage has been deprecated. - This needs a CI test This builds on PR BLAST-WarpX#6966, allowing Python defined MultiFabs to be added to checkpoint/restart.

BLAST-WarpX#7025 ("Allow any MultiFab to be written to the diagnostics"), BLAST-WarpX#7023 (checkpoint_restart flag on the registry) and BLAST-WarpX#6966 (allocdata callback) merged without a test under Examples/. Nothing in the suite registers a field from Python and writes it to a diagnostic by its registered name: - no Examples/ input uses the "allocdata" callback; - the one input that calls alloc_init (rz_spacecraft_charging) uses the fields internally via saxpy/copymf and never puts them in a data_list; - the one input using additional_fields_to_plot (3d_particle_fields_diags) lists rho/rho_electrons/rho_protons, which are matched by explicit branches rather than the registry lookup. Add that coverage to an existing test rather than a new one. The 3D ohm-solver cylinder-compression case does a 2x uniform Bz compression of a diamagnetic plasma column, so the magnetic pressure Bz^2/(2 mu0) is a physically meaningful quantity for it: Bz is excluded from the core and piles up outside, and the resulting pressure is what confines the column radially against the plasma pressure the hybrid solver already tracks. (The field is essentially purely axial -- |B_perp|/|Bz| < 1% -- while the compression it drives is radial; the field name refers to the component the pressure is built from, not to the direction of the compression.) It is registered from Python in an "allocdata" callback, refreshed after each E-solve, and requested in data_list by its registered name -- so the test exercises alloc_init, the allocdata hook, and the MultiFabRegister name lookup in FullDiagnostics together. No new input file and no measurable extra runtime. The field is registered on the Bz BoxArray, so it inherits that staggering and is exact where it is computed; the diagnostic's CellCenterFunctor does the face-to-centre averaging on output, as it does for Bz itself. Note that the cell-centred output is therefore avg(Bz^2)/(2 mu0), which is >= (avg Bz)^2/(2 mu0) by Jensen's inequality -- the two agree to 1e-16 in the uniform region and differ by up to 9.7% where Bz varies most sharply. Verified the peak value: 321.06 Pa against Bz_max = 0.0284 T, i.e. Bz_max^2/(2 mu0) = 321.06 Pa. The step-0 dump shows the allocation value (zeros) because "afterInitEsolve" fires before AddExternalFields(), so Bfield_fp does not yet hold the applied compression field, and there is no callback between AddExternalFields() and the initial diagnostic write. Every dump from step 1 on carries the current value. The benchmark gains the one new key; every other recorded value is unchanged. Depends on the preceding commit: without the PICMI pass-through the registered name is dropped in Python and never reaches the C++ lookup.

This adds a new call back to the Python interface. This call back, `allocdata`, is called during initialization, both from scratch and from restart. New MultiFabs should be allocated at this point. This is associated with the PR BLAST-WarpX#6954 which allows writing of any MultiFabs to the diagnostics and to checkpoint files. The new MultiFabs would need to be added in the `allocdata` call back to be included in the diagnostics. --------- Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

This allows new MultiFabs and user defined MultiFabs to be added to the checkpoint/restart. Comments: - All of the specific writes in `FlushFormatCheckpoint::WriteToFile` and read in `WarpX::InitFromCheckpoint` could be removed by setting the new flag on those MultiFabs. However, this would break any existing restart dumps since the names of the MultiFabs would change. - The change was not implemented in `Python/pywarpx/fields.py` since that usage has been deprecated. - This needs a CI test This builds on PR BLAST-WarpX#6966, allowing Python defined MultiFabs to be added to checkpoint/restart.

BLAST-WarpX#7025 ("Allow any MultiFab to be written to the diagnostics"), BLAST-WarpX#7023 (checkpoint_restart flag on the registry) and BLAST-WarpX#6966 (allocdata callback) merged without a test under Examples/. Nothing in the suite registers a field from Python and writes it to a diagnostic by its registered name: - no Examples/ input uses the "allocdata" callback; - the one input that calls alloc_init (rz_spacecraft_charging) uses the fields internally via saxpy/copymf and never puts them in a data_list; - the one input using additional_fields_to_plot (3d_particle_fields_diags) lists rho/rho_electrons/rho_protons, which are matched by explicit branches rather than the registry lookup. Add that coverage to an existing test rather than a new one. The 3D ohm-solver cylinder-compression case does a 2x uniform Bz compression of a diamagnetic plasma column, so the magnetic pressure Bz^2/(2 mu0) is a physically meaningful quantity for it: Bz is excluded from the core and piles up outside, and the resulting pressure is what confines the column radially against the plasma pressure the hybrid solver already tracks. (The field is essentially purely axial -- |B_perp|/|Bz| < 1% -- while the compression it drives is radial; the field name refers to the component the pressure is built from, not to the direction of the compression.) It is registered from Python in an "allocdata" callback, refreshed after each E-solve, and requested in data_list by its registered name -- so the test exercises alloc_init, the allocdata hook, and the MultiFabRegister name lookup in FullDiagnostics together. No new input file and no measurable extra runtime. The field is registered on the Bz BoxArray, so it inherits that staggering and is exact where it is computed; the diagnostic's CellCenterFunctor does the face-to-centre averaging on output, as it does for Bz itself. Note that the cell-centred output is therefore avg(Bz^2)/(2 mu0), which is >= (avg Bz)^2/(2 mu0) by Jensen's inequality (Jensen, Acta Mathematica 30, 175-193, 1906, https://doi.org/10.1007/BF02418571) -- the two agree to 1e-16 in the uniform region and differ by up to 9.7% where Bz varies most sharply. Jensen bounds the sign of that gap but not its size; for the two-point face average used here the exact identity is mean(B^2) - mean(B)^2 = (a-b)^2/4, i.e. a ratio of 1 + ((a-b)/(a+b))^2, so 9.7% corresponds to a face-pair contrast of 0.31. Verified the peak value: 321.06 Pa against Bz_max = 0.0284 T, i.e. Bz_max^2/(2 mu0) = 321.06 Pa. The step-0 dump shows the allocation value (zeros) because "afterInitEsolve" fires before AddExternalFields(), so Bfield_fp does not yet hold the applied compression field, and there is no callback between AddExternalFields() and the initial diagnostic write. Every dump from step 1 on carries the current value. The benchmark gains the one new key; every other recorded value is unchanged. Depends on the preceding commit: without the PICMI pass-through the registered name is dropped in Python and never reaches the C++ lookup.

Add allocdata callback

96ae435

dpgrote added component: Python Python layer component: initialization Changes related to the initialization of the simulation labels Jun 23, 2026

dpgrote mentioned this pull request Jun 23, 2026

Add fusion particle production diagnostic #6954

Merged

ax3l self-requested a review June 23, 2026 19:16

ax3l reviewed Jun 23, 2026

View reviewed changes

Comment thread Source/Diagnostics/WarpXIO.cpp Outdated

ax3l reviewed Jun 23, 2026

View reviewed changes

Comment thread Source/Initialization/WarpXInitData.cpp Outdated

Fix compile

d12c3bb

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

ax3l approved these changes Jun 23, 2026

View reviewed changes

ax3l self-assigned this Jun 23, 2026

ax3l enabled auto-merge (squash) June 23, 2026 19:18

Fix include

56cc2ff

ax3l merged commit 2c3eb90 into BLAST-WarpX:development Jun 23, 2026
51 of 52 checks passed

dpgrote deleted the add_allocdata_python_callback branch June 23, 2026 23:09

n01r reviewed Jun 24, 2026

View reviewed changes

dpgrote mentioned this pull request Jul 7, 2026

Add checkpoint_restart flag to MultiFab registry #7023

Merged

n01r mentioned this pull request Aug 12, 2026

Tests: cover Python-registered MultiFab output in the ohm-solver test #7158

Open

Navigation Menu

Uh oh!

Add allocdata Python callback - #6966

Add allocdata Python callback#6966
ax3l merged 3 commits into
BLAST-WarpX:developmentfrom
dpgrote:add_allocdata_python_callback

dpgrote commented Jun 23, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

ax3l left a comment

Uh oh!

Uh oh!

n01r left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

dpgrote commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ax3l left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

n01r left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

dpgrote commented Jun 23, 2026 •

edited

Loading