WarpX
LatticeElementFinder.H
Go to the documentation of this file.
1 /* Copyright 2022 David Grote
2  *
3  * This file is part of WarpX.
4  *
5  * License: BSD-3-Clause-LBNL
6  */
7 #ifndef WARPX_ACCELERATORLATTICE_LATTICEELEMENTS_LATTICEELEMENTFINDER_H_
8 #define WARPX_ACCELERATORLATTICE_LATTICEELEMENTS_LATTICEELEMENTFINDER_H_
9 
14 
15 #include <AMReX_REAL.H>
16 #include <AMReX_GpuContainers.H>
17 
18 class AcceleratorLattice;
20 
21 // Instances of the LatticeElementFinder class are saved in the AcceleratorLattice class
22 // as the objects in a LayoutData.
23 // The LatticeElementFinder handles the lookup needed to find the lattice elements at
24 // particle locations.
25 
27 {
28 
36  void InitElementFinder (int lev, amrex::MFIter const& a_mfi,
37  AcceleratorLattice const& accelerator_lattice);
38 
44  void AllocateIndices (AcceleratorLattice const& accelerator_lattice);
45 
53  void UpdateIndices (int lev, amrex::MFIter const& a_mfi,
54  AcceleratorLattice const& accelerator_lattice);
55 
56  /* Define the location and size of the index lookup table */
57  /* Use the type Real to be consistent with the way the main grid is defined */
58  int m_nz;
59  amrex::Real m_zmin;
60  amrex::Real m_dz;
61 
62  /* Parameters needed for the Lorentz transforms into and out of the boosted frame */
63  /* The time for m_time is consistent with the main time variable */
64  amrex::ParticleReal m_gamma_boost;
65  amrex::ParticleReal m_uz_boost;
66  amrex::Real m_time;
67 
76  WarpXParIter const& a_pti, int a_offset, AcceleratorLattice const& accelerator_lattice) const;
77 
78  /* The index lookup tables for each lattice element type */
81 
92  amrex::Gpu::DeviceVector<int> & indices) const;
93 };
94 
100 {
101 
110  void
111  InitLatticeElementFinderDevice (WarpXParIter const& a_pti, int a_offset,
112  AcceleratorLattice const& accelerator_lattice,
113  LatticeElementFinder const & h_finder);
114 
115  /* Size and location of the index lookup table */
116  amrex::Real m_zmin;
117  amrex::Real m_dz;
118  amrex::Real m_dt;
119 
120  /* Parameters needed for the Lorentz transforms into and out of the boosted frame */
121  amrex::ParticleReal m_gamma_boost;
122  amrex::ParticleReal m_uz_boost;
123  amrex::Real m_time;
124 
126  const amrex::ParticleReal* AMREX_RESTRICT m_ux = nullptr;
127  const amrex::ParticleReal* AMREX_RESTRICT m_uy = nullptr;
128  const amrex::ParticleReal* AMREX_RESTRICT m_uz = nullptr;
129 
130  /* Device level instances for each lattice element type */
133 
134  /* Device level index lookup tables for each element type */
135  int const* d_quad_indices_arr = nullptr;
136  int const* d_plasmalens_indices_arr = nullptr;
137 
145  void operator () (const long i,
146  amrex::ParticleReal& field_Ex,
147  amrex::ParticleReal& field_Ey,
148  amrex::ParticleReal& field_Ez,
149  amrex::ParticleReal& field_Bx,
150  amrex::ParticleReal& field_By,
151  amrex::ParticleReal& field_Bz) const noexcept
152  {
153 
154  using namespace amrex::literals;
155 
156  amrex::ParticleReal x, y, z;
157  m_get_position(i, x, y, z);
158 
159  // Find location of partice in the indices grid
160  // (which is in the boosted frame)
161  const int iz = static_cast<int>((z - m_zmin)/m_dz);
162 
163  constexpr amrex::ParticleReal inv_c2 = 1._prt/(PhysConst::c*PhysConst::c);
164  amrex::ParticleReal const gamma = std::sqrt(1._prt + (m_ux[i]*m_ux[i] + m_uy[i]*m_uy[i] + m_uz[i]*m_uz[i])*inv_c2);
165  amrex::ParticleReal const vzp = m_uz[i]/gamma;
166 
167  amrex::ParticleReal zpvdt = z + vzp*m_dt;
168 
169  // The position passed to the get_field methods needs to be in the lab frame.
170  if (m_gamma_boost > 1._prt) {
172  zpvdt = m_gamma_boost*zpvdt + m_uz_boost*(m_time + m_dt);
173  }
174 
175  amrex::ParticleReal Ex_sum = 0._prt;
176  amrex::ParticleReal Ey_sum = 0._prt;
177  const amrex::ParticleReal Ez_sum = 0._prt;
178  amrex::ParticleReal Bx_sum = 0._prt;
179  amrex::ParticleReal By_sum = 0._prt;
180  const amrex::ParticleReal Bz_sum = 0._prt;
181 
182  if (d_quad.nelements > 0) {
183  if (d_quad_indices_arr[iz] > -1) {
184  const auto ielement = d_quad_indices_arr[iz];
185  amrex::ParticleReal Ex, Ey, Bx, By;
186  d_quad.get_field(ielement, x, y, z, zpvdt, Ex, Ey, Bx, By);
187  Ex_sum += Ex;
188  Ey_sum += Ey;
189  Bx_sum += Bx;
190  By_sum += By;
191  }
192  }
193 
194  if (d_plasmalens.nelements > 0) {
195  if (d_plasmalens_indices_arr[iz] > -1) {
196  const auto ielement = d_plasmalens_indices_arr[iz];
197  amrex::ParticleReal Ex, Ey, Bx, By;
198  d_plasmalens.get_field(ielement, x, y, z, zpvdt, Ex, Ey, Bx, By);
199  Ex_sum += Ex;
200  Ey_sum += Ey;
201  Bx_sum += Bx;
202  By_sum += By;
203  }
204  }
205 
206  if (m_gamma_boost > 1._prt) {
207  // The fields returned from get_field is in the lab frame
208  // Transform the fields to the boosted frame
209  const amrex::ParticleReal Ex_boost = m_gamma_boost*Ex_sum - m_uz_boost*By_sum;
210  const amrex::ParticleReal Ey_boost = m_gamma_boost*Ey_sum + m_uz_boost*Bx_sum;
211  const amrex::ParticleReal Bx_boost = m_gamma_boost*Bx_sum + m_uz_boost*Ey_sum*inv_c2;
212  const amrex::ParticleReal By_boost = m_gamma_boost*By_sum - m_uz_boost*Ex_sum*inv_c2;
213  Ex_sum = Ex_boost;
214  Ey_sum = Ey_boost;
215  Bx_sum = Bx_boost;
216  By_sum = By_boost;
217  }
218 
219  field_Ex += Ex_sum;
220  field_Ey += Ey_sum;
221  field_Ez += Ez_sum;
222  field_Bx += Bx_sum;
223  field_By += By_sum;
224  field_Bz += Bz_sum;
225 
226  }
227 
228 };
229 
230 #endif // WARPX_ACCELERATORLATTICE_LATTICEELEMENTS_LATTICEELEMENTFINDER_H_
#define AMREX_FORCE_INLINE
#define AMREX_RESTRICT
#define AMREX_GPU_HOST_DEVICE
Definition: AcceleratorLattice.H:21
Definition: WarpXParticleContainer.H:53
static constexpr auto c
vacuum speed of light [m/s]
Definition: constant.H:44
constexpr int iz
i
Definition: check_interp_points_and_weights.py:174
int gamma
boosted frame
Definition: stencil.py:431
Definition: HardEdgedPlasmaLens.H:64
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void get_field(const int ielement, const amrex::ParticleReal x, const amrex::ParticleReal y, const amrex::ParticleReal z, const amrex::ParticleReal zpvdt, amrex::ParticleReal &Ex, amrex::ParticleReal &Ey, amrex::ParticleReal &Bx, amrex::ParticleReal &By) const
Fetch the field of the specified element at the given location.
Definition: HardEdgedPlasmaLens.H:90
int nelements
Definition: HardEdgedPlasmaLens.H:73
Definition: HardEdgedQuadrupole.H:64
int nelements
Definition: HardEdgedQuadrupole.H:73
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void get_field(const int ielement, const amrex::ParticleReal x, const amrex::ParticleReal y, const amrex::ParticleReal z, const amrex::ParticleReal zpvdt, amrex::ParticleReal &Ex, amrex::ParticleReal &Ey, amrex::ParticleReal &Bx, amrex::ParticleReal &By) const
Fetch the field of the specified element at the given location.
Definition: HardEdgedQuadrupole.H:90
The lattice element finder class that can be trivially copied to the device. This only has simple dat...
Definition: LatticeElementFinder.H:100
const amrex::ParticleReal *AMREX_RESTRICT m_ux
Definition: LatticeElementFinder.H:126
const amrex::ParticleReal *AMREX_RESTRICT m_uy
Definition: LatticeElementFinder.H:127
int const * d_plasmalens_indices_arr
Definition: LatticeElementFinder.H:136
amrex::Real m_time
Definition: LatticeElementFinder.H:123
HardEdgedPlasmaLensDevice d_plasmalens
Definition: LatticeElementFinder.H:132
amrex::Real m_dt
Definition: LatticeElementFinder.H:118
void InitLatticeElementFinderDevice(WarpXParIter const &a_pti, int a_offset, AcceleratorLattice const &accelerator_lattice, LatticeElementFinder const &h_finder)
Initialize the data needed to do the lookups.
Definition: LatticeElementFinder.cpp:89
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator()(const long i, amrex::ParticleReal &field_Ex, amrex::ParticleReal &field_Ey, amrex::ParticleReal &field_Ez, amrex::ParticleReal &field_Bx, amrex::ParticleReal &field_By, amrex::ParticleReal &field_Bz) const noexcept
Gather the field for the particle from the lattice elements.
Definition: LatticeElementFinder.H:145
amrex::Real m_dz
Definition: LatticeElementFinder.H:117
int const * d_quad_indices_arr
Definition: LatticeElementFinder.H:135
amrex::ParticleReal m_uz_boost
Definition: LatticeElementFinder.H:122
const amrex::ParticleReal *AMREX_RESTRICT m_uz
Definition: LatticeElementFinder.H:128
amrex::Real m_zmin
Definition: LatticeElementFinder.H:116
HardEdgedQuadrupoleDevice d_quad
Definition: LatticeElementFinder.H:131
GetParticlePosition< PIdx > m_get_position
Definition: LatticeElementFinder.H:125
amrex::ParticleReal m_gamma_boost
Definition: LatticeElementFinder.H:121
Definition: LatticeElementFinder.H:27
LatticeElementFinderDevice GetFinderDeviceInstance(WarpXParIter const &a_pti, int a_offset, AcceleratorLattice const &accelerator_lattice) const
Get the device level instance associated with this instance.
Definition: LatticeElementFinder.cpp:80
void setup_lattice_indices(amrex::Gpu::DeviceVector< amrex::ParticleReal > const &zs, amrex::Gpu::DeviceVector< amrex::ParticleReal > const &ze, amrex::Gpu::DeviceVector< int > &indices) const
Fill in the index lookup tables This loops over the grid (in z) and finds the lattice element closest...
Definition: LatticeElementFinder.cpp:125
amrex::Gpu::DeviceVector< int > d_quad_indices
Definition: LatticeElementFinder.H:79
int m_nz
Definition: LatticeElementFinder.H:58
amrex::Real m_zmin
Definition: LatticeElementFinder.H:59
amrex::Gpu::DeviceVector< int > d_plasmalens_indices
Definition: LatticeElementFinder.H:80
amrex::ParticleReal m_gamma_boost
Definition: LatticeElementFinder.H:64
amrex::Real m_time
Definition: LatticeElementFinder.H:66
void UpdateIndices(int lev, amrex::MFIter const &a_mfi, AcceleratorLattice const &accelerator_lattice)
Update the index lookup tables for each element type, filling in the values.
Definition: LatticeElementFinder.cpp:54
amrex::Real m_dz
Definition: LatticeElementFinder.H:60
amrex::ParticleReal m_uz_boost
Definition: LatticeElementFinder.H:65
void InitElementFinder(int lev, amrex::MFIter const &a_mfi, AcceleratorLattice const &accelerator_lattice)
Initialize the element finder at the level and grid.
Definition: LatticeElementFinder.cpp:18
void AllocateIndices(AcceleratorLattice const &accelerator_lattice)
Allocate the index lookup tables for each element type.
Definition: LatticeElementFinder.cpp:39