Add focusing position to Gaussian beam initialization (#4639)

* add focusing position * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * removed prints * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix stashed * generalized to focal distance * first test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix formatting in benchmark checksums * updated test and start docs * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * updated test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * focal distance is optional * updated test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixed unused variables vbulk * updated docs * updated test * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reset benchmark * vx vy mistake * fix conversion error clang tidy * fix clang tidy error * fix typo * Update Docs/source/usage/parameters.rst Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Examples/Tests/gaussian_beam/analysis_focusing_beam.py Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Examples/Tests/gaussian_beam/analysis_focusing_beam.py Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Source/Particles/PhysicalParticleContainer.cpp Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Examples/Tests/gaussian_beam/analysis_focusing_beam.py Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Source/Particles/PhysicalParticleContainer.cpp Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update Source/Particles/PhysicalParticleContainer.cpp Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Source/Particles/PhysicalParticleContainer.cpp Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * integrate comments by remi * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove unused variable * reset benchmark * fix v_dot_n * Update Source/Particles/PhysicalParticleContainer.cpp Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Source/Particles/PhysicalParticleContainer.cpp Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * Update Source/Particles/PhysicalParticleContainer.cpp Co-authored-by: Remi Lehe <remi.lehe@normalesup.org> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tools <warpx@lbl.gov> Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>
BLAST-WarpX · Feb 23, 2024 · d45d173 · d45d173
1 parent 07738d3
commit d45d173
Show file tree

Hide file tree

Showing 9 changed files with 282 additions and 10 deletions.
diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst
@@ -779,7 +779,7 @@ Particle initialization
 
       * ``<species_name>.q_tot`` (beam charge),
 
-      * ``<species_name>.npart`` (number of particles in the beam),
+      * ``<species_name>.npart`` (number of macroparticles in the beam),
 
       * ``<species_name>.x/y/z_m`` (average position in `x/y/z`),
 
@@ -797,6 +797,10 @@ Particle initialization
       If set to 4, symmetrization is in the x and y direction, (x,y) (-x,y) (x,-y) (-x,-y).
       If set to 8, symmetrization is also done with x and y exchanged, (y,x), (-y,x), (y,-x), (-y,-x)).
 
+      * ``<species_name>.focal_distance`` (optional, distance between the beam centroid and the position of the focal plane of the beam, along the direction of the beam mean velocity; space charge is ignored in the initialization of the particles)
+
+      If ``<species_name>.focal_distance`` is specified, ``x_rms``, ``y_rms`` and ``z_rms`` are the size of the beam in the focal plane. Since the beam is not necessarily initialized close to its focal plane, the initial size of the beam will differ from ``x_rms``, ``y_rms``, ``z_rms``.
+
     * ``external_file``: Inject macroparticles with properties (mass, charge, position, and momentum - :math:`\gamma \beta m c`) read from an external openPMD file.
       With it users can specify the additional arguments:
 

