Crash with 3D space charge and compact particle distributions

[nikitakuklev]

Noticed some crashes during space charge model testing. Resizing mesh on particle add did not help, but increasing beam size did. I know small beams are not physical / not intended use case of 3D space charge, but it should raise an error instead of crashing. This is probably AMReX issue more so than ImpactX one?

(minimal repro by Claude):

"""
Minimal reproducer: ImpactX 3D space charge crashes in AMReX
computeRoundoffDomain when beam coordinate scales are small (~1e-6).

The crash is in amrex::Geometry::computeRoundoffDomain() — the domain
becomes so small that floating-point roundoff prevents computing valid
geometry bounds during ResizeMesh() inside HandleSpacecharge().

This is NOT specific to add_n_particles — add_particles from a
distribution triggers the same crash with the same coordinate scales.

Error:
  Assertion `rtmp > std::numeric_limits<ParticleReal>::lowest()
            && iters < maxiters'
  failed, file "/project/amrex/Src/Base/AMReX_Geometry.cpp", line 584

Run:
  python repro_impactx_crash.py          # crashes (sigma=1e-6)
  python repro_impactx_crash.py 1e-3     # works  (sigma=1e-3)

ImpactX version: 26.02
AMReX version: 26.02
"""

import sys

from impactx import ImpactX, distribution, elements

sigma = float(sys.argv[1]) if len(sys.argv) > 1 else 1e-6

sim = ImpactX()
sim.particle_shape = 2
sim.space_charge = "3D"
sim.verbose = 0
sim.tiny_profiler = False
sim.init_grids()

ref = sim.particle_container().ref_particle()
ref.set_charge_qe(1.0)
ref.set_mass_MeV(938.272)
ref.set_kin_energy_MeV(2.5)

distr = distribution.Gaussian(
    lambdaX=sigma,
    lambdaY=sigma,
    lambdaT=sigma,
    lambdaPx=sigma,
    lambdaPy=sigma,
    lambdaPt=sigma,
)
sim.add_particles(1e-10, distr, 100)

# Stable FODO cell
sim.lattice.extend([
    elements.ExactDrift(name="d0", ds=0.25),
    elements.Quad(name="qf", ds=0.5, k=1.2),
    elements.ExactDrift(name="d1", ds=0.5),
    elements.Quad(name="qd", ds=0.5, k=-1.2),
    elements.ExactDrift(name="d2", ds=0.25),
])

sim.track_particles()
sim.finalize()
print(f"OK (sigma={sigma})")

[ax3l]

Hi @nikitakuklev,

Thank you for reporting this! Yes our goal is to not crash but rather throw an exception.

The error that you see is from the AMReX Geometry::computeRoundoffDomain, and it should rather throw instead of aborting the process.

Host-side asserts like those are defined here
https://github.com/AMReX-Codes/amrex/blob/26.03/Src/Base/AMReX_BLassert.H
https://github.com/AMReX-Codes/amrex/blob/26.03/Src/Base/AMReX.H#L227-L230

What we need to set in ImpactX is that AMReX' system::throw_exception option (amrex.throw_exception = 1) is set when we run from Python, so asserts become exceptions:
https://github.com/AMReX-Codes/amrex/blob/26.03/Src/Base/AMReX.cpp#L301

We already do that in our CI runners, but it should rather be done properly when ImpactX is imported in Python.

[ax3l]

I started a PR in #1350 but will have to pause for a few days due to deadlines. I'll try to validate it next week and try to ship it in the 26.04 release cycle.