Merge branch 'development' into add_insulator_bc

.github/workflows/cuda.yml

@@ -115,7 +115,7 @@ jobs:
         which nvcc || echo "nvcc not in PATH!"
 
         git clone https://github.com/AMReX-Codes/amrex.git ../amrex
-        cd ../amrex && git checkout --detach 28b010126a1b39297d8a496ba81f171d8563953b && cd -
+        cd ../amrex && git checkout --detach 24.06 && cd -
         make COMP=gcc QED=FALSE USE_MPI=TRUE USE_GPU=TRUE USE_OMP=FALSE USE_FFT=TRUE USE_CCACHE=TRUE -j 4
 
         ccache -s

CMakeLists.txt

@@ -1,7 +1,7 @@
 # Preamble ####################################################################
 #
 cmake_minimum_required(VERSION 3.20.0)
-project(WarpX VERSION 24.05)
+project(WarpX VERSION 24.06)
 
 include(${WarpX_SOURCE_DIR}/cmake/WarpXFunctions.cmake)
 
@@ -79,6 +79,7 @@ option(WarpX_LIB           "Build WarpX as a library"                   OFF)
 option(WarpX_MPI           "Multi-node support (message-passing)"       ON)
 option(WarpX_OPENPMD       "openPMD I/O (HDF5, ADIOS)"                  ON)
 option(WarpX_FFT           "FFT-based solvers"                          OFF)
+option(WarpX_HEFFTE        "Multi-node FFT-based solvers"               OFF)
 option(WarpX_PYTHON        "Python bindings"                            OFF)
 option(WarpX_SENSEI        "SENSEI in situ diagnostics"                 OFF)
 option(WarpX_QED           "QED support (requires PICSAR)"              ON)
@@ -136,6 +137,10 @@ mark_as_advanced(WarpX_MPI_THREAD_MULTIPLE)
 
 option(WarpX_amrex_internal                    "Download & build AMReX" ON)
 
+if(WarpX_HEFFTE AND NOT WarpX_MPI)
+    message(FATAL_ERROR "WarpX_HEFFTE (${WarpX_HEFFTE}) can only be used if WarpX_MPI is ON.")
+endif()
+
 # change the default build type to Release (or RelWithDebInfo) instead of Debug
 set_default_build_type("Release")
 
@@ -174,6 +179,10 @@ option(ABLASTR_FFT "compile AnyFFT wrappers" ${WarpX_FFT})
 if(WarpX_FFT)
     set(ABLASTR_FFT ON CACHE STRING "FFT-based solvers" FORCE)
 endif()
+option(ABLASTR_HEFFTE "compile AnyFFT wrappers" ${WarpX_HEFFTE})
+if(WarpX_HEFFTE)
+    set(ABLASTR_HEFFTE ON CACHE STRING "Multi-Node FFT-based solvers" FORCE)
+endif()
 
 # this defined the variable BUILD_TESTING which is ON by default
 #include(CTest)
@@ -215,6 +224,23 @@ if(WarpX_FFT)
     endif()
 endif()
 
+# multi-node FFT
+if(WarpX_HEFFTE)
+    if(WarpX_COMPUTE STREQUAL CUDA)
+        set(_heFFTe_COMPS CUDA)
+    elseif(WarpX_COMPUTE STREQUAL HIP)
+        set(_heFFTe_COMPS ROCM)
+    elseif(WarpX_COMPUTE STREQUAL SYCL)
+        set(_heFFTe_COMPS ONEAPI)
+    else()  # NOACC, OMP
+        set(_heFFTe_COMPS FFTW)  # or MKL
+    endif()
+    # note: we could also enforce GPUAWARE for CUDA and HIP, which can still be
+    #       disabled at runtime
+
+    find_package(Heffte REQUIRED COMPONENTS ${_heFFTe_COMPS})
+endif()
+
 # Python
 if(WarpX_PYTHON)
     find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
@@ -455,6 +481,10 @@ foreach(D IN LISTS WarpX_DIMS)
         endif()
     endif()
 
