Add rotation to Gaussian beam initialization

Docs/source/usage/parameters.rst

@@ -951,6 +951,17 @@ Particle initialization
 
           \sigma_{x,y}(z) &= \sigma^*_{x,y} \sqrt{1 + \left( \frac{z - z^*}{\beta^*_{x,y}} \right)^2}
 
+      * ``<species_name>.do_rotation`` (`bool`, optional) the beam is rotated around its centroid.
+
+      If ``do_rotation = 1`` then the user needs to specify:
+
+          * ``<species_name>.rotation_axis``: (list of 3 `doubles`) axis around which the rotation takes place
+
+          * ``<species_name>.rotation_angle``: (`double`) angle of rotation around the specified axis, in radians.
+
+
+      Note that the other beam parameters (e.g. ``<species_name>.x/y/z_rms``, etc.) are used in the initialization process `before` performing the rotation.
+      Therefore, the user should define the beam size, cuts, and focal distance for the beam pre-rotatation, hence aligned to the Cartesian axes.
 
     * ``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:

Examples/Tests/gaussian_beam/CMakeLists.txt

@@ -20,3 +20,13 @@ add_warpx_test(
     diags/diag1000010  # output
     OFF  # dependency
 )
+
+add_warpx_test(
+    test_3d_rotated_gaussian_beam  # name
+    3  # dims
+    2  # nprocs
+    inputs_test_3d_rotated_gaussian_beam  # inputs
+    analysis_rotated_beam.py  # analysis
+    diags/diag1000000  # output
+    OFF  # dependency
+)

Examples/Tests/gaussian_beam/analysis_rotated_beam.py

@@ -0,0 +1,96 @@
+#!/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 openpmd_viewer import OpenPMDTimeSeries
+from scipy.constants import c, eV, m_e, micro, milli
+
+sys.path.insert(1, "../../../../warpx/Regression/Checksum/")
+import checksumAPI
+
+sigmax = 2 * micro
+sigmay = 1 * micro
+sigmaz = 3.0 * micro
+emitx = 100 * milli
+emity = 20 * milli
+gamma = 125 * 1e9 * eV / (m_e * c**2)
+x_m = 10 * sigmax
+y_m = 10 * sigmay
+z_m = 10 * sigmaz
+focal_distance = 4 * sigmaz
+theta = 0.25 * np.pi
+
+
+def s(z, sigma, 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(sigma**2 + emit**2 * (z - z_m - focal_distance) ** 2 / (sigma**2))
+
+
+filename = sys.argv[1]
+
+series = OpenPMDTimeSeries("./diags/openpmd/")
+
+# Rotated beam
+xrot, yrot, zrot, w, uxrot, uyrot, uzrot = series.get_particle(
+    ["x", "y", "z", "w", "ux", "uy", "uz"], species="beam1", iteration=0, plot=False
+)
+
+# Rotate the beam back so that it propagates along z
+x = x_m + np.cos(-theta) * (xrot - x_m) + np.sin(-theta) * (zrot - z_m)
+y = yrot
+z = z_m - np.sin(-theta) * (xrot - x_m) + np.cos(-theta) * (zrot - z_m)
+
+
+# Compute the size of the beam in different z slices
+gridz = np.linspace(z_m - focal_distance * 0.8, z_m + focal_distance * 0.8, 64)
+tol = 0.5 * (gridz[1] - gridz[0])
+sx, sy = [], []
+for z_g in gridz:
+    i = np.abs(z - z_g) < 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])))
+    else:
+        sx.append(0.0)
+        sy.append(0.0)
+
+# Theoretical prediction for the size of the beam in each z slice
+sx_theory = s(gridz, sigmax, emitx / gamma)
+sy_theory = s(gridz, sigmay, emity / gamma)
+
+assert np.allclose(sx, sx_theory, rtol=0.086, atol=0)
+assert np.allclose(sy, sy_theory, rtol=0.056, atol=0)
+
+# Rotate the beam back so that it propagates along z
+ux = np.cos(-theta) * uxrot + np.sin(-theta) * uzrot
+uy = uyrot
+uz = -np.sin(-theta) * uxrot + np.cos(-theta) * uzrot
+
+
+uz_m = np.average(uz, weights=w)
+uz_th = np.sqrt(np.average((uz - uz_m) ** 2, weights=w))
+
+ux_m = np.average(ux, weights=w)
+ux_th = np.sqrt(np.average((ux - ux_m) ** 2, weights=w))
+
+
+assert np.isclose(uz_m, gamma, rtol=1e-8, atol=0)
+assert np.abs(uz_th) < 1e-8
+assert np.abs(ux_m) < 4
+assert np.isclose(ux_th, emitx / sigmax, rtol=1e-3, atol=0)
+
+
+test_name = os.path.split(os.getcwd())[1]
+checksumAPI.evaluate_checksum(test_name, filename)

