Python: Init Beam Distributions

examples/CMakeLists.txt

@@ -18,6 +18,20 @@ if(ImpactX_MPI)
     )
 endif()
 
+function(impactx_test_set_pythonpath test_name)
+    if(WIN32)
+        string(REPLACE ";" "\\;" WIN_PYTHONPATH "$ENV{PYTHONPATH}")
+        string(REGEX REPLACE "/" "\\\\" WIN_PYTHON_OUTPUT_DIRECTORY ${CMAKE_PYTHON_OUTPUT_DIRECTORY})
+        set_tests_properties(TEST ${test_name}
+            PROPERTIES ENVIRONMENT "PYTHONPATH=${WIN_PYTHON_OUTPUT_DIRECTORY}\;${WIN_PYTHONPATH}"
+        )
+    else()
+        set_tests_properties(${test_name}
+            PROPERTIES ENVIRONMENT "PYTHONPATH=${CMAKE_PYTHON_OUTPUT_DIRECTORY}:$ENV{PYTHONPATH}"
+        )
+    endif()
+endfunction()
+
 
 # FODO Cell ###################################################################
 #
@@ -42,17 +56,17 @@ set_tests_properties(FODO.plot PROPERTIES DEPENDS "FODO.run")
 #
 if(ImpactX_MPI)
     add_test(NAME FODO.MPI.run
-            COMMAND ${MPI_TEST_EXE}
-                $<TARGET_FILE:app> ${ImpactX_SOURCE_DIR}/examples/fodo/input_fodo.in
-            WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+             COMMAND ${MPI_TEST_EXE}
+                 $<TARGET_FILE:app> ${ImpactX_SOURCE_DIR}/examples/fodo/input_fodo.in
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
     )
     add_test(NAME FODO.MPI.analysis
-            COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py
-            WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+             COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
     )
     add_test(NAME FODO.MPI.plot
-            COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png
-            WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+             COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
     )
 
     set_tests_properties(FODO.MPI.run PROPERTIES ENVIRONMENT "OMP_NUM_THREADS=1")
@@ -61,6 +75,53 @@ if(ImpactX_MPI)
 endif()
 
 
