ImpactX
SoftQuad.H
Go to the documentation of this file.
1 /* Copyright 2022-2023 The Regents of the University of California, through Lawrence
2  * Berkeley National Laboratory (subject to receipt of any required
3  * approvals from the U.S. Dept. of Energy). All rights reserved.
4  *
5  * This file is part of ImpactX.
6  *
7  * Authors: Chad Mitchell, Axel Huebl
8  * License: BSD-3-Clause-LBNL
9  */
10 #ifndef IMPACTX_SOFTQUAD_H
11 #define IMPACTX_SOFTQUAD_H
12 
15 #include "mixin/alignment.H"
16 #include "mixin/beamoptic.H"
17 #include "mixin/thick.H"
18 
19 #include <ablastr/constant.H>
20 
21 #include <AMReX.H>
22 #include <AMReX_Array.H>
23 #include <AMReX_Extension.H>
24 #include <AMReX_REAL.H>
25 
26 #include <array>
27 #include <cmath>
28 #include <stdexcept>
29 #include <tuple>
30 #include <vector>
31 
32 
33 namespace impactx
34 {
54  {
56  0.834166514794446,
57  0.598104328994702,
58  0.141852844428785,
59  -0.118211272182381,
60  -9.056149864743113E-002,
61  1.803476331179615E-002,
62  4.464887700797893E-002,
63  7.364410636252136E-003,
64  -1.697541023436736E-002,
65  -9.012679515542771E-003,
66  4.367667630047725E-003,
67  5.444030542119803E-003,
68  -5.889959910931886E-005,
69  -2.409098101409192E-003,
70  -7.962712154165590E-004,
71  7.855814707106538E-004,
72  6.174930463182168E-004,
73  -1.340154094301854E-004,
74  -3.167213724698439E-004,
75  -4.925292460592617E-005,
76  1.221580597451921E-004,
77  6.331025910961789E-005,
78  -3.202416719002774E-005,
79  -3.872103053895529E-005,
80  8.212882937116278E-007
81  };
82 
84  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
85  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
86  0, 0, 0, 0, 0
87  };
88  };
89 
96 namespace SoftQuadrupoleData
97 {
99  static inline int next_id = 0;
100 
102  static inline std::map<int, std::vector<amrex::ParticleReal>> h_cos_coef = {};
104  static inline std::map<int, std::vector<amrex::ParticleReal>> h_sin_coef = {};
105 
107  static inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_cos_coef = {};
109  static inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_sin_coef = {};
110 
111 } // namespace SoftQuadrupoleData
112 
114  : public elements::BeamOptic<SoftQuadrupole>,
115  public elements::Thick,
116  public elements::Alignment
117  {
118  static constexpr auto name = "SoftQuadrupole";
120 
135  amrex::ParticleReal ds,
136  amrex::ParticleReal gscale,
137  std::vector<amrex::ParticleReal> cos_coef,
138  std::vector<amrex::ParticleReal> sin_coef,
139  amrex::ParticleReal dx = 0,
140  amrex::ParticleReal dy = 0,
141  amrex::ParticleReal rotation_degree = 0,
142  int mapsteps = 1,
143  int nslice = 1
144  )
145  : Thick(ds, nslice),
146  Alignment(dx, dy, rotation_degree),
147  m_gscale(gscale), m_mapsteps(mapsteps), m_id(SoftQuadrupoleData::next_id)
148  {
149  // next created soft quad has another id for its data
151 
152  // validate sin and cos coefficients are the same length
153  m_ncoef = int(cos_coef.size());
154  if (m_ncoef != int(sin_coef.size()))
155  throw std::runtime_error("SoftQuadrupole: cos and sin coefficients must have same length!");
156 
157  // host data
162 
163  // device data
167  cos_coef.begin(), cos_coef.end(),
170  sin_coef.begin(), sin_coef.end(),
173 
174  // low-level objects we can use on device
177  }
178 
180  using BeamOptic::operator();
181 
196  amrex::ParticleReal & AMREX_RESTRICT x,
197  amrex::ParticleReal & AMREX_RESTRICT y,
198  amrex::ParticleReal & AMREX_RESTRICT t,
199  amrex::ParticleReal & AMREX_RESTRICT px,
200  amrex::ParticleReal & AMREX_RESTRICT py,
201  amrex::ParticleReal & AMREX_RESTRICT pt,
202  [[maybe_unused]] uint64_t & AMREX_RESTRICT idcpu,
203  [[maybe_unused]] RefPart const & refpart
204  ) const
205  {
206  using namespace amrex::literals; // for _rt and _prt
207 
208  // shift due to alignment errors of the element
209  shift_in(x, y, px, py);
210 
211  // intialize output values
212  amrex::ParticleReal xout = x;
213  amrex::ParticleReal yout = y;
214  amrex::ParticleReal tout = t;
215 
216  // initialize output values of momenta
217  amrex::ParticleReal pxout = px;
218  amrex::ParticleReal pyout = py;
219  amrex::ParticleReal ptout = pt;
220 
221  // get the linear map
223 
224  // symplectic linear map for a quadrupole is computed using the
225  // Hamiltonian formalism as described in:
226  // https://uspas.fnal.gov/materials/09UNM/ComputationalMethods.pdf .
227  // R denotes the transfer matrix in the basis (x,px,y,py,t,pt),
228  // so that, e.g., R(3,4) = dyf/dpyi.
229 
230  // push particles using the linear map
231  xout = R(1,1)*x + R(1,2)*px + R(1,3)*y
232  + R(1,4)*py + R(1,5)*t + R(1,6)*pt;
233  pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y
234  + R(2,4)*py + R(2,5)*t + R(2,6)*pt;
235  yout = R(3,1)*x + R(3,2)*px + R(3,3)*y
236  + R(3,4)*py + R(3,5)*t + R(3,6)*pt;
237  pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y
238  + R(4,4)*py + R(4,5)*t + R(4,6)*pt;
239  tout = R(5,1)*x + R(5,2)*px + R(5,3)*y
240  + R(5,4)*py + R(5,5)*t + R(5,6)*pt;
241  ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y
242  + R(6,4)*py + R(6,5)*t + R(6,6)*pt;
243 
244  // assign updated values
245  x = xout;
246  y = yout;
247  t = tout;
248  px = pxout;
249  py = pyout;
250  pt = ptout;
251 
252  // undo shift due to alignment errors of the element
253  shift_out(x, y, px, py);
254  }
255 
261  void operator() (RefPart & AMREX_RESTRICT refpart) const
262  {
263  using namespace amrex::literals; // for _rt and _prt
264 
265  // assign input reference particle values
266  amrex::ParticleReal const x = refpart.x;
267  amrex::ParticleReal const px = refpart.px;
268  amrex::ParticleReal const y = refpart.y;
269  amrex::ParticleReal const py = refpart.py;
270  amrex::ParticleReal const z = refpart.z;
271  amrex::ParticleReal const pz = refpart.pz;
272  amrex::ParticleReal const pt = refpart.pt;
273  amrex::ParticleReal const s = refpart.s;
274  amrex::ParticleReal const sedge = refpart.sedge;
275 
276  // initialize linear map (deviation) values
277  for (int i=1; i<7; i++) {
278  for (int j=1; j<7; j++) {
279  auto const default_value = (i == j) ? 1.0_prt : 0.0_prt;
280  refpart.map(i, j) = default_value;
281  }
282  }
283 
284  // length of the current slice
285  amrex::ParticleReal const slice_ds = m_ds / nslice();
286 
287  // compute intial value of beta*gamma
288  amrex::ParticleReal const bgi = sqrt(pow(pt, 2) - 1.0_prt);
289 
290  // call integrator to advance (t,pt)
291  amrex::ParticleReal const zin = s - sedge;
292  amrex::ParticleReal const zout = zin + slice_ds;
293  int const nsteps = m_mapsteps;
294 
295  integrators::symp2_integrate(refpart,zin,zout,nsteps,*this);
296  amrex::ParticleReal const ptf = refpart.pt;
297 
298  /*
299  // print computed linear map:
300  for(int i=1; i<7; ++i){
301  for(int j=1; j<7; ++j){
302  amrex::PrintToFile("QuadMap.txt") << i << " " <<
303  j << " " << refpart.map(i,j) << "\n";
304  }
305  }
306  //
307  */
308 
309  // advance position (x,y,z)
310  refpart.x = x + slice_ds*px/bgi;
311  refpart.y = y + slice_ds*py/bgi;
312  refpart.z = z + slice_ds*pz/bgi;
313 
314  // compute final value of beta*gamma
315  amrex::ParticleReal const bgf = sqrt(pow(ptf, 2) - 1.0_prt);
316 
317  // advance momentum (px,py,pz)
318  refpart.px = px*bgf/bgi;
319  refpart.py = py*bgf/bgi;
320  refpart.pz = pz*bgf/bgi;
321 
322  // advance integrated path length
323  refpart.s = s + slice_ds;
324  }
325 
332  std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal>
334  Quad_Bfield (amrex::ParticleReal const zeval) const
335  {
336  using namespace amrex::literals; // for _rt and _prt
337 
338  // pick the right data depending if we are on the host side
339  // (reference particle push) or device side (particles):
340 #if AMREX_DEVICE_COMPILE
341  amrex::ParticleReal* cos_data = m_cos_d_data;
342  amrex::ParticleReal* sin_data = m_sin_d_data;
343 #else
344  amrex::ParticleReal* cos_data = m_cos_h_data;
345  amrex::ParticleReal* sin_data = m_sin_h_data;
346 #endif
347 
348  // specify constants
350  amrex::ParticleReal const zlen = m_ds;
351  amrex::ParticleReal const zmid = zlen / 2.0_prt;
352 
353  // compute on-axis magnetic field (z is relative to quadrupole midpoint)
354  amrex::ParticleReal bfield = 0.0;
355  amrex::ParticleReal bfieldp = 0.0;
356  amrex::ParticleReal bfieldint = 0.0;
357  amrex::ParticleReal const z = zeval - zmid;
358 
359  if (std::abs(z) <= zmid)
360  {
361  bfield = 0.5_prt*cos_data[0];
362  bfieldint = z*bfield;
363  for (int j=1; j < m_ncoef; ++j)
364  {
365  bfield = bfield + cos_data[j] * cos(j * 2 * pi * z / zlen) +
366  sin_data[j] * sin(j * 2 * pi * z / zlen);
367  bfieldp = bfieldp - j * 2 * pi * cos_data[j] * sin(j * 2 * pi * z / zlen) / zlen +
368  j * 2 * pi * sin_data[j] * cos(j * 2 * pi * z / zlen) / zlen;
369  bfieldint = bfieldint + zlen * cos_data[j] * sin(j * 2 * pi * z / zlen) / (j * 2 * pi) -
370  zlen * sin_data[j] * cos(j * 2 * pi * z / zlen) / (j * 2 * pi);
371  }
372  }
373  return std::make_tuple(bfield, bfieldp, bfieldint);
374  }
375 
385  void map1 (amrex::ParticleReal const tau,
386  RefPart & refpart,
387  [[maybe_unused]] amrex::ParticleReal & zeval) const
388  {
389  using namespace amrex::literals; // for _rt and _prt
390 
391  // push the reference particle
392  amrex::ParticleReal const t = refpart.t;
393  amrex::ParticleReal const pt = refpart.pt;
394  amrex::ParticleReal const z = zeval;
395 
396  if (pt < -1.0_prt) {
397  refpart.t = t + tau/sqrt(1.0_prt - pow(pt, -2));
398  refpart.pt = pt;
399  }
400  else {
401  refpart.t = t;
402  refpart.pt = pt;
403  }
404 
405  zeval = z + tau;
406 
407  // push the linear map equations
409  amrex::ParticleReal const betgam = refpart.beta_gamma();
410 
411  refpart.map(1,1) = R(1,1) + tau*R(2,1);
412  refpart.map(1,2) = R(1,2) + tau*R(2,2);
413  refpart.map(1,3) = R(1,3) + tau*R(2,3);
414  refpart.map(1,4) = R(1,4) + tau*R(2,4);
415 
416  refpart.map(3,1) = R(3,1) + tau*R(4,1);
417  refpart.map(3,2) = R(3,2) + tau*R(4,2);
418  refpart.map(3,3) = R(3,3) + tau*R(4,3);
419  refpart.map(3,4) = R(3,4) + tau*R(4,4);
420 
421  refpart.map(5,5) = R(5,5) + tau*R(6,5)/pow(betgam,2);
422  refpart.map(5,6) = R(5,6) + tau*R(6,6)/pow(betgam,2);
423 
424  }
425 
435  void map2 (amrex::ParticleReal const tau,
436  RefPart & refpart,
437  amrex::ParticleReal & zeval) const
438  {
439  using namespace amrex::literals; // for _rt and _prt
440 
441  amrex::ParticleReal const t = refpart.t;
442  amrex::ParticleReal const pt = refpart.pt;
443 
444  // Define parameters and intermediate constants
445  amrex::ParticleReal const G0 = m_gscale;
446 
447  // push the reference particle
448  auto [bz, bzp, bzint] = Quad_Bfield(zeval);
449  amrex::ignore_unused(bzp, bzint);
450 
451  refpart.t = t;
452  refpart.pt = pt;
453 
454  // push the linear map equations
456  amrex::ParticleReal const alpha = G0*bz;
457 
458  refpart.map(2,1) = R(2,1) - tau*alpha*R(1,1);
459  refpart.map(2,2) = R(2,2) - tau*alpha*R(1,2);
460  refpart.map(2,3) = R(2,3) - tau*alpha*R(1,3);
461  refpart.map(2,4) = R(2,4) - tau*alpha*R(1,4);
462 
463  refpart.map(4,1) = R(4,1) + tau*alpha*R(3,1);
464  refpart.map(4,2) = R(4,2) + tau*alpha*R(3,2);
465  refpart.map(4,3) = R(4,3) + tau*alpha*R(3,3);
466  refpart.map(4,4) = R(4,4) + tau*alpha*R(3,4);
467 
468  }
469 
472  void
474  {
475  // remove from unique data map
480 
485  }
486 
487  amrex::ParticleReal m_gscale;
489  int m_id;
490 
491  int m_ncoef = 0;
492  amrex::ParticleReal* m_cos_h_data = nullptr;
493  amrex::ParticleReal* m_sin_h_data = nullptr;
494  amrex::ParticleReal* m_cos_d_data = nullptr;
495  amrex::ParticleReal* m_sin_d_data = nullptr;
496  };
497 
498 } // namespace impactx
499 
500 #endif // IMPACTX_SOFTQUAD_H
#define AMREX_FORCE_INLINE
#define AMREX_GPU_HOST_DEVICE
#define AMREX_GPU_HOST
static constexpr amrex::Real pi
void copyAsync(HostToDevice, InIter begin, InIter end, OutIter result) noexcept
static constexpr HostToDevice hostToDevice
void streamSynchronize() noexcept
constexpr std::enable_if_t< std::is_floating_point< T >::value, T > pi()
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void ignore_unused(const Ts &...)
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE GpuComplex< T > pow(const GpuComplex< T > &a_z, const T &a_y) noexcept
const int[]
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE GpuComplex< T > sqrt(const GpuComplex< T > &a_z) noexcept
i
static std::map< int, amrex::Gpu::DeviceVector< amrex::ParticleReal > > d_sin_coef
device: sine coefficients in Fourier expansion of on-axis magnetic field Bz
Definition: SoftQuad.H:109
static int next_id
last used id for a created soft quad
Definition: SoftQuad.H:99
static std::map< int, std::vector< amrex::ParticleReal > > h_sin_coef
host: sine coefficients in Fourier expansion of on-axis magnetic field Bz
Definition: SoftQuad.H:104
static std::map< int, std::vector< amrex::ParticleReal > > h_cos_coef
host: cosine coefficients in Fourier expansion of on-axis magnetic field Bz
Definition: SoftQuad.H:102
static std::map< int, amrex::Gpu::DeviceVector< amrex::ParticleReal > > d_cos_coef
device: cosine coefficients in Fourier expansion of on-axis magnetic field Bz
Definition: SoftQuad.H:107
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void symp2_integrate(RefPart &refpart, amrex::ParticleReal const zin, amrex::ParticleReal const zout, int const nsteps, T_Element const &element)
Definition: Integrators.H:22
Definition: ImpactX.cpp:33
@ t
fixed t as the independent variable
s
int nsteps
int count
Definition: SoftQuad.H:54
amrex::Vector< amrex::ParticleReal > default_cos_coef
Definition: SoftQuad.H:55
amrex::Vector< amrex::ParticleReal > default_sin_coef
Definition: SoftQuad.H:83
Definition: ReferenceParticle.H:30
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::ParticleReal beta_gamma() const
Definition: ReferenceParticle.H:79
amrex::ParticleReal pt
energy, normalized by rest energy
Definition: ReferenceParticle.H:39
amrex::Array2D< amrex::ParticleReal, 1, 6, 1, 6 > map
linearized map
Definition: ReferenceParticle.H:44
amrex::ParticleReal t
clock time * c in meters
Definition: ReferenceParticle.H:35
Definition: SoftQuad.H:117
amrex::ParticleReal * m_sin_d_data
non-owning pointer to device cosine coefficients
Definition: SoftQuad.H:495
amrex::ParticleReal * m_sin_h_data
non-owning pointer to host cosine coefficients
Definition: SoftQuad.H:493
amrex::ParticleReal m_gscale
Definition: SoftQuad.H:487
int m_ncoef
unique soft quad id used for data lookup map
Definition: SoftQuad.H:491
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator()(amrex::ParticleReal &AMREX_RESTRICT x, amrex::ParticleReal &AMREX_RESTRICT y, amrex::ParticleReal &AMREX_RESTRICT t, amrex::ParticleReal &AMREX_RESTRICT px, amrex::ParticleReal &AMREX_RESTRICT py, amrex::ParticleReal &AMREX_RESTRICT pt, [[maybe_unused]] uint64_t &AMREX_RESTRICT idcpu, [[maybe_unused]] RefPart const &refpart) const
Definition: SoftQuad.H:195
ImpactXParticleContainer::ParticleType PType
Definition: SoftQuad.H:119
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void map2(amrex::ParticleReal const tau, RefPart &refpart, amrex::ParticleReal &zeval) const
Definition: SoftQuad.H:435
amrex::ParticleReal * m_cos_d_data
non-owning pointer to host sine coefficients
Definition: SoftQuad.H:494
amrex::ParticleReal * m_cos_h_data
number of Fourier coefficients
Definition: SoftQuad.H:492
SoftQuadrupole(amrex::ParticleReal ds, amrex::ParticleReal gscale, std::vector< amrex::ParticleReal > cos_coef, std::vector< amrex::ParticleReal > sin_coef, amrex::ParticleReal dx=0, amrex::ParticleReal dy=0, amrex::ParticleReal rotation_degree=0, int mapsteps=1, int nslice=1)
Definition: SoftQuad.H:134
int m_id
number of map integration steps per slice
Definition: SoftQuad.H:489
std::tuple< amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal > AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE Quad_Bfield(amrex::ParticleReal const zeval) const
Definition: SoftQuad.H:334
static constexpr auto name
Definition: SoftQuad.H:118
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void map1(amrex::ParticleReal const tau, RefPart &refpart, [[maybe_unused]] amrex::ParticleReal &zeval) const
Definition: SoftQuad.H:385
void finalize()
Definition: SoftQuad.H:473
int m_mapsteps
scaling factor for quad field gradient
Definition: SoftQuad.H:488
Definition: alignment.H:27
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void shift_out(amrex::ParticleReal &AMREX_RESTRICT x, amrex::ParticleReal &AMREX_RESTRICT y, amrex::ParticleReal &AMREX_RESTRICT px, amrex::ParticleReal &AMREX_RESTRICT py) const
Definition: alignment.H:91
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::ParticleReal dx() const
Definition: alignment.H:120
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void shift_in(amrex::ParticleReal &AMREX_RESTRICT x, amrex::ParticleReal &AMREX_RESTRICT y, amrex::ParticleReal &AMREX_RESTRICT px, amrex::ParticleReal &AMREX_RESTRICT py) const
Definition: alignment.H:61
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::ParticleReal dy() const
Definition: alignment.H:130
Definition: beamoptic.H:149
Definition: thick.H:24
Thick(amrex::ParticleReal ds, int nslice)
Definition: thick.H:30
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::ParticleReal ds() const
Definition: thick.H:53
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE int nslice() const
Definition: thick.H:43
amrex::ParticleReal m_ds
Definition: thick.H:58