Examples/Tests/gaussian_beam/inputs_test_3d_rotated_gaussian_beam

@@ -0,0 +1,118 @@
+#################################
+########## MY CONSTANTS #########
+#################################
+my_constants.mc2 = m_e*clight*clight
+my_constants.micro = 1.e-6
+my_constants.milli = 1.e-3
+my_constants.GeV = q_e*1.e9
+
+# BEAM
+my_constants.energy = 125.*GeV
+my_constants.gamma = energy/mc2
+my_constants.npart = 2.e10
+my_constants.nmacropart = 2000000
+my_constants.charge = q_e * npart
+my_constants.sigmax = 2*micro
+my_constants.sigmay = 1*micro
+my_constants.sigmaz = 3.*micro
+my_constants.mux = 0.5*Lx
+my_constants.muy = 0.5*Ly
+my_constants.muz = 0.5*Lz
+my_constants.ux = 0.0
+my_constants.uy = 0.0
+my_constants.uz = gamma
+my_constants.emitx = 100*milli
+my_constants.emity = 20*milli
+my_constants.dux = emitx / sigmax
+my_constants.duy = emity / sigmay
+my_constants.duz = 0.
+my_constants.focal_distance = 4*sigmaz
+my_constants.nmacropart = 1e6
+my_constants.rotation_angle = pi*0.25
+
+# BOX
+my_constants.Lx = 20*sigmax
+my_constants.Ly = 20*sigmay
+my_constants.Lz = 20*sigmaz
+my_constants.nx = 128
+my_constants.ny = 128
+my_constants.nz = 128
+
+#################################
+####### 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 0 0
+geometry.prob_hi = Lx  Ly  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 ###########
+#################################
+warpx.do_electrostatic = relativistic
+warpx.const_dt = sigmaz/clight/20
+warpx.grid_type = collocated
+algo.particle_shape = 3
+algo.particle_pusher = vay
+
+#################################
+########### 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 = mux
+beam1.y_m = muy
+beam1.z_m = muz
+beam1.focal_distance = focal_distance
+beam1.npart = nmacropart
+beam1.q_tot = -charge
+beam1.do_rotation = 1
+beam1.rotation_angle = rotation_angle
+beam1.rotation_axis = 0 1 0
+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 = openpmd diag1
+
+diag1.intervals = 0
+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 ux uy uz
+openpmd.fields_to_plot = none
+openpmd.format = openpmd
+openpmd.dump_last_timestep = 1

Regression/Checksum/benchmarks_json/test_3d_rotated_gaussian_beam.json

@@ -0,0 +1,14 @@
+{
+  "lev=0": {
+    "rho_beam1": 140000330514.43195
+  },
+  "beam1": {
+    "particle_momentum_x": 4.723661353801595e-11,
+    "particle_momentum_y": 4.353803030332835e-12,
+    "particle_momentum_z": 4.7237908112561805e-11,
+    "particle_position_x": 19.99700982207277,
+    "particle_position_y": 9.99844596216089,
+    "particle_position_z": 29.99550729505918,
+    "particle_weight": 20000000000.0
+  }
+}

Source/Initialization/PlasmaInjector.H

@@ -114,6 +114,9 @@ public:
     int symmetrization_order = 4;
     bool do_focusing = false;
     amrex::Real focal_distance;
+    bool do_rotation = false;
+    amrex::Real rotation_angle;
+    amrex::Vector<amrex::Real> rotation_axis;
 
     bool external_file = false; //! initialize from an openPMD file
     amrex::Real z_shift = 0.0; //! additional z offset for particle positions

Source/Initialization/PlasmaInjector.cpp

@@ -233,6 +233,13 @@ void PlasmaInjector::setupGaussianBeam (amrex::ParmParse const& pp_species)
     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");
+    utils::parser::queryWithParser(pp_species, source_name, "do_rotation", do_rotation);
+
+    if(do_rotation){
+        utils::parser::queryWithParser(pp_species, source_name, "rotation_angle", rotation_angle);
+        utils::parser::getArrWithParser(pp_species, source_name, "rotation_axis", rotation_axis, 0, 3);
+    }
+
     if(focusing_is_specified){
         do_focusing = true;
         utils::parser::queryWithParser(pp_species, source_name, "focal_distance", focal_distance);