+    if(ABLASTR_HEFFTE)
+        target_link_libraries(ablastr_${SD} PUBLIC Heffte::Heffte)
+    endif()
+
     if(WarpX_PYTHON)
         target_link_libraries(pyWarpX_${SD} PRIVATE pybind11::module pybind11::windows_extras)
         if(WarpX_PYTHON_IPO)
@@ -539,6 +569,13 @@ foreach(D IN LISTS WarpX_DIMS)
         target_compile_definitions(ablastr_${SD} PUBLIC ABLASTR_USE_FFT)
     endif()
 
+    if(WarpX_HEFFTE)
+        target_compile_definitions(ablastr_${SD} PUBLIC WARPX_USE_HEFFTE)
+    endif()
+    if(ABLASTR_HEFFTE)
+        target_compile_definitions(ablastr_${SD} PUBLIC ABLASTR_USE_HEFFTE)
+    endif()
+
     if(WarpX_PYTHON AND pyWarpX_VERSION_INFO)
         # for module __version__
         target_compile_definitions(pyWarpX_${SD} PRIVATE

Docs/source/conf.py

@@ -103,9 +103,9 @@ def __init__(self, *args, **kwargs):
 # built documents.
 #
 # The short X.Y version.
-version = u'24.05'
+version = u'24.06'
 # The full version, including alpha/beta/rc tags.
-release = u'24.05'
+release = u'24.06'
 
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.

Docs/source/install/cmake.rst

@@ -97,6 +97,7 @@ CMake Option                  Default & Values                             Descr
 ``WarpX_PRECISION``           SINGLE/**DOUBLE**                            Floating point precision (single/double)
 ``WarpX_PARTICLE_PRECISION``  SINGLE/**DOUBLE**                            Particle floating point precision (single/double), defaults to WarpX_PRECISION value if not set
 ``WarpX_FFT``                 ON/**OFF**                                   FFT-based solvers
+``WarpX_HEFFTE``              ON/**OFF**                                   Multi-Node FFT-based solvers
 ``WarpX_PYTHON``              ON/**OFF**                                   Python bindings
 ``WarpX_QED``                 **ON**/OFF                                   QED support (requires PICSAR)
 ``WarpX_QED_TABLE_GEN``       ON/**OFF**                                   QED table generation support (requires PICSAR and Boost)
@@ -271,6 +272,7 @@ Environment Variable          Default & Values                             Descr
 ``WARPX_PRECISION``           SINGLE/**DOUBLE**                            Floating point precision (single/double)
 ``WARPX_PARTICLE_PRECISION``  SINGLE/**DOUBLE**                            Particle floating point precision (single/double), defaults to WarpX_PRECISION value if not set
 ``WARPX_FFT``                 ON/**OFF**                                   FFT-based solvers
+``WARPX_HEFFTE``              ON/**OFF**                                   Multi-Node FFT-based solvers
 ``WARPX_QED``                 **ON**/OFF                                   PICSAR QED (requires PICSAR)
 ``WARPX_QED_TABLE_GEN``       ON/**OFF**                                   QED table generation (requires PICSAR and Boost)
 ``BUILD_PARALLEL``            ``2``                                        Number of threads to use for parallel builds

Docs/source/install/dependencies.rst

@@ -23,12 +23,13 @@ Optional dependencies include:
 - for on-node accelerated compute *one of either*:
 
   - `OpenMP 3.1+ <https://www.openmp.org>`__: for threaded CPU execution or
-  - `CUDA Toolkit 11.7+ <https://developer.nvidia.com/cuda-downloads>`__: for Nvidia GPU support (see `matching host-compilers <https://gist.github.com/ax3l/9489132>`_) or
+  - `CUDA Toolkit 11.7+ <https://developer.nvidia.com/cuda-downloads>`__: for Nvidia GPU support (see `matching host-compilers <https://gist.github.com/ax3l/9489132>`__) or
   - `ROCm 5.2+ (5.5+ recommended) <https://gpuopen.com/learn/amd-lab-notes/amd-lab-notes-rocm-installation-readme/>`__: for AMD GPU support
-- `FFTW3 <http://www.fftw.org>`_: for spectral solver (PSATD) support when running on CPU or SYCL
+- `FFTW3 <http://www.fftw.org>`__: for spectral solver (PSATD or IGF) support when running on CPU or SYCL
 
   - also needs the ``pkg-config`` tool on Unix
-- `BLAS++ <https://github.com/icl-utk-edu/blaspp>`_ and `LAPACK++ <https://github.com/icl-utk-edu/lapackpp>`_: for spectral solver (PSATD) support in RZ geometry
+- `heFFTe 2.4.0+ <https://github.com/icl-utk-edu/heffte`__: for multi-node spectral solver (IGF) support
+- `BLAS++ <https://github.com/icl-utk-edu/blaspp>`__ and `LAPACK++ <https://github.com/icl-utk-edu/lapackpp>`__: for spectral solver (PSATD) support in RZ geometry
 - `Boost 1.66.0+ <https://www.boost.org/>`__: for QED lookup tables generation support
 - `openPMD-api 0.15.1+ <https://github.com/openPMD/openPMD-api>`__: we automatically download and compile a copy of openPMD-api for openPMD I/O support
 

Docs/source/install/hpc/perlmutter.rst

@@ -153,7 +153,7 @@ Use the following :ref:`cmake commands <building-cmake>` to compile the applicat
          cd $HOME/src/warpx
          rm -rf build_pm_gpu
 
-         cmake -S . -B build_pm_gpu -DWarpX_COMPUTE=CUDA -DWarpX_FFT=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_DIMS="1;2;RZ;3"
+         cmake -S . -B build_pm_gpu -DWarpX_COMPUTE=CUDA -DWarpX_FFT=ON -DWarpX_HEFFTE=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_DIMS="1;2;RZ;3"
          cmake --build build_pm_gpu -j 16
 
       The WarpX application executables are now in ``$HOME/src/warpx/build_pm_gpu/bin/``.
@@ -164,7 +164,7 @@ Use the following :ref:`cmake commands <building-cmake>` to compile the applicat
          cd $HOME/src/warpx
          rm -rf build_pm_gpu_py
 
-         cmake -S . -B build_pm_gpu_py -DWarpX_COMPUTE=CUDA -DWarpX_FFT=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_APP=OFF -DWarpX_PYTHON=ON -DWarpX_DIMS="1;2;RZ;3"
+         cmake -S . -B build_pm_gpu_py -DWarpX_COMPUTE=CUDA -DWarpX_FFT=ON -DWarpX_HEFFTE=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_APP=OFF -DWarpX_PYTHON=ON -DWarpX_DIMS="1;2;RZ;3"
          cmake --build build_pm_gpu_py -j 16 --target pip_install
 
    .. tab-item:: CPU Nodes
@@ -174,7 +174,7 @@ Use the following :ref:`cmake commands <building-cmake>` to compile the applicat
          cd $HOME/src/warpx
          rm -rf build_pm_cpu
 
-         cmake -S . -B build_pm_cpu -DWarpX_COMPUTE=OMP -DWarpX_FFT=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_DIMS="1;2;RZ;3"
+         cmake -S . -B build_pm_cpu -DWarpX_COMPUTE=OMP -DWarpX_FFT=ON -DWarpX_HEFFTE=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_DIMS="1;2;RZ;3"
          cmake --build build_pm_cpu -j 16
 
       The WarpX application executables are now in ``$HOME/src/warpx/build_pm_cpu/bin/``.
@@ -184,7 +184,7 @@ Use the following :ref:`cmake commands <building-cmake>` to compile the applicat
 
          rm -rf build_pm_cpu_py
 
-         cmake -S . -B build_pm_cpu_py -DWarpX_COMPUTE=OMP -DWarpX_FFT=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_APP=OFF -DWarpX_PYTHON=ON -DWarpX_DIMS="1;2;RZ;3"
+         cmake -S . -B build_pm_cpu_py -DWarpX_COMPUTE=OMP -DWarpX_FFT=ON -DWarpX_HEFFTE=ON -DWarpX_QED_TABLE_GEN=ON -DWarpX_APP=OFF -DWarpX_PYTHON=ON -DWarpX_DIMS="1;2;RZ;3"
          cmake --build build_pm_cpu_py -j 16 --target pip_install
 
 Now, you can :ref:`submit Perlmutter compute jobs <running-cpp-perlmutter>` for WarpX :ref:`Python (PICMI) scripts <usage-picmi>` (:ref:`example scripts <usage-examples>`).

Docs/source/usage/parameters.rst

@@ -2724,7 +2724,9 @@ In-situ capabilities can be used by turning on Sensei or Ascent (provided they a
 
 * ``<diag_name>.fields_to_plot`` (list of `strings`, optional)
     Fields written to output.
-    Possible scalar fields: ``part_per_cell`` ``rho`` ``phi`` ``F`` ``part_per_grid`` ``divE`` ``divB`` and ``rho_<species_name>``, where ``<species_name>`` must match the name of one of the available particle species. Note that ``phi`` will only be written out when do_electrostatic==labframe. Also, note that for ``<diag_name>.diag_type = BackTransformed``, the only scalar field currently supported is ``rho``.
+    Possible scalar fields: ``part_per_cell`` ``rho`` ``phi`` ``F`` ``part_per_grid`` ``divE`` ``divB`` ``rho_<species_name>`` and ``T_<species_name>``, where ``<species_name>`` must match the name of one of the available particle species.
+    ``T_<species_name>`` is the temperature in eV.
+    Note that ``phi`` will only be written out when do_electrostatic==labframe. Also, note that for ``<diag_name>.diag_type = BackTransformed``, the only scalar field currently supported is ``rho``.
     Possible vector field components in Cartesian geometry: ``Ex`` ``Ey`` ``Ez`` ``Bx`` ``By`` ``Bz`` ``jx`` ``jy`` ``jz``.
     Possible vector field components in RZ geometry: ``Er`` ``Et`` ``Ez`` ``Br`` ``Bt`` ``Bz`` ``jr`` ``jt`` ``jz``.
     The default is ``<diag_name>.fields_to_plot = Ex Ey Ez Bx By Bz jx jy jz`` in Cartesian geometry and ``<diag_name>.fields_to_plot = Er Et Ez Br Bt Bz jr jt jz`` in RZ geometry.

Examples/Physics_applications/laser_acceleration/PICMI_inputs_3d.py

@@ -55,7 +55,9 @@
 electrons = picmi.Species(
     particle_type = 'electron',
     name = 'electrons',
-    initial_distribution = uniform_distribution)
+    initial_distribution = uniform_distribution,
+    warpx_add_int_attributes = {'regionofinterest': "(z>12.0e-6) * (z<13.0e-6)"},
+    warpx_add_real_attributes = {'initialenergy': "ux*ux + uy*uy + uz*uz"})
 
 # Particles: beam electrons
 q_tot = 1e-12

Examples/Physics_applications/laser_acceleration/inputs_3d

@@ -55,7 +55,7 @@ electrons.do_continuous_injection = 1
 electrons.addIntegerAttributes = regionofinterest
 electrons.attribute.regionofinterest(x,y,z,ux,uy,uz,t) = "(z>12.0e-6) * (z<13.0e-6)"
 electrons.addRealAttributes = initialenergy
-electrons.attribute.initialenergy(x,y,z,ux,uy,uz,t) = " ux*ux + uy*uy + uz*uz"
+electrons.attribute.initialenergy(x,y,z,ux,uy,uz,t) = "ux*ux + uy*uy + uz*uz"
 
 #################################
 ############ LASER  #############

Examples/Physics_applications/laser_ion/PICMI_inputs_2d.py

@@ -91,18 +91,17 @@
     rms_velocity=[c*ux_th, 0., c*uz_th]  # thermal velocity spread in m/s
 )
 
-# TODO: add additional attributes orig_x and orig_z
 electrons = picmi.Species(
     particle_type='electron',
     name='electrons',
     initial_distribution=slab_with_ramp_dist_electrons,
 )
 
-# TODO: add additional attributes orig_x and orig_z
 hydrogen = picmi.Species(
     particle_type='proton',
     name='hydrogen',
-    initial_distribution=slab_with_ramp_dist_hydrogen
+    initial_distribution=slab_with_ramp_dist_hydrogen,
+    warpx_add_real_attributes = {"orig_x": "x", "orig_z": "z"}
 )
 
 # Laser

Examples/Tests/collision/inputs_3d

@@ -75,7 +75,7 @@ collision3.ndt = 10
 diagnostics.diags_names = diag1 diag_parser_filter diag_uniform_filter diag_random_filter
 diag1.intervals = 10
 diag1.diag_type = Full
-diag1.fields_to_plot = Ex Ey Ez Bx By Bz
+diag1.fields_to_plot = Ex Ey Ez Bx By Bz T_electron T_ion
 
 ## diag_parser_filter is a diag used to test the particle filter function.
 diag_parser_filter.intervals = 150:150:

Examples/Tests/embedded_circle/inputs_2d

@@ -5,7 +5,7 @@
 max_step = 11
 warpx.const_dt = 3.99e-13
 warpx.do_electrostatic = labframe
-warpx.self_fields_required_precision = 1e-06
+warpx.self_fields_required_precision = 2e-06
 warpx.eb_implicit_function = -((x-0.00005)**2+(z-0.00005)**2-1e-05**2)
 warpx.eb_potential(x,y,z,t) = -10
 warpx.self_fields_absolute_tolerance = 0.02

Examples/Tests/ionization/PICMI_inputs_2d.py

@@ -47,12 +47,14 @@
     fill_in = True)
 electrons = picmi.Species(
     particle_type = 'electron',
-    name = 'electrons')
+    name = 'electrons',
+    warpx_add_real_attributes = {'orig_z': 'z'})
 ions = picmi.Species(
     particle_type = 'N',
     name = 'ions',
     charge_state = 2,
-    initial_distribution = uniform_distribution)
+    initial_distribution = uniform_distribution,
+    warpx_add_real_attributes = {'orig_z': 'z'})
 
 # Field ionization
 nitrogen_ionization = picmi.FieldIonization(
@@ -88,7 +90,7 @@
     name = 'diag1',
     period = 10000,
     species = [electrons, ions],
-    data_list = ['ux', 'uy', 'uz', 'x', 'z', 'weighting'],
+    data_list = ['ux', 'uy', 'uz', 'x', 'z', 'weighting', 'orig_z'],
     write_dir = '.',
     warpx_file_prefix = 'Python_ionization_plt')
 field_diag = picmi.FieldDiagnostic(

Examples/Tests/ionization/analysis_ionization.py

@@ -52,7 +52,7 @@
     ad = ds.all_data()
 
     # Plot ions with ionization levels
-    species = 'ions';
+    species = 'ions'
     xi = ad[species, 'particle_position_x'].v
     zi = ad[species, 'particle_position_y'].v
     ii = ad[species, 'particle_ionizationLevel'].v
@@ -75,7 +75,7 @@
     plt.colorbar()
 
     # Plot electrons
-    species = 'electrons';
+    species = 'electrons'
     if species in [x[0] for x in ds.field_list]:
         xe = ad[species, 'particle_position_x'].v
         ze = ad[species, 'particle_position_y'].v
@@ -96,10 +96,11 @@
 # Check that the user runtime component (if it exists) worked as expected
 try:
     orig_z = ad['electrons', 'particle_orig_z'].v
-    assert np.all( (orig_z > 0) & (orig_z < 1.5e-5) )
+    print(f"orig_z: min = {np.min(orig_z)}, max = {np.max(orig_z)}")
+    assert np.all( (orig_z > 0.0) & (orig_z < 1.5e-5) )
     print('particle_orig_z has reasonable values')
 except yt.utilities.exceptions.YTFieldNotFound:
-    pass # Some of the tested script to not have the quantity orig_z
+    pass # The backtransformed diagnostic version of the test does not have orig_z
 
 test_name = os.path.split(os.getcwd())[1]
 checksumAPI.evaluate_checksum(test_name, filename)

Examples/Tests/ionization/inputs_2d_bf_rt

@@ -48,7 +48,7 @@ electrons.zmax =  50.e-6
 electrons.profile = constant
 electrons.density = 2.
 electrons.momentum_distribution_type = at_rest
-electrons.do_continuous_injection=1
+electrons.do_continuous_injection = 1
 
 lasers.names        = laser1
 laser1.profile      = Gaussian