diff --git a/Examples/Tests/gaussian_beam/analysis_focusing_beam.py b/Examples/Tests/gaussian_beam/analysis_focusing_beam.py
@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+
+# Copyright 2024 Arianna Formenti
+#
+# This file is part of WarpX.
+#
+# License: BSD-3-Clause-LBNL
+
+
+import os
+import sys
+
+import numpy as np
+from scipy.constants import c, eV, m_e, micro, nano
+
+sys.path.insert(1, '../../../../warpx/Regression/Checksum/')
+
+import checksumAPI
+from openpmd_viewer import OpenPMDTimeSeries
+
+GeV=1e9*eV
+energy = 125.*GeV
+gamma = energy/(m_e*c**2)
+sigmax = 516.0*nano
+sigmay = 7.7*nano
+sigmaz = 300.*micro
+nz = 256
+Lz = 20*sigmaz
+gridz = np.linspace(-0.5*Lz, 0.5*Lz, nz)
+tol = gridz[1] - gridz[0]
+emitx = 50*micro
+emity = 20*nano
+focal_distance = 4*sigmaz
+
+def s(z, sigma0, emit):
+    '''The theoretical size of a focusing beam (in the absence of space charge),
+    at position z, given its emittance and size at focus.'''
+    return np.sqrt(sigma0**2 + emit**2 * (z - focal_distance)**2 / sigma0**2)
+
+filename = sys.argv[1]
+
+ts = OpenPMDTimeSeries('./diags/openpmd/')
+
+x, y, z, w, = ts.get_particle( ['x', 'y', 'z', 'w'], species='beam1', iteration=0, plot=False)
+
+imin = np.argmin(np.sqrt((gridz+0.8*focal_distance)**2))
+imax = np.argmin(np.sqrt((gridz-0.8*focal_distance)**2))
+
+sx, sy = [], []
+# Compute the size of the beam in each z slice
+subgrid = gridz[imin:imax]
+for d in subgrid:
+    i = np.sqrt((z - d)**2) < tol
+    if (np.sum(i)!=0):
+        mux = np.average(x[i], weights=w[i])
+        muy = np.average(y[i], weights=w[i])
+        sx.append(np.sqrt(np.average((x[i]-mux)**2, weights=w[i])))
+        sy.append(np.sqrt(np.average((y[i]-muy)**2, weights=w[i])))
+
+# Theoretical prediction for the size of the beam in each z slice
+sx_theory = s(subgrid, sigmax, emitx/gamma)
+sy_theory = s(subgrid, sigmay, emity/gamma)
+
+assert(np.allclose(sx, sx_theory, rtol=0.051, atol=0))
+assert(np.allclose(sy, sy_theory, rtol=0.038, atol=0))
+
+test_name = os.path.split(os.getcwd())[1]
+checksumAPI.evaluate_checksum(test_name, filename)
diff --git a/Examples/Tests/gaussian_beam/inputs_focusing_beam b/Examples/Tests/gaussian_beam/inputs_focusing_beam
@@ -0,0 +1,120 @@
+#################################
+########## MY CONSTANTS #########
+#################################
+my_constants.mc2 = m_e*clight*clight
+my_constants.nano = 1.0e-9
+my_constants.micro = 1.e-6
+my_constants.GeV = q_e*1.e9
+
+# BEAMS
+my_constants.energy = 125.*GeV
+my_constants.gamma = energy/mc2
+my_constants.npart = 2.e10
+my_constants.nmacropart = 1000000
+my_constants.charge = q_e * npart
+
+my_constants.sigmax = 516.0*nano
+my_constants.sigmay = 7.7*nano
+my_constants.sigmaz = 300.*micro
+
+my_constants.mux = 0.0
+my_constants.muy = 0.0
+my_constants.muz = 0.0
+
+my_constants.ux = 0.0
+my_constants.uy = 0.0
+my_constants.uz = gamma
+
+my_constants.emitx = 50*micro
+my_constants.emity = 20*nano
+my_constants.emitz = 0.
+
+my_constants.dux = emitx / sigmax
+my_constants.duy = emity / sigmay
+my_constants.duz = emitz / sigmaz
+
+my_constants.focal_distance = 4*sigmaz
+
+# BOX
+my_constants.Lx = 20*sigmax
+my_constants.Ly = 20*sigmay
+my_constants.Lz = 20*sigmaz
+my_constants.nx = 256
+my_constants.ny = 256
+my_constants.nz = 256
+
+
+
+#################################
+####### GENERAL PARAMETERS ######
+#################################
+max_step = 0
+amr.n_cell = nx ny nz
+amr.max_grid_size = 256
+amr.blocking_factor = 2
+amr.max_level = 0
+geometry.dims = 3
+geometry.prob_lo = -0.5*Lx -0.5*Ly -0.5*Lz
+geometry.prob_hi =  0.5*Lx  0.5*Ly  0.5*Lz
+
+#################################
+######## BOUNDARY CONDITION #####
+#################################
+boundary.field_lo = PEC PEC PEC
+boundary.field_hi = PEC PEC PEC
+boundary.particle_lo = Absorbing Absorbing Absorbing
+boundary.particle_hi = Absorbing Absorbing Absorbing
+
+#################################
+############ NUMERICS ###########
+#################################
+algo.particle_shape = 3
+
+#################################
+########### PARTICLES ###########
+#################################
+particles.species_names = beam1
+
+beam1.species_type = electron
+beam1.injection_style = gaussian_beam
+beam1.x_rms = sigmax
+beam1.y_rms = sigmay
+beam1.z_rms = sigmaz
+beam1.x_m = muy
+beam1.y_m = mux
+beam1.z_m = muz
+beam1.focal_distance = focal_distance
+beam1.npart = nmacropart
+beam1.q_tot = -charge
+
+beam1.momentum_distribution_type = gaussian
+beam1.ux_m = ux
+beam1.uy_m = uy
+beam1.uz_m = uz
+beam1.ux_th = dux
+beam1.uy_th = duy
+beam1.uz_th = duz
+
+#################################
+######### DIAGNOSTICS ###########
+#################################
+# FULL
+diagnostics.diags_names =  diag1 openpmd
+
+diag1.intervals = 1
+diag1.diag_type = Full
+diag1.write_species = 1
+diag1.species = beam1
+diag1.fields_to_plot = rho_beam1
+diag1.format = plotfile
+diag1.dump_last_timestep = 1
+
+
+openpmd.intervals = 1
+openpmd.diag_type = Full
+openpmd.write_species = 1
+openpmd.species = beam1
+openpmd.beam1.variables = w x y z
+openpmd.fields_to_plot = none
+openpmd.format = openpmd
+openpmd.dump_last_timestep = 1
diff --git a/Regression/Checksum/benchmarks_json/focusing_gaussian_beam.json b/Regression/Checksum/benchmarks_json/focusing_gaussian_beam.json
@@ -0,0 +1,14 @@
+{
+  "lev=0": {
+    "rho_beam1": 5637659324885.731
+  },
+  "beam1": {
+    "particle_momentum_x": 2.113335069047891e-14,
+    "particle_momentum_y": 5.653782702472731e-16,
+    "particle_momentum_z": 6.68023056405556e-11,
+    "particle_position_x": 0.5640698094900541,
+    "particle_position_y": 0.011947117763454793,
+    "particle_position_z": 239.4325333696459,
+    "particle_weight": 19999620000.0
+  }
+}
diff --git a/Regression/WarpX-tests.ini b/Regression/WarpX-tests.ini
@@ -4627,3 +4627,21 @@ compileTest = 0
 doVis = 0
 compareParticles = 1
 analysisRoutine = Examples/Tests/Implicit/analysis_1d.py
