MPI Domain Decomposition Artifacts in Charging Simulation

[ashwynsam]

Issue: MPI Domain Decomposition Artifacts in Spacecraft Charging Simulation

Description

I'm extending the spacecraft charging example (Examples/Physics_applications/spacecraft_charging) to simulate plasma flow over an embedded boundary to determine steady-state charging potentials. While the simulation runs successfully in serial mode, running with MPI produces visible artifacts at subdomain boundaries, leading to incorrect results.

Reproducible Setup

• Base Example: Extended from Examples/Physics_applications/spacecraft_charging
• Implementation: Using PICMI interface
• Problem Type: Plasma-body interaction with embedded boundary

Observed Behavior

• Serial Execution: Runs correctly with smooth potential distribution
• MPI Execution:
  • Simulation completes but shows discontinuities at subdomain boundaries
  • The potential field (phi) shows clear artifacts at domain interfaces
  • These discontinuities appear to affect the overall charging solution

Visual Evidence

The attached image shows the potential field at t=2.12e-06s (iteration 100):

• Clear discontinuities visible in the potential field
• Sharp transitions occur at what appear to be MPI domain boundaries
• The central region shows the charged body, but the surrounding potential field has unphysical steps

System Information

• WarpX Version: 24.09-65-ge3e6ab8690e9
• MPI Implementation: 2 GPU nodes
• Number of MPI ranks used: 4 ranks/node; total of 8 MPI processes
• Platform: Perlmutter

Input Script

The full script can be found below. The key components are:

• Using PICMI interface with embedded boundary for spacecraft charging
• Multiple particle species (electrons and protons) with both uniform and flux distributions
• Custom field correction implementation for charge collection
• Running with MPI decomposition (causing the observed artifacts)

#!/usr/bin/env python3

from pywarpx import picmi
from pywarpx.fields import ExWrapper, EyWrapper, EzWrapper, PhiFPWrapper, RhoFPWrapper
from pywarpx.callbacks import installafterInitEsolve, installafterEsolve
from pywarpx.particle_containers import ParticleBoundaryBufferWrapper
import scipy.constants as scc
import numpy as np
from mpi4py import MPI as mpi

# Utilities
class SpaceChargeFieldCorrector(object):

    def __init__(self):
        self.saved_first_iteration_fields = False
        self.spacecraft_potential = 1. # Initial voltage: 1V
        self.spacecraft_capacitance = None

    def correct_space_charge_fields(self, q=None):
        """
        Function that will be called at each iteration,
        after each electrostatic solve in WarpX
        """
        assert self.saved_first_iteration_fields

        # Compute the charge that WarpX thinks there is on the spacecraft
        # from phi and rho after the Poisson solver
        q_v = compute_virtual_charge_on_spacecraft()
        if q is None:
            q = compute_actual_charge_on_spacecraft()

        # Correct fields so as to recover the actual charge
        Ex = ExWrapper(include_ghosts=True)
        Ex[...] += (q - q_v)*self.normalized_Ex[...]
        Ey = EyWrapper(include_ghosts=True)
        Ey[...] += (q - q_v)*self.normalized_Ey[...]
        Ez = EzWrapper(include_ghosts=True)
        Ez[...] += (q - q_v)*self.normalized_Ez[...]
        phi = PhiFPWrapper(include_ghosts=True)
        phi[...] += (q - q_v)*self.normalized_phi[...]
        # Update the potential at the boundary of the spacecraft
        self.spacecraft_potential += (q - q_v)*self.spacecraft_capacitance
        sim.extension.warpx.set_potential_on_eb( "%f" %self.spacecraft_potential )
        print('Setting potential to %f' %self.spacecraft_potential)

        # Confirm that the charge on the spacecraft is now correct
        compute_virtual_charge_on_spacecraft()

    def save_normalized_vacuum_Efields(self):
        # Compute the charge that WarpX thinks there is on the spacecraft
        # from phi and rho after the Poisson solver
        q_v = compute_virtual_charge_on_spacecraft()
        self.spacecraft_capacitance = 1./q_v # the potential was set to 1V

        # Check that this iteration corresponded to a vacuum solve
        rho = RhoFPWrapper(include_ghosts=False)
        # In principle, we should check that `rho` is exactly 0
        # However, due to machine precision errors when adding the charge
        # of ions and electrons, this can be slightly different than 0
        assert np.all( abs(rho[...]) < 1.e-12 )

        # Record fields
        Ex = ExWrapper(include_ghosts=True)
        self.normalized_Ex = Ex[...].copy()/q_v
        Ey = EyWrapper(include_ghosts=True)
        self.normalized_Ey = Ey[...].copy()/q_v
        Ez = EzWrapper(include_ghosts=True)
        self.normalized_Ez = Ez[...].copy()/q_v
        phi = PhiFPWrapper(include_ghosts=True)
        self.normalized_phi = phi[...].copy()/q_v

        self.saved_first_iteration_fields = True
        self.correct_space_charge_fields(q=0)

