Reference Particle: Store Charge & Mass (#229)

* Reference Particle: Store Charge & Mass Store the mass and charge of the reference particle. This simpliifes init logic and is useful for later I/O tasks of meta data. * Update Examples * Update Docs * Fix: Python Reference Lifetime
ECP-WarpX · Aug 31, 2022 · 4483972 · 4483972
1 parent c276e5b
commit 4483972
Show file tree

Hide file tree

Showing 14 changed files with 234 additions and 132 deletions.
diff --git a/docs/source/usage/python.rst b/docs/source/usage/python.rst
@@ -63,13 +63,13 @@ General
 
       This must come first, before particle beams and lattice elements are initialized.
 
-   .. py:method:: add_particles(qm_qeeV, charge_C, distr, npart)
+   .. py:method:: add_particles(charge_C, distr, npart)
 
       Generate and add n particles to the particle container.
+      Note: Set the reference particle properties (charge, mass, energy) first.
 
       Will also resize the geometry based on the updated particle distribution's extent and then redistribute particles in according AMReX grid boxes.
 
-      :param float qm_qeeV: charge/mass ratio (q_e/eV)
       :param float charge_C: bunch charge (C)
       :param distr: distribution function to draw from (object from :py:mod:`impactx.distribution`)
       :param int npart: number of particles to draw
@@ -224,7 +224,19 @@ Particles
 
       Read-only: Get reference particle beta*gamma
 
-   .. py:method:: set_energy_MeV(energy_MeV, massE_MeV)
+   .. py:property:: qm_qeeV
+
+      Read-only: Get reference particle charge to mass ratio (elementary charge/eV)
+
+   .. py:method:: set_charge_qe(charge_qe)
+
+      Write-only: Set reference particle charge in (positive) elementary charges.
+
+   .. py:method:: set_mass_MeV(massE)
+
+      Write-only: Set reference particle rest mass (MeV/c^2).
+
+   .. py:method:: set_energy_MeV(energy_MeV)
 
       Write-only: Set reference particle energy.
 

diff --git a/examples/chicane/run_chicane.py b/examples/chicane/run_chicane.py
@@ -23,11 +23,14 @@
 # load a 5 GeV electron beam with an initial
 # normalized transverse rms emittance of 1 um
 energy_MeV = 5.0e3  # reference energy
-charge_C = 1.0e-9  # used with space charge
-mass_MeV = 0.510998950
-qm_qeeV = -1.0e-6 / mass_MeV  # charge/mass
+bunch_charge_C = 1.0e-9  # used with space charge
 npart = 10000  # number of macro particles
 
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_energy_MeV(energy_MeV)
+
+#   particle bunch
 distr = distribution.Waterbag(
     sigmaX=2.2951017632e-5,
     sigmaY=1.3084093142e-5,
@@ -39,10 +42,7 @@
     muypy=0.933345606203060,
     mutpt=0.999999961419755,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(energy_MeV, mass_MeV)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 # design the accelerator lattice
 ns = 25  # number of slices per ds in the element

diff --git a/examples/fodo/run_fodo.py b/examples/fodo/run_fodo.py
@@ -23,11 +23,14 @@
 # load a 2 GeV electron beam with an initial
 # unnormalized rms emittance of 2 nm
 energy_MeV = 2.0e3  # reference energy
-charge_C = 1.0e-9  # used with space charge
-mass_MeV = 0.510998950
-qm_qeeV = -1.0e-6 / mass_MeV  # charge/mass
+bunch_charge_C = 1.0e-9  # used with space charge
 npart = 10000  # number of macro particles
 
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_energy_MeV(energy_MeV)
+
+#   particle bunch
 distr = distribution.Waterbag(
     sigmaX=3.9984884770e-5,
     sigmaY=3.9984884770e-5,
@@ -39,10 +42,7 @@
     muypy=0.846574929020762,
     mutpt=0.0,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(energy_MeV, mass_MeV)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 # design the accelerator lattice
 ns = 25  # number of slices per ds in the element

diff --git a/examples/iota_lattice/run_iotalattice.py b/examples/iota_lattice/run_iotalattice.py
@@ -22,11 +22,14 @@
 
 # init particle beam
 energy_MeV = 2.5
-charge_C = 1.0e-9  # assign zero weighting to particles
-mass_MeV = 938.27208816
-qm_qeeV = 1.0e-6 / mass_MeV
+bunch_charge_C = 1.0e-9  # used with space charge
 npart = 10000
 
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_energy_MeV(energy_MeV)
+
+#   particle bunch
 distr = distribution.Waterbag(
     sigmaX=1.588960728035e-3,
     sigmaY=2.496625268437e-3,
@@ -35,10 +38,7 @@
     sigmaPy=1.802433091137e-3,
     sigmaPt=0.0,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(energy_MeV, mass_MeV)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 # init accelerator lattice
 ns = 10  # number of slices per ds in the element

diff --git a/examples/iota_lens/run_iotalens.py b/examples/iota_lens/run_iotalens.py
@@ -22,26 +22,23 @@
 
 # load a 2.5 MeV proton beam
 energy_MeV = 2.5  # reference energy
-charge_C = 1.0e-9  # used with space charge
-mass_MeV = 938.27208816
-qm_qeeV = 1.0e-6 / mass_MeV  # charge/mass
+bunch_charge_C = 1.0e-9  # used with space charge
 npart = 10000  # number of macro particles
 
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_energy_MeV(energy_MeV)
+
+#   particle bunch
 distr = distribution.Waterbag(
     sigmaX=2.0e-3,
     sigmaY=2.0e-3,
     sigmaT=1.0e-3,
     sigmaPx=3.0e-4,
     sigmaPy=3.0e-4,
     sigmaPt=0.0,
-    muxpx=0.0,
-    muypy=0.0,
-    mutpt=0.0,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(energy_MeV, mass_MeV)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 constEnd = elements.ConstF(ds=0.0025, kx=1.0, ky=1.0, kt=1.0e-12)
 nllens = elements.NonlinearLens(knll=2.0e-7, cnll=0.01)

diff --git a/examples/kurth/run_kurth.py b/examples/kurth/run_kurth.py
@@ -24,26 +24,23 @@
 # unnormalized rms emittance of 1 um in each
 # coordinate plane
 energy_MeV = 2.0e3  # reference energy
-charge_C = 1.0e-9  # used with space charge
-mass_MeV = 938.27208816
-qm_qeeV = 1.0e-6 / mass_MeV  # charge/mass
+bunch_charge_C = 1.0e-9  # used with space charge
 npart = 10000  # number of macro particles
 
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(1.0).set_mass_MeV(938.27208816).set_energy_MeV(energy_MeV)
+
+#   particle bunch
 distr = distribution.Kurth6D(
     sigmaX=1.0e-3,
     sigmaY=1.0e-3,
     sigmaT=3.369701494258956e-4,
     sigmaPx=1.0e-3,
     sigmaPy=1.0e-3,
     sigmaPt=2.9676219145931020e-3,
-    muxpx=0.0,
-    muypy=0.0,
-    mutpt=0.0,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(energy_MeV, mass_MeV)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 # design the accelerator lattice: here we just assign a single element
 sim.lattice.append(elements.ConstF(ds=2.0, kx=1.0, ky=1.0, kt=1.0))

diff --git a/examples/multipole/run_multipole.py b/examples/multipole/run_multipole.py
@@ -23,11 +23,14 @@
 # load a 2 GeV electron beam with an initial
 # unnormalized rms emittance of  nm
 energy_MeV = 2.0e3  # reference energy
-charge_C = 1.0e-9  # used without space charge
-mass_MeV = 0.510998950
-qm_qeeV = -1.0e-6 / mass_MeV  # charge/mass
+bunch_charge_C = 1.0e-9  # used without space charge
 npart = 10000  # number of macro particles
 
+#   reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_energy_MeV(energy_MeV)
+
+#   particle bunch
 distr = distribution.Waterbag(
     sigmaX=4.0e-3,
     sigmaY=4.0e-3,
@@ -36,10 +39,7 @@
     sigmaPy=3.0e-4,
     sigmaPt=2.0e-3,
 )
-sim.add_particles(qm_qeeV, charge_C, distr, npart)
-
-# set the energy in the reference particle
-sim.particle_container().ref_particle().set_energy_MeV(energy_MeV, mass_MeV)
+sim.add_particles(bunch_charge_C, distr, npart)
 
 # design the accelerator lattice
 multipole = [

diff --git a/src/ImpactX.H b/src/ImpactX.H
@@ -71,14 +71,12 @@ namespace impactx
          * distribution's extent and then redistribute particles in according
          * AMReX grid boxes.
          *
-         * @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
          */
         void
         add_particles (
-            amrex::ParticleReal qm,
             amrex::ParticleReal bunch_charge,
             distribution::KnownDistributions distr,
             int npart

diff --git a/src/initialization/InitDistribution.cpp b/src/initialization/InitDistribution.cpp
@@ -24,14 +24,21 @@ namespace impactx
 {
     void
     ImpactX::add_particles (
-        amrex::ParticleReal qm,
         amrex::ParticleReal bunch_charge,
         distribution::KnownDistributions distr,
         int npart
     )
     {
         BL_PROFILE("ImpactX::add_particles");
 
+        auto const & ref = m_particle_container->GetRefParticle();
+        AMREX_ASSERT_WITH_MESSAGE(ref.charge_qe() != 0.0,
+                                  "add_particles: Reference particle charge not yet set!");
+        AMREX_ASSERT_WITH_MESSAGE(ref.mass_MeV() != 0.0,
+                                  "add_particles: Reference particle mass not yet set!");
+        AMREX_ASSERT_WITH_MESSAGE(ref.energy_MeV() != 0.0,
+                                  "add_particles: Reference particle energy not yet set!");
+
         amrex::Vector<amrex::ParticleReal> x, y, t;
         amrex::Vector<amrex::ParticleReal> px, py, pt;
         amrex::ParticleReal ix, iy, it, ipx, ipy, ipt;
@@ -70,7 +77,8 @@ namespace impactx
 
         int const lev = 0;
         m_particle_container->AddNParticles(lev, x, y, t, px, py, pt,
-                                            qm, bunch_charge * rel_part_this_proc);
+                                            ref.qm_qeeV(),
+                                            bunch_charge * rel_part_this_proc);
 
         // Resize the mesh to fit the spatial extent of the beam and then
         // redistribute particles, so they reside on the MPI rank that is
@@ -97,16 +105,24 @@ namespace impactx
         std::string particle_type;  // Particle type
         pp_dist.get("particle", particle_type);
 
-        amrex::ParticleReal qm = 0.0; // charge/mass ratio (q_e/eV)
+        amrex::ParticleReal qe;     // charge (elementary charge)
+        amrex::ParticleReal massE;  // MeV/c^2
         if(particle_type == "electron") {
-            qm = -1.0/0.510998950e6;
+            qe = -1.0;
+            massE = 0.510998950;
         } else if(particle_type == "proton") {
-            qm = 1.0/938.27208816e6;
+            qe = 1.0;
+            massE = 938.27208816;
         }
-        else {
-            qm = 0.0;
+        else {  // default to electron
+            qe = -1.0;
+            massE = 0.510998950;
         }
 
+        // set charge and mass and energy of ref particle
+        m_particle_container->GetRefParticle()
+            .set_charge_qe(qe).set_mass_MeV(massE).set_energy_MeV(energy);
+
         int npart = 1;  // Number of simulation particles
         pp_dist.get("npart", npart);
 
@@ -134,7 +150,7 @@ namespace impactx
               sigpx, sigpy, sigpt,
               muxpx, muypy, mutpt));
 
-          add_particles(qm, bunch_charge, waterbag, npart);
+          add_particles(bunch_charge, waterbag, npart);
 
         } else if (distribution_type == "kurth6d") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -154,7 +170,7 @@ namespace impactx
             sigpx, sigpy, sigpt,
             muxpx, muypy, mutpt));
 
-          add_particles(qm, bunch_charge, kurth6D, npart);
+          add_particles(bunch_charge, kurth6D, npart);
 
         } else if (distribution_type == "gaussian") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -174,7 +190,7 @@ namespace impactx
             sigpx, sigpy, sigpt,
             muxpx, muypy, mutpt));
 
-          add_particles(qm, bunch_charge, gaussian, npart);
+          add_particles(bunch_charge, gaussian, npart);
 
         } else if (distribution_type == "kvdist") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -194,7 +210,7 @@ namespace impactx
             sigpx, sigpy, sigpt,
             muxpx, muypy, mutpt));
 
