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/beamoptic.H"
16 #include "mixin/thick.H"
17 
18 #include <ablastr/constant.H>
19 
20 #include <AMReX.H>
21 #include <AMReX_Array.H>
22 #include <AMReX_Extension.H>
23 #include <AMReX_REAL.H>
24 
25 #include <array>
26 #include <cmath>
27 #include <stdexcept>
28 #include <tuple>
29 #include <vector>
30 
31 
32 namespace impactx
33 {
53  {
55  0.834166514794446,
56  0.598104328994702,
57  0.141852844428785,
58  -0.118211272182381,
59  -9.056149864743113E-002,
60  1.803476331179615E-002,
61  4.464887700797893E-002,
62  7.364410636252136E-003,
63  -1.697541023436736E-002,
64  -9.012679515542771E-003,
65  4.367667630047725E-003,
66  5.444030542119803E-003,
67  -5.889959910931886E-005,
68  -2.409098101409192E-003,
69  -7.962712154165590E-004,
70  7.855814707106538E-004,
71  6.174930463182168E-004,
72  -1.340154094301854E-004,
73  -3.167213724698439E-004,
74  -4.925292460592617E-005,
75  1.221580597451921E-004,
76  6.331025910961789E-005,
77  -3.202416719002774E-005,
78  -3.872103053895529E-005,
79  8.212882937116278E-007
80  };
81 
83  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
84  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
85  0, 0, 0, 0, 0
86  };
87  };
88 
95 namespace SoftQuadrupoleData
96 {
98  static inline int next_id = 0;
99 
101  static inline std::map<int, std::vector<amrex::ParticleReal>> h_cos_coef = {};
103  static inline std::map<int, std::vector<amrex::ParticleReal>> h_sin_coef = {};
104 
106  static inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_cos_coef = {};
108  static inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_sin_coef = {};
109 
110 } // namespace SoftQuadrupoleData
111 
113  : public elements::BeamOptic<SoftQuadrupole>,
114  public elements::Thick
115  {
116  static constexpr auto name = "SoftQuadrupole";
118 
130  amrex::ParticleReal ds,
131  amrex::ParticleReal gscale,
132  std::vector<amrex::ParticleReal> cos_coef,
133  std::vector<amrex::ParticleReal> sin_coef,
134  int mapsteps = 1,
135  int nslice = 1
136  )
137  : Thick(ds, nslice),
138  m_gscale(gscale), m_mapsteps(mapsteps), m_id(SoftQuadrupoleData::next_id)
139  {
140  // next created soft quad has another id for its data
142 
143  // validate sin and cos coefficients are the same length
144  m_ncoef = cos_coef.size();
145  if (m_ncoef != int(sin_coef.size()))
146  throw std::runtime_error("SoftQuadrupole: cos and sin coefficients must have same length!");
147 
148  // host data
153 
154  // device data
158  cos_coef.begin(), cos_coef.end(),
161  sin_coef.begin(), sin_coef.end(),
164 
165  // low-level objects we can use on device
168  }
169 
171  using BeamOptic::operator();
172 
184  PType& AMREX_RESTRICT p,
185  amrex::ParticleReal & AMREX_RESTRICT px,
186  amrex::ParticleReal & AMREX_RESTRICT py,
187  amrex::ParticleReal & AMREX_RESTRICT pt,
188  [[maybe_unused]] RefPart const & refpart
189  ) const
190  {
191  using namespace amrex::literals; // for _rt and _prt
192 
193  // access AoS data such as positions and cpu/id
194  amrex::ParticleReal const x = p.pos(RealAoS::x);
195  amrex::ParticleReal const y = p.pos(RealAoS::y);
196  amrex::ParticleReal const t = p.pos(RealAoS::t);
197 
198  // initialize output values of momenta
199  amrex::ParticleReal pxout = px;
200  amrex::ParticleReal pyout = py;
201  amrex::ParticleReal ptout = pt;
202 
203  // get the linear map
205 
206  // symplectic linear map for a quadrupole is computed using the
207  // Hamiltonian formalism as described in:
208  // https://uspas.fnal.gov/materials/09UNM/ComputationalMethods.pdf .
209  // R denotes the transfer matrix in the basis (x,px,y,py,t,pt),
210  // so that, e.g., R(3,4) = dyf/dpyi.
211 
212  // push particles using the linear map
213  p.pos(RealAoS::x) = R(1,1)*x + R(1,2)*px + R(1,3)*y
214  + R(1,4)*py + R(1,5)*t + R(1,6)*pt;
215  pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y
216  + R(2,4)*py + R(2,5)*t + R(2,6)*pt;
217  p.pos(RealAoS::y) = R(3,1)*x + R(3,2)*px + R(3,3)*y
218  + R(3,4)*py + R(3,5)*t + R(3,6)*pt;
219  pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y
220  + R(4,4)*py + R(4,5)*t + R(4,6)*pt;
221  p.pos(RealAoS::t) = R(5,1)*x + R(5,2)*px + R(5,3)*y
222  + R(5,4)*py + R(5,5)*t + R(5,6)*pt;
223  ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y
224  + R(6,4)*py + R(6,5)*t + R(6,6)*pt;
225 
226  // assign updated momenta
227  px = pxout;
228  py = pyout;
229  pt = ptout;
230  }
231 
237  void operator() (RefPart & AMREX_RESTRICT refpart) const
238  {
239  using namespace amrex::literals; // for _rt and _prt
240 
241  // assign input reference particle values
242  amrex::ParticleReal const x = refpart.x;
243  amrex::ParticleReal const px = refpart.px;
244  amrex::ParticleReal const y = refpart.y;
245  amrex::ParticleReal const py = refpart.py;
246  amrex::ParticleReal const z = refpart.z;
247  amrex::ParticleReal const pz = refpart.pz;
248  amrex::ParticleReal const pt = refpart.pt;
249  amrex::ParticleReal const s = refpart.s;
250  amrex::ParticleReal const sedge = refpart.sedge;
251 
252  // initialize linear map (deviation) values
253  for (int i=1; i<7; i++) {
254  for (int j=1; j<7; j++) {
255  auto const default_value = (i == j) ? 1.0_prt : 0.0_prt;
256  refpart.map(i, j) = default_value;
257  }
258  }
259 
260  // length of the current slice
261  amrex::ParticleReal const slice_ds = m_ds / nslice();
262 
263  // compute intial value of beta*gamma
264  amrex::ParticleReal const bgi = sqrt(pow(pt, 2) - 1.0_prt);
265 
266  // call integrator to advance (t,pt)
267  amrex::ParticleReal const zin = s - sedge;
268  amrex::ParticleReal const zout = zin + slice_ds;
269  int const nsteps = m_mapsteps;
270 
271  integrators::symp2_integrate(refpart,zin,zout,nsteps,*this);
272  amrex::ParticleReal const ptf = refpart.pt;
273 
274  /*
275  // print computed linear map:
276  for(int i=1; i<7; ++i){
277  for(int j=1; j<7; ++j){
278  amrex::PrintToFile("QuadMap.txt") << i << " " <<
279  j << " " << refpart.map(i,j) << "\n";
280  }
281  }
282  //
283  */
284 
285  // advance position (x,y,z)
286  refpart.x = x + slice_ds*px/bgi;
287  refpart.y = y + slice_ds*py/bgi;
288  refpart.z = z + slice_ds*pz/bgi;
289 
290  // compute final value of beta*gamma
291  amrex::ParticleReal const bgf = sqrt(pow(ptf, 2) - 1.0_prt);
292 
293  // advance momentum (px,py,pz)
294  refpart.px = px*bgf/bgi;
295  refpart.py = py*bgf/bgi;
296  refpart.pz = pz*bgf/bgi;
297 
298  // advance integrated path length
299  refpart.s = s + slice_ds;
300  }
301 
308  std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal>
310  Quad_Bfield (amrex::ParticleReal const zeval) const
311  {
312  using namespace amrex::literals; // for _rt and _prt
313 
314  // pick the right data depending if we are on the host side
315  // (reference particle push) or device side (particles):
316 #if AMREX_DEVICE_COMPILE
317  amrex::ParticleReal* cos_data = m_cos_d_data;
318  amrex::ParticleReal* sin_data = m_sin_d_data;
319 #else
320  amrex::ParticleReal* cos_data = m_cos_h_data;
321  amrex::ParticleReal* sin_data = m_sin_h_data;
322 #endif
323 
324  // specify constants
326  amrex::ParticleReal const zlen = m_ds;
327  amrex::ParticleReal const zmid = zlen / 2.0_prt;
328 
329  // compute on-axis magnetic field (z is relative to quadrupole midpoint)
330  amrex::ParticleReal bfield = 0.0;
331  amrex::ParticleReal bfieldp = 0.0;
332  amrex::ParticleReal bfieldint = 0.0;
333  amrex::ParticleReal const z = zeval - zmid;
334 
335  if (std::abs(z) <= zmid)
336  {
337  bfield = 0.5_prt*cos_data[0];
338  bfieldint = z*bfield;
339  for (int j=1; j < m_ncoef; ++j)
340  {
341  bfield = bfield + cos_data[j] * cos(j * 2 * pi * z / zlen) +
342  sin_data[j] * sin(j * 2 * pi * z / zlen);
343  bfieldp = bfieldp - j * 2 * pi * cos_data[j] * sin(j * 2 * pi * z / zlen) / zlen +
344  j * 2 * pi * sin_data[j] * cos(j * 2 * pi * z / zlen) / zlen;
345  bfieldint = bfieldint + zlen * cos_data[j] * sin(j * 2 * pi * z / zlen) / (j * 2 * pi) -
346  zlen * sin_data[j] * cos(j * 2 * pi * z / zlen) / (j * 2 * pi);
347  }
348  }
349  return std::make_tuple(bfield, bfieldp, bfieldint);
350  }
351 
361  void map1 (amrex::ParticleReal const tau,
362  RefPart & refpart,
363  [[maybe_unused]] amrex::ParticleReal & zeval) const
364  {
365  using namespace amrex::literals; // for _rt and _prt
366 
367  // push the reference particle
368  amrex::ParticleReal const t = refpart.t;
369  amrex::ParticleReal const pt = refpart.pt;
370  amrex::ParticleReal const z = zeval;
371 
372  if (pt < -1.0_prt) {
373  refpart.t = t + tau/sqrt(1.0_prt - pow(pt, -2));
374  refpart.pt = pt;
375  }
376  else {
377  refpart.t = t;
378  refpart.pt = pt;
379  }
380 
381  zeval = z + tau;
382 
383  // push the linear map equations
385  amrex::ParticleReal const betgam = refpart.beta_gamma();
386 
387  refpart.map(1,1) = R(1,1) + tau*R(2,1);
388  refpart.map(1,2) = R(1,2) + tau*R(2,2);
389  refpart.map(1,3) = R(1,3) + tau*R(2,3);
390  refpart.map(1,4) = R(1,4) + tau*R(2,4);
391 
392  refpart.map(3,1) = R(3,1) + tau*R(4,1);
393  refpart.map(3,2) = R(3,2) + tau*R(4,2);
394  refpart.map(3,3) = R(3,3) + tau*R(4,3);
395  refpart.map(3,4) = R(3,4) + tau*R(4,4);
396 
397  refpart.map(5,5) = R(5,5) + tau*R(6,5)/pow(betgam,2);
398  refpart.map(5,6) = R(5,6) + tau*R(6,6)/pow(betgam,2);
399 
400  }
401 
411  void map2 (amrex::ParticleReal const tau,
412  RefPart & refpart,
413  amrex::ParticleReal & zeval) const
414  {
415  using namespace amrex::literals; // for _rt and _prt
416 
417  amrex::ParticleReal const t = refpart.t;
418  amrex::ParticleReal const pt = refpart.pt;
419 
420  // Define parameters and intermediate constants
421  amrex::ParticleReal const G0 = m_gscale;
422 
423  // push the reference particle
424  auto [bz, bzp, bzint] = Quad_Bfield(zeval);
425  amrex::ignore_unused(bzp, bzint);
426 
427  refpart.t = t;
428  refpart.pt = pt;
429 
430  // push the linear map equations
432  amrex::ParticleReal const alpha = G0*bz;
433 
434  refpart.map(2,1) = R(2,1) - tau*alpha*R(1,1);
435  refpart.map(2,2) = R(2,2) - tau*alpha*R(1,2);
436  refpart.map(2,3) = R(2,3) - tau*alpha*R(1,3);
437  refpart.map(2,4) = R(2,4) - tau*alpha*R(1,4);
438 
439  refpart.map(4,1) = R(4,1) + tau*alpha*R(3,1);
440  refpart.map(4,2) = R(4,2) + tau*alpha*R(3,2);
441  refpart.map(4,3) = R(4,3) + tau*alpha*R(3,3);
442  refpart.map(4,4) = R(4,4) + tau*alpha*R(3,4);
443 
444  }
445 
448  void
450  {
451  // remove from unique data map
456 
461  }
462 
463  private:
464  amrex::ParticleReal m_gscale;
466  int m_id;
467 
468  int m_ncoef = 0;
469  amrex::ParticleReal* m_cos_h_data = nullptr;
470  amrex::ParticleReal* m_sin_h_data = nullptr;
471  amrex::ParticleReal* m_cos_d_data = nullptr;
472  amrex::ParticleReal* m_sin_d_data = nullptr;
473  };
474 
475 } // namespace impactx
476 
477 #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
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:108
static int next_id
last used id for a created soft quad
Definition: SoftQuad.H:98
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:103
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:101
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:106
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:24
Definition: ImpactX.cpp:33
s
int nsteps
int count
Definition: SoftQuad.H:53
amrex::Vector< amrex::ParticleReal > default_cos_coef
Definition: SoftQuad.H:54
amrex::Vector< amrex::ParticleReal > default_sin_coef
Definition: SoftQuad.H:82
@ x
position in x [m] (at fixed s OR fixed t)
Definition: ImpactXParticleContainer.H:42
@ y
position in y [m] (at fixed s OR fixed t)
Definition: ImpactXParticleContainer.H:43
@ t
c * time-of-flight [m] (at fixed s)
Definition: ImpactXParticleContainer.H:44
Definition: ReferenceParticle.H:30
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::ParticleReal beta_gamma() const
Definition: ReferenceParticle.H:79
amrex::ParticleReal pt
energy deviation, 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:115
amrex::ParticleReal * m_sin_d_data
non-owning pointer to device cosine coefficients
Definition: SoftQuad.H:472
amrex::ParticleReal * m_sin_h_data
non-owning pointer to host cosine coefficients
Definition: SoftQuad.H:470
amrex::ParticleReal m_gscale
Definition: SoftQuad.H:464
int m_ncoef
unique soft quad id used for data lookup map
Definition: SoftQuad.H:468
SoftQuadrupole(amrex::ParticleReal ds, amrex::ParticleReal gscale, std::vector< amrex::ParticleReal > cos_coef, std::vector< amrex::ParticleReal > sin_coef, int mapsteps=1, int nslice=1)
Definition: SoftQuad.H:129
ImpactXParticleContainer::ParticleType PType
Definition: SoftQuad.H:117
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void map2(amrex::ParticleReal const tau, RefPart &refpart, amrex::ParticleReal &zeval) const
Definition: SoftQuad.H:411
amrex::ParticleReal * m_cos_d_data
non-owning pointer to host sine coefficients
Definition: SoftQuad.H:471
amrex::ParticleReal * m_cos_h_data
number of Fourier coefficients
Definition: SoftQuad.H:469
int m_id
number of map integration steps per slice
Definition: SoftQuad.H:466
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:310
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void operator()(PType &AMREX_RESTRICT p, amrex::ParticleReal &AMREX_RESTRICT px, amrex::ParticleReal &AMREX_RESTRICT py, amrex::ParticleReal &AMREX_RESTRICT pt, [[maybe_unused]] RefPart const &refpart) const
Definition: SoftQuad.H:183
static constexpr auto name
Definition: SoftQuad.H:116
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void map1(amrex::ParticleReal const tau, RefPart &refpart, [[maybe_unused]] amrex::ParticleReal &zeval) const
Definition: SoftQuad.H:361
void finalize()
Definition: SoftQuad.H:449
int m_mapsteps
scaling factor for quad field gradient
Definition: SoftQuad.H:465
Definition: beamoptic.H:135
Definition: thick.H:24
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::ParticleReal ds() const
Definition: thick.H:50
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE int nslice() const
Definition: thick.H:40
amrex::ParticleReal m_ds
Definition: thick.H:56
Thick(amrex::ParticleReal const ds, int const nslice)
Definition: thick.H:30