def compute_virtual_charge_on_spacecraft():
    """
    Given that we asked WarpX to solve the Poisson
    equation with phi=1 on the spacecraft and phi=0
    on the boundary of the domain, compute the charge
    that WarpX thinks there should be on the spacecraft.
    """
    # Get global array for the whole domain (across MPI ranks)
    phi = PhiFPWrapper(include_ghosts=False)[:,:,:].copy()
    rho = RhoFPWrapper(include_ghosts=False)[:,:,:].copy()

    # Check that this codes correspond to the global size of the box
    assert phi.shape == (nx+1, ny+1, nz+1)
    assert rho.shape == (nx+1, ny+1, nz+1)

    dx, dy, dz = sim.extension.warpx.Geom(lev=0).data().CellSize()

    # Compute integral of grad phi over surfaces of the domain
    grad_phi_integral = \
        1./dx*(phi[-1,:,:]-phi[-2,:,:]).sum() * dy*dz \
      + 1./dx*(phi[0,:,:]-phi[1,:,:]).sum() * dy*dz \
      + 1./dy*(phi[:,-1,:]-phi[:,-2,:]).sum() * dx*dy \
      + 1./dy*(phi[:,0,:]-phi[:,1,:]).sum() * dx*dy \
      + 1./dz*(phi[:,:,-1]-phi[:,:,-2]).sum() * dy*dz \
      + 1./dz*(phi[:,:,0]-phi[:,:,1]).sum() * dy*dz

    # Compute integral of rho over volume of the domain
    # (i.e. total charge of the plasma particles)
    rho_integral = rho[1:-2, 1:-2, 1:-2].sum() * dx*dy*dz

    # Compute charge of the spacecraft, based on Gauss theorem
    q_spacecraft = - rho_integral - scc.epsilon_0 * grad_phi_integral
    print('Virtual charge on the spacecraft: %e' %q_spacecraft)
    return float(q_spacecraft)  # Convert from `nump`/`cupy` array to float

def compute_actual_charge_on_spacecraft():
    """
    Compute the actual charge on the spacecraft,
    by counting how many electrons and protons
    where collected by the WarpX embedded boundary (EB)
    """
    charge= {'electrons': -scc.e,
             'electrons2': -scc.e,
             'electrons3': -scc.e,
             'protons': scc.e,
             'protons2': scc.e}
    q_spacecraft = 0

    particle_buffer = ParticleBoundaryBufferWrapper()

    for species in charge.keys():
        weights = particle_buffer.get_particle_boundary_buffer(species, 'eb', 'w', 0)
        sum_weights_over_tiles = sum([w.sum() for w in weights])
        # Reduce across all MPI ranks
        ntot = mpi.COMM_WORLD.allreduce(sum_weights_over_tiles, op=mpi.SUM)
        print('Total number of %s collected on spacecraft: %e'%(species, ntot))
        q_spacecraft += ntot * charge[species]
    print('Actual charge on the spacecraft: %e' %q_spacecraft)
    return float(q_spacecraft)  # Convert from `nump`/`cupy` array to float

##########################
# numerics parameters
##########################

dt = 2.11865e-8

# --- Nb time steps

max_steps = 100
diagnostic_interval = 10

# --- grid

nx = 128
ny = 128
nz = 128

xmin = -5.0
xmax = 5.0
ymin = -5.0
ymax = 5.0
zmin = -5.0
zmax = 5.0