-          add_particles(qm, bunch_charge, kvDist, npart);
+          add_particles(bunch_charge, kvDist, npart);
 
         } else if (distribution_type == "kurth4d") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
@@ -214,7 +230,7 @@ namespace impactx
             sigpx, sigpy, sigpt,
             muxpx, muypy, mutpt));
 
-          add_particles(qm, bunch_charge, kurth4D, npart);
+          add_particles(bunch_charge, kurth4D, npart);
         } else if (distribution_type == "semigaussian") {
           amrex::ParticleReal sigx,sigy,sigt,sigpx,sigpy,sigpt;
           amrex::ParticleReal muxpx = 0.0, muypy = 0.0, mutpt = 0.0;
@@ -233,24 +249,11 @@ namespace impactx
             sigpx, sigpy, sigpt,
             muxpx, muypy, mutpt));
 
-          add_particles(qm, bunch_charge, semigaussian, npart);
+          add_particles(bunch_charge, semigaussian, npart);
         } else {
             amrex::Abort("Unknown distribution: " + distribution_type);
         }
 
-        // reference particle
-        amrex::ParticleReal massE;  // MeV
-        if (particle_type == "electron") {
-            massE = 0.510998950;
-        } else if (particle_type == "proton") {
-            massE = 938.27208816;
-        } else {
-            massE = 0.510998950;  // default to electron
-        }
-
-        // set the energy in the reference particle
-        m_particle_container->GetRefParticle().set_energy_MeV(energy, massE);
-
         // print information on the initialized beam
         amrex::Print() << "Beam kinetic energy (MeV): " << energy << std::endl;
         amrex::Print() << "Bunch charge (C): " << bunch_charge << std::endl;