+# Python: FODO Cell ###########################################################
+#
+if(ImpactX_PYTHON)
+    add_test(NAME FODO.py.run
+             COMMAND ${Python_EXECUTABLE}
+                 ${ImpactX_SOURCE_DIR}/examples/fodo/run_fodo.py
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+    )
+    add_test(NAME FODO.py.analysis
+             COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+    )
+    add_test(NAME FODO.py.plot
+             COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+    )
+
+    set_tests_properties(FODO.py.analysis PROPERTIES DEPENDS "FODO.py.run")
+    set_tests_properties(FODO.py.plot PROPERTIES DEPENDS "FODO.py.run")
+    impactx_test_set_pythonpath(FODO.py.run)
+endif()
+
+
+# Python: MPI-parallel FODO Cell ##############################################
+#
+if(ImpactX_PYTHON AND ImpactX_MPI)
+    add_test(NAME FODO.py.MPI.run
+             COMMAND ${MPI_TEST_EXE} ${Python_EXECUTABLE}
+                 ${ImpactX_SOURCE_DIR}/examples/fodo/run_fodo.py
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+    )
+    add_test(NAME FODO.py.MPI.analysis
+             COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/analysis_fodo.py
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+    )
+    add_test(NAME FODO.py.MPI.plot
+             COMMAND ${ImpactX_SOURCE_DIR}/examples/fodo/plot_fodo.py --save-png
+             WORKING_DIRECTORY ${ImpactX_RUNTIME_OUTPUT_DIRECTORY}
+    )
+
+    set_tests_properties(FODO.py.MPI.run PROPERTIES ENVIRONMENT "OMP_NUM_THREADS=1")
+    set_tests_properties(FODO.py.MPI.analysis PROPERTIES DEPENDS "FODO.py.MPI.run")
+    set_tests_properties(FODO.py.MPI.plot PROPERTIES DEPENDS "FODO.py.MPI.run")
+    impactx_test_set_pythonpath(FODO.py.MPI.run)
+endif()
+
+
 # Chicane #####################################################################
 #
 add_test(NAME chicane.run

examples/fodo/README.rst

@@ -19,7 +19,39 @@ The second moments of the particle distribution after the FODO cell should coinc
 In this test, the initial and final values of :math:`\sigma_x`, :math:`\sigma_y`, :math:`\sigma_t`, :math:`\epsilon_x`, :math:`\epsilon_y`, and :math:`\epsilon_t` must agree with nominal values.
 
 
+Run
+---
+
+This example an either be run with an inputs file (``impactx input_fodo.in``):
 
 .. literalinclude:: input_fodo.in
    :language: ini
    :caption: You can copy this file from ``examples/fodo/input_fodo.in``.
+
+Or as a Python script (``python3 run_fodo.py``):
+
+.. literalinclude:: run_fodo.py
+   :language: python3
+   :caption: You can copy this file from ``examples/fodo/run_fodo.py``.
+
+Both execution modes can also be prefixed with an MPI executor, such as ``mpiexec -n 4 ...`` or ``srun -n 4 ...``, depending on the system.
+
+
+Analyze
+-------
+
+We run the following script to analyze correctness:
+
+.. literalinclude:: analysis_fodo.py
+   :language: python3
+   :caption: You can copy this file from ``examples/fodo/analysis_fodo.py``.
+
+
+Visualize
+---------
+
+You can run the following script to visualize the beam evolution over time:
+
+.. literalinclude:: plot_fodo.py
+   :language: python3
+   :caption: You can copy this file from ``examples/fodo/plot_fodo.py``.

examples/fodo/run_fodo.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+#
+# Copyright 2022 ImpactX contributors
+# Authors: Axel Huebl, Chad Mitchell
+# License: BSD-3-Clause-LBNL
+#
+# -*- coding: utf-8 -*-
+
+import amrex
+from impactx import ImpactX, RefPart, distribution, elements
+
+impactX = ImpactX()
+
+impactX.set_particle_shape(2)
+impactX.set_diags_slice_step_diagnostics(True)
+impactX.init_grids()
+
+# init particle beam
+energy_MeV = 2.0e3
+charge_C = 0.0  # assign zero weighting to particles
+mass_MeV = 0.510998950
+qm_qeeV = -1.0/0.510998950e6
+npart = 10000
+
+distr = distribution.Waterbag(
+    sigmaX = 3.9984884770e-5,
+    sigmaY = 3.9984884770e-5,
+    sigmaT = 1.0e-3,
+    sigmaPx = 2.6623538760e-5,
+    sigmaPy = 2.6623538760e-5,
+    sigmaPt = 2.0e-3,
+    muxpx = -0.846574929020762,
+    muypy = 0.846574929020762,
+    mutpt = 0.0)
+distribution.generate_add_particles(
+    impactX.particle_container(), qm_qeeV, charge_C, distr, npart)
+
+# init reference particle
+refPart = RefPart()
+# make the next two lines a helper function?
+refPart.pt = -energy_MeV / mass_MeV - 1.0
+refPart.pz = (refPart.pt**2 - 1.0)**0.5
+impactX.particle_container().set_ref_particle(refPart)
+
+# init accelerator lattice
+ns = 25  # number of slices per ds in the element
+fodo = [
+    elements.Drift(ds=0.25, nslice=ns),
+    elements.Quad(ds=1.0, k=1.0, nslice=ns),
+    elements.Drift(ds=0.5, nslice=ns),
+    elements.Quad(ds=1.0, k=-1.0, nslice=ns),
+    elements.Drift(ds=0.25, nslice=ns)
+]
+#  assign a fodo segment
+impactX.lattice.extend(fodo)
+
+# run simulation
+impactX.evolve()
+
+# clean shutdown
+del impactX
+amrex.finalize()

src/ImpactX.cpp

@@ -59,6 +59,20 @@ namespace impactx
         //   before we start the evolve loop, we are in "step 0" (initial state)
         int global_step = 0;
 
+        // count particles - if no particles are found in our particle container, then a lot of
+        // AMReX routines over ParIter won't work and we have nothing to do here anyways
+        {
+            int const nLevelPC = finestLevel();
+            amrex::Long nParticles = 0;
+            for (int lev = 0; lev <= nLevelPC; ++lev) {
+                nParticles += m_particle_container->NumberOfParticlesAtLevel(lev);
+            }
+            if (nParticles == 0) {
+                amrex::Abort("No particles found. Cannot run evolve without a beam.");
+                return;
+            }
+        }
+
         amrex::ParmParse pp_diag("diag");
         int file_min_digits = 6;
         {

src/initialization/CMakeLists.txt

@@ -1,6 +1,7 @@
 target_sources(ImpactX
   PRIVATE
     AmrCoreData.cpp
+    InitAMReX.cpp
     InitDistribution.cpp
     InitElement.cpp
     InitMeshRefinement.cpp

src/initialization/InitAMReX.H

@@ -0,0 +1,32 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#ifndef IMPACT_DEFAULT_INIT_AMREX_H
+#define IMPACT_DEFAULT_INIT_AMREX_H
+
+#include "AmrCoreData.H"
+
+
+namespace impactx::initialization
+{
+    /** Initialize AMReX
+     *
+     * Initializes AMReX if not already done.
+     */
+    void default_init_AMReX (int argc, char* argv[]);
+
+    /** Initialize AMReX
+     *
+     * Initializes AMReX if not already done.
+     */
+    void default_init_AMReX ();
+
+} // namespace impactx::initialization
+
+#endif // IMPACT_DEFAULT_INIT_AMREX_H

src/initialization/InitAMReX.cpp

@@ -0,0 +1,47 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl, Chad Mitchell, Ji Qiang
+ * License: BSD-3-Clause-LBNL
+ */
+#include "InitAMReX.H"
+
+#include "initialization/InitParser.H"
+
+#include <AMReX.H>
+#include <AMReX_ParallelDescriptor.H>
+
+#if defined(AMREX_USE_MPI)
+#   include <mpi.h>
+#else
+#   include <AMReX_ccse-mpi.H>
+#endif
+
+
+namespace impactx::initialization
+{
+    void
+    default_init_AMReX (int argc, char* argv[])
+    {
+        if (!amrex::Initialized())
+        {
+            bool const build_parm_parse = true;
+            amrex::Initialize(
+                    argc,
+                    argv,
+                    build_parm_parse,
+                    MPI_COMM_WORLD,
+                    impactx::initialization::overwrite_amrex_parser_defaults
+            );
+        }
+    }
+
+    void
+    default_init_AMReX ()
+    {
+        default_init_AMReX(0, nullptr);
+    }
+} // namespace impactx::initialization

src/initialization/InitDistribution.H

@@ -0,0 +1,79 @@
+/* Copyright 2022 The Regents of the University of California, through Lawrence
+ *           Berkeley National Laboratory (subject to receipt of any required
+ *           approvals from the U.S. Dept. of Energy). All rights reserved.
+ *
+ * This file is part of ImpactX.
+ *
+ * Authors: Axel Huebl
+ * License: BSD-3-Clause-LBNL
+ */
+#include "particles/ImpactXParticleContainer.H"
+
+#include <AMReX.H>
+#include <AMReX_REAL.H>
+#include <AMReX_ParmParse.H>
+#include <AMReX_Print.H>
+#include <AMReX_Random.H>
+#include <AMReX_Vector.H>
+
+#include <string>
+
+
+namespace impactx
+{
+    /** Generate and add n particles to a particle container
+     *
+     * @tparam T_Distribution distribution function to draw from (type)
+     * @param pc a particle container to add particles to
+     * @param qm charge/mass ratio (q_e/eV)
+     * @param bunch_charge bunch charge (C)
+     * @param distr distribution function to draw from (object)
+     * @param npart number of particles to draw
+     */
+    template <typename T_Distribution>
+    void
+    generate_add_particles (
+        impactx::ImpactXParticleContainer & pc,
+        amrex::ParticleReal qm,
+        amrex::ParticleReal bunch_charge,
+        T_Distribution & distr,
+        int npart
+    )
+    {
+        amrex::Vector<amrex::ParticleReal> x, y, t;
+        amrex::Vector<amrex::ParticleReal> px, py, pt;
+        amrex::ParticleReal ix, iy, it, ipx, ipy, ipt;
+        amrex::RandomEngine rng;
+
+        // Logic: We initialize 1/Nth of particles, independent of their
+        // position, per MPI rank. We then measure the distribution's spatial
+        // extent, create a grid, resize it to fit the beam, and then
+        // redistribute particles so that they reside on the correct MPI rank.
+        int myproc = amrex::ParallelDescriptor::MyProc();
+        int nprocs = amrex::ParallelDescriptor::NProcs();
+        int navg = npart / nprocs;
+        int nleft = npart - navg * nprocs;
+        int npart_this_proc = (myproc < nleft) ? navg+1 : navg;
+
+        x.reserve(npart_this_proc);
+        y.reserve(npart_this_proc);
+        t.reserve(npart_this_proc);
+        px.reserve(npart_this_proc);
+        py.reserve(npart_this_proc);
+        pt.reserve(npart_this_proc);
+
+        for (amrex::Long i = 0; i < npart_this_proc; ++i) {
+            distr(ix, iy, it, ipx, ipy, ipt, rng);
+            x.push_back(ix);
+            y.push_back(iy);
+            t.push_back(it);
+            px.push_back(ipx);
+            py.push_back(ipy);
+            pt.push_back(ipt);
+        }
+
+        int const lev = 0;
+        pc.AddNParticles(lev, x, y, t, px, py, pt,
+                         qm, bunch_charge);
+    }
+} // namespace impactx