number_per_cell = 1
ppc = [1, 1, 1]

##########################
# physics components
##########################

n = 7.0e9 #plasma density #particles/m^3
Te = 85 #Electron temp in eV
Ti = 82 #Ion temp in eV
qe = picmi.constants.q_e #elementary charge
m_e = picmi.constants.m_e #electron mass
m_i = 1836.0 * m_e #mass of ion
v_eth = (qe * Te / m_e) ** 0.5
v_pth = (qe * Ti / m_i) ** 0.5
v_mean = 300.e3

e_dist = picmi.UniformFluxDistribution(
    flux=n*v_eth/(2*np.pi)**.5, # Flux for Gaussian with negligible mean velocity
    surface_flux_position=-4.9,
    flux_direction=+1, flux_normal_axis='x',
    directed_velocity= [v_mean, 0, 0],
    rms_velocity=[v_eth, v_eth, v_eth],
    gaussian_flux_momentum_distribution=True )
electrons = picmi.Species(particle_type='electron', name='electrons',
                          initial_distribution=e_dist,
                          warpx_save_particles_at_eb=1)

e_dist3 = picmi.UniformFluxDistribution(
    flux=n*v_eth/(2*np.pi)**.5, # Flux for Gaussian with negligible mean velocity
    surface_flux_position=4.9,
    flux_direction=-1, flux_normal_axis='x',
    directed_velocity= [v_mean, 0, 0],
    rms_velocity=[v_eth, v_eth, v_eth],
    gaussian_flux_momentum_distribution=True )
electrons3 = picmi.Species(particle_type='electron', name='electrons3',
                          initial_distribution=e_dist3,
                          warpx_save_particles_at_eb=1)

p_dist = picmi.UniformFluxDistribution(
    flux=n*v_mean, # Flux for Gaussian with negligible temperature
    surface_flux_position=-4.9,
    flux_direction=+1, flux_normal_axis='x',
    directed_velocity=[v_mean, 0, 0],
    rms_velocity=[v_pth, v_pth, v_pth],
    gaussian_flux_momentum_distribution=True )
protons = picmi.Species(particle_type='proton', name='protons',
                        initial_distribution=p_dist,
                        warpx_save_particles_at_eb=1)

e_dist2 = picmi.UniformDistribution(density = n, rms_velocity=[v_eth, v_eth, v_eth], directed_velocity = [v_mean, 0, 0])
electrons2 = picmi.Species(particle_type='electron', name='electrons2',
             initial_distribution=e_dist2, warpx_save_particles_at_eb=1)

p_dist2 = picmi.UniformDistribution(density = n, rms_velocity=[v_pth, v_pth, v_pth], directed_velocity = [v_mean, 0, 0])
protons2 = picmi.Species(particle_type='proton', name='protons2',
           initial_distribution=p_dist2, warpx_save_particles_at_eb=1)

##########################
# numerics components
##########################

grid = picmi.Cartesian3DGrid(
    number_of_cells = [nx, ny, nz],
    lower_bound = [xmin, ymin, zmin],
    upper_bound = [xmax, ymax, zmax],
    lower_boundary_conditions = ['dirichlet', 'dirichlet', 'dirichlet'],
    upper_boundary_conditions =  ['dirichlet', 'dirichlet', 'dirichlet'],
    lower_boundary_conditions_particles = ['absorbing', 'reflecting', 'reflecting'],
    upper_boundary_conditions_particles =  ['absorbing', 'reflecting', 'reflecting'],
    warpx_blocking_factor=8,
    warpx_max_grid_size = 256
)

solver = picmi.ElectrostaticSolver(
    grid=grid, method='Multigrid',
    required_precision=1e-7,
    warpx_absolute_tolerance=1e-6
)

embedded_boundary = picmi.EmbeddedBoundary(
    implicit_function= "-0.5 + (x >= -2.0) * (x <= -1.0) * (y >= -0.5) * (y <= 0.5) * (z >= -0.5) * (z <= 0.5)",
    potential=1., # arbitrary value ; this will be corrected by a callback function
)

##########################
# diagnostics
##########################