+
+[focusing_gaussian_beam]
+buildDir = .
+inputFile = Examples/Tests/gaussian_beam/inputs_focusing_beam
+runtime_params =
+dim = 3
+addToCompileString = USE_OPENPMD=TRUE
+cmakeSetupOpts = -DWarpX_DIMS=3 -DWarpX_OPENPMD=ON
+restartTest = 0
+useMPI = 1
+numprocs = 2
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+compareParticles = 0
+doComparison = 0
+analysisRoutine = Examples/Tests/gaussian_beam/analysis_focusing_beam.py
diff --git a/Source/Initialization/PlasmaInjector.H b/Source/Initialization/PlasmaInjector.H
@@ -112,6 +112,8 @@ public:
     long npart;
     int do_symmetrize = 0;
     int symmetrization_order = 4;
+    bool do_focusing = false;
+    amrex::Real focal_distance;
 
     bool external_file = false; //! initialize from an openPMD file
     amrex::Real z_shift = 0.0; //! additional z offset for particle positions

diff --git a/Source/Initialization/PlasmaInjector.cpp b/Source/Initialization/PlasmaInjector.cpp
@@ -232,6 +232,11 @@ void PlasmaInjector::setupGaussianBeam (amrex::ParmParse const& pp_species)
     utils::parser::getWithParser(pp_species, source_name, "npart", npart);
     utils::parser::queryWithParser(pp_species, source_name, "do_symmetrize", do_symmetrize);
     utils::parser::queryWithParser(pp_species, source_name, "symmetrization_order", symmetrization_order);