field_diag = picmi.FieldDiagnostic(
    name = 'diag1',
    grid = grid,
    period = diagnostic_interval,
    data_list = ['Ex', 'Ey', 'Ez', 'phi', 'rho',
                 'rho_electrons', 'rho_electrons2', 'rho_electrons3',
                 'rho_protons', 'rho_protons2'],
    warpx_format = 'openpmd'
)

part_diag = picmi.ParticleDiagnostic(name = 'diag1',
    period = diagnostic_interval,
    species = [electrons, electrons2, electrons3,
               protons, protons2],
    warpx_format = 'openpmd'
)

##########################
# simulation setup
##########################

sim = picmi.Simulation(
    solver = solver,
    time_step_size = dt,
    max_steps = max_steps,
    warpx_embedded_boundary=embedded_boundary,
    warpx_amrex_the_arena_is_managed=1
)

sim.add_species(electrons,
                layout = picmi.PseudoRandomLayout(n_macroparticles_per_cell=number_per_cell, grid=grid))

sim.add_species(electrons3,
                layout = picmi.PseudoRandomLayout(n_macroparticles_per_cell=number_per_cell, grid=grid))

sim.add_species(protons,
                layout = picmi.PseudoRandomLayout(n_macroparticles_per_cell=number_per_cell, grid=grid))

sim.add_species(electrons2,
                layout = picmi.GriddedLayout(n_macroparticle_per_cell=ppc, grid=grid))

sim.add_species(protons2,
                layout = picmi.GriddedLayout(n_macroparticle_per_cell=ppc, grid=grid))

sim.add_diagnostic(field_diag)
sim.add_diagnostic(part_diag)

##########################
# simulation run
##########################

spc = SpaceChargeFieldCorrector()
installafterInitEsolve( spc.save_normalized_vacuum_Efields )
installafterEsolve( spc.correct_space_charge_fields )

sim.step(max_steps)

Sbatch Script

#!/bin/bash -l

#SBATCH -t 00:05:00
#SBATCH -N 2
#SBATCH -J WarpX
#    note: <proj> must end on _g
#SBATCH -A m4667_g
#SBATCH -q debug
# A100 40GB (most nodes)
#SBATCH -C gpu
# A100 80GB (256 nodes)
#S BATCH -C gpu&hbm80g
#SBATCH --exclusive
# ideally single:1, but NERSC cgroups issue
#SBATCH --gpu-bind=none
#SBATCH --ntasks-per-node=4
#SBATCH --gpus-per-node=4
#SBATCH -o WarpX.o%j
#SBATCH -e WarpX.e%j

# pin to closest NIC to GPU
export MPICH_OFI_NIC_POLICY=GPU

# threads for OpenMP and threaded compressors per MPI rank
#   note: 16 avoids hyperthreading (32 virtual cores, 16 physical)
export SRUN_CPUS_PER_TASK=16
export OMP_NUM_THREADS=${SRUN_CPUS_PER_TASK}

# GPU-aware MPI optimizations
GPU_AWARE_MPI="amrex.use_gpu_aware_mpi=1"
# GPU_AWARE_MPI="warpx_amrex_use_gpu_aware_mpi=1"

# CUDA visible devices are ordered inverse to local task IDs
#   Reference: nvidia-smi topo -m
srun --cpu-bind=cores bash -c "
    export CUDA_VISIBLE_DEVICES=\$((3-SLURM_LOCALID));
    python3 PICMI_inputs_3d.py ${GPU_AWARE_MPI}" \
  > output.txt

[RemiLehe]

Thanks for reporting this issue. I was able to reproduce it.

A few remarks:

• The issue also occurs when using only 2 GPUs (instead of 8), but not when using only 1 GPU.
• The issue can be reproduced with 1 GPU, by forcing a decomposition into 2 boxes, with:

  warpx_max_grid_size_x = 128,
  warpx_max_grid_size_y = 128,
  warpx_max_grid_size_z = 64

• It also occurs without GPU-aware MPI (amrex.use_gpu_aware_mpi=0)
• The issue does not seem to occur when callback function are commented out:

[RemiLehe]

It seems that setting include_ghosts=False everywhere in the script fixes the issue. @ashwynsam Could you confirm?