+    const bool focusing_is_specified = pp_species.contains("focal_distance");
+    if(focusing_is_specified){
+        do_focusing = true;
+        utils::parser::queryWithParser(pp_species, source_name, "focal_distance", focal_distance);
+    }
     const std::set<int> valid_symmetries = {4,8};
     WARPX_ALWAYS_ASSERT_WITH_MESSAGE( valid_symmetries.count(symmetrization_order),
         "Error: Symmetrization only supported to orders 4 or 8 ");
@@ -240,6 +245,7 @@ void PlasmaInjector::setupGaussianBeam (amrex::ParmParse const& pp_species)
                                 ux_parser, uy_parser, uz_parser,
                                 ux_th_parser, uy_th_parser, uz_th_parser,
                                 h_mom_temp, h_mom_vel);
+
 #if defined(WARPX_DIM_XZ)
     WARPX_ALWAYS_ASSERT_WITH_MESSAGE( y_rms > 0._rt,
         "Error: Gaussian beam y_rms must be strictly greater than 0 in 2D "

diff --git a/Source/Particles/PhysicalParticleContainer.H b/Source/Particles/PhysicalParticleContainer.H
@@ -231,7 +231,8 @@ public:
         amrex::Real x_m, amrex::Real y_m, amrex::Real z_m,
         amrex::Real x_rms, amrex::Real y_rms, amrex::Real z_rms,
         amrex::Real x_cut, amrex::Real y_cut, amrex::Real z_cut,
-        amrex::Real q_tot, long npart, int do_symmetrize, int symmetrization_order);
+        amrex::Real q_tot, long npart, int do_symmetrize, int symmetrization_order,
+        amrex::Real focal_distance);
 
     /** Load a particle beam from an external file
      * @param[in] the PlasmaInjector instance holding the input parameters

diff --git a/Source/Particles/PhysicalParticleContainer.cpp b/Source/Particles/PhysicalParticleContainer.cpp
@@ -528,7 +528,7 @@ PhysicalParticleContainer::AddGaussianBeam (
     const Real x_cut, const Real y_cut, const Real z_cut,
     const Real q_tot, long npart,
     const int do_symmetrize,
-    const int symmetrization_order) {
+    const int symmetrization_order, const Real focal_distance) {
 
     // Declare temporary vectors on the CPU
     Gpu::HostVector<ParticleReal> particle_x;
@@ -549,28 +549,65 @@ PhysicalParticleContainer::AddGaussianBeam (
         for (long i = 0; i < npart; ++i) {
 #if defined(WARPX_DIM_3D) || defined(WARPX_DIM_RZ)
             const Real weight = q_tot/(npart*charge);
-            const Real x = amrex::RandomNormal(x_m, x_rms);
-            const Real y = amrex::RandomNormal(y_m, y_rms);
-            const Real z = amrex::RandomNormal(z_m, z_rms);
+            Real x = amrex::RandomNormal(x_m, x_rms);
+            Real y = amrex::RandomNormal(y_m, y_rms);
+            Real z = amrex::RandomNormal(z_m, z_rms);
 #elif defined(WARPX_DIM_XZ)
             const Real weight = q_tot/(npart*charge*y_rms);
-            const Real x = amrex::RandomNormal(x_m, x_rms);
+            Real x = amrex::RandomNormal(x_m, x_rms);
             constexpr Real y = 0._prt;
-            const Real z = amrex::RandomNormal(z_m, z_rms);
+            Real z = amrex::RandomNormal(z_m, z_rms);
 #elif defined(WARPX_DIM_1D_Z)
             const Real weight = q_tot/(npart*charge*x_rms*y_rms);
             constexpr Real x = 0._prt;
             constexpr Real y = 0._prt;
-            const Real z = amrex::RandomNormal(z_m, z_rms);
+            Real z = amrex::RandomNormal(z_m, z_rms);
 #endif
             if (plasma_injector.insideBounds(x, y, z)  &&
                 std::abs( x - x_m ) <= x_cut * x_rms     &&
                 std::abs( y - y_m ) <= y_cut * y_rms     &&
                 std::abs( z - z_m ) <= z_cut * z_rms   ) {
                 XDim3 u = plasma_injector.getMomentum(x, y, z);
+
+            if (plasma_injector.do_focusing){
+                XDim3 u_bulk = plasma_injector.getInjectorMomentumHost()->getBulkMomentum(x,y,z);
+                Real u_bulk_norm = std::sqrt( u_bulk.x*u_bulk.x+u_bulk.y*u_bulk.y+u_bulk.z*u_bulk.z );
+
+                // Compute the position of the focal plane
+                // (it is located at a distance `focal_distance` from the beam centroid, in the direction of the bulk velocity)
+                Real n_x = u_bulk.x/u_bulk_norm;
+                Real n_y = u_bulk.y/u_bulk_norm;
+                Real n_z = u_bulk.z/u_bulk_norm;
+                Real x_f = x_m + focal_distance * n_x;
+                Real y_f = y_m + focal_distance * n_y;
+                Real z_f = z_m + focal_distance * n_z;
+                Real gamma = std::sqrt( 1._rt + (u.x*u.x+u.y*u.y+u.z*u.z) );
+
+                Real v_x = u.x / gamma * PhysConst::c;
+                Real v_y = u.y / gamma * PhysConst::c;
+                Real v_z = u.z / gamma * PhysConst::c;
+
+                // Compute the time at which the particle will cross the focal plane
+                Real v_dot_n = v_x * n_x + v_y * n_y + v_z * n_z;
+                Real t = ((x_f-x)*n_x + (y_f-y)*n_y + (z_f-z)*n_z) / v_dot_n;
+
+                // Displace particles in the direction orthogonal to the beam bulk momentum
+                // i.e. orthogonal to (n_x, n_y, n_z)
+#if defined(WARPX_DIM_3D) || defined(WARPX_DIM_RZ)
+                x = x - (v_x - v_dot_n*n_x) * t;
+                y = y - (v_y - v_dot_n*n_y) * t;
+                z = z - (v_z - v_dot_n*n_z) * t;
+#elif defined(WARPX_DIM_XZ)
+                x = x - (v_x - v_dot_n*n_x) * t;
+                z = z - (v_z - v_dot_n*n_z) * t;
+#elif defined(WARPX_DIM_1D_Z)
+                z = z - (v_z - v_dot_n*n_z) * t;
+#endif
+            }
                 u.x *= PhysConst::c;
                 u.y *= PhysConst::c;
                 u.z *= PhysConst::c;
+
                 if (do_symmetrize && symmetrization_order == 8){
                     // Add eight particles to the beam:
                     CheckAndAddParticle(x, y, z, u.x, u.y, u.z, weight/8._rt,
@@ -634,6 +671,7 @@ PhysicalParticleContainer::AddGaussianBeam (
     }
     // Add the temporary CPU vectors to the particle structure
     auto const np = static_cast<long>(particle_z.size());
+
     amrex::Vector<ParticleReal> xp(particle_x.data(), particle_x.data() + np);
     amrex::Vector<ParticleReal> yp(particle_y.data(), particle_y.data() + np);
     amrex::Vector<ParticleReal> zp(particle_z.data(), particle_z.data() + np);
@@ -893,7 +931,8 @@ PhysicalParticleContainer::AddParticles (int lev)
                             plasma_injector->q_tot,
                             plasma_injector->npart,
                             plasma_injector->do_symmetrize,
-                            plasma_injector->symmetrization_order);
+                            plasma_injector->symmetrization_order,
+                            plasma_injector->focal_distance);
         }
 
         if (plasma_injector->external_file) {