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 
193  PType& AMREX_RESTRICT p,
194  amrex::ParticleReal & AMREX_RESTRICT px,
195  amrex::ParticleReal & AMREX_RESTRICT py,
196  amrex::ParticleReal & AMREX_RESTRICT pt,
197  [[maybe_unused]] RefPart const & refpart
198  ) const
199  {
200  using namespace amrex::literals; // for _rt and _prt
201 
202  // shift due to alignment errors of the element
203  shift_in(p.pos(RealAoS::x), p.pos(RealAoS::y), px, py);
204 
205  // access AoS data such as positions and cpu/id
206  amrex::ParticleReal const x = p.pos(RealAoS::x);
207  amrex::ParticleReal const y = p.pos(RealAoS::y);
208  amrex::ParticleReal const t = p.pos(RealAoS::t);
209 
210  // initialize output values of momenta
211  amrex::ParticleReal pxout = px;
212  amrex::ParticleReal pyout = py;
213  amrex::ParticleReal ptout = pt;
214 
215  // get the linear map
217 
218  // symplectic linear map for a quadrupole is computed using the
219  // Hamiltonian formalism as described in:
220  // https://uspas.fnal.gov/materials/09UNM/ComputationalMethods.pdf .
221  // R denotes the transfer matrix in the basis (x,px,y,py,t,pt),
222  // so that, e.g., R(3,4) = dyf/dpyi.
223 
224  // push particles using the linear map
225  p.pos(RealAoS::x) = R(1,1)*x + R(1,2)*px + R(1,3)*y
226  + R(1,4)*py + R(1,5)*t + R(1,6)*pt;
227  pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y
228  + R(2,4)*py + R(2,5)*t + R(2,6)*pt;
229  p.pos(RealAoS::y) = R(3,1)*x + R(3,2)*px + R(3,3)*y
230  + R(3,4)*py + R(3,5)*t + R(3,6)*pt;
231  pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y
232  + R(4,4)*py + R(4,5)*t + R(4,6)*pt;
233  p.pos(RealAoS::t) = R(5,1)*x + R(5,2)*px + R(5,3)*y
234  + R(5,4)*py + R(5,5)*t + R(5,6)*pt;
235  ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y
236  + R(6,4)*py + R(6,5)*t + R(6,6)*pt;
237 
238  // assign updated momenta
239  px = pxout;
240  py = pyout;
241  pt = ptout;
242 
243  // undo shift due to alignment errors of the element
244  shift_out(p.pos(RealAoS::x), p.pos(RealAoS::y), px, py);
245  }
246 
252  void operator() (RefPart & AMREX_RESTRICT refpart) const
253  {
254  using namespace amrex::literals; // for _rt and _prt
255 
256  // assign input reference particle values
257  amrex::ParticleReal const x = refpart.x;
258  amrex::ParticleReal const px = refpart.px;
259  amrex::ParticleReal const y = refpart.y;
260  amrex::ParticleReal const py = refpart.py;
261  amrex::ParticleReal const z = refpart.z;
262  amrex::ParticleReal const pz = refpart.pz;
263  amrex::ParticleReal const pt = refpart.pt;
264  amrex::ParticleReal const s = refpart.s;
265  amrex::ParticleReal const sedge = refpart.sedge;
266 
267  // initialize linear map (deviation) values
268  for (int i=1; i<7; i++) {
269  for (int j=1; j<7; j++) {
270  auto const default_value = (i == j) ? 1.0_prt : 0.0_prt;
271  refpart.map(i, j) = default_value;
272  }
273  }
274 
275  // length of the current slice
276  amrex::ParticleReal const slice_ds = m_ds / nslice();
277 
278  // compute intial value of beta*gamma
279  amrex::ParticleReal const bgi = sqrt(pow(pt, 2) - 1.0_prt);
280 
281  // call integrator to advance (t,pt)
282  amrex::ParticleReal const zin = s - sedge;
283  amrex::ParticleReal const zout = zin + slice_ds;
284  int const nsteps = m_mapsteps;
285 
286  integrators::symp2_integrate(refpart,zin,zout,nsteps,*this);
287  amrex::ParticleReal const ptf = refpart.pt;
288 
289  /*
290  // print computed linear map:
291  for(int i=1; i<7; ++i){
292  for(int j=1; j<7; ++j){
293  amrex::PrintToFile("QuadMap.txt") << i << " " <<
294  j << " " << refpart.map(i,j) << "\n";
295  }
296  }
297  //
298  */
299 
300  // advance position (x,y,z)
301  refpart.x = x + slice_ds*px/bgi;
302  refpart.y = y + slice_ds*py/bgi;
303  refpart.z = z + slice_ds*pz/bgi;
304 
305  // compute final value of beta*gamma
306  amrex::ParticleReal const bgf = sqrt(pow(ptf, 2) - 1.0_prt);
307 
308  // advance momentum (px,py,pz)
309  refpart.px = px*bgf/bgi;
310  refpart.py = py*bgf/bgi;
311  refpart.pz = pz*bgf/bgi;
312 
313  // advance integrated path length
314  refpart.s = s + slice_ds;
315  }
316 
323  std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal>
325  Quad_Bfield (amrex::ParticleReal const zeval) const
326  {
327  using namespace amrex::literals; // for _rt and _prt
328 
329  // pick the right data depending if we are on the host side
330  // (reference particle push) or device side (particles):
331 #if AMREX_DEVICE_COMPILE
332  amrex::ParticleReal* cos_data = m_cos_d_data;
333  amrex::ParticleReal* sin_data = m_sin_d_data;
334 #else
335  amrex::ParticleReal* cos_data = m_cos_h_data;
336  amrex::ParticleReal* sin_data = m_sin_h_data;
337 #endif
338 
339  // specify constants
341  amrex::ParticleReal const zlen = m_ds;
342  amrex::ParticleReal const zmid = zlen / 2.0_prt;
343 
344  // compute on-axis magnetic field (z is relative to quadrupole midpoint)
345  amrex::ParticleReal bfield = 0.0;
346  amrex::ParticleReal bfieldp = 0.0;
347  amrex::ParticleReal bfieldint = 0.0;
348  amrex::ParticleReal const z = zeval - zmid;
349 
350  if (std::abs(z) <= zmid)
351  {
352  bfield = 0.5_prt*cos_data[0];
353  bfieldint = z*bfield;
354  for (int j=1; j < m_ncoef; ++j)
355  {
356  bfield = bfield + cos_data[j] * cos(j * 2 * pi * z / zlen) +
357  sin_data[j] * sin(j * 2 * pi * z / zlen);
358  bfieldp = bfieldp - j * 2 * pi * cos_data[j] * sin(j * 2 * pi * z / zlen) / zlen +
359  j * 2 * pi * sin_data[j] * cos(j * 2 * pi * z / zlen) / zlen;
360  bfieldint = bfieldint + zlen * cos_data[j] * sin(j * 2 * pi * z / zlen) / (j * 2 * pi) -
361  zlen * sin_data[j] * cos(j * 2 * pi * z / zlen) / (j * 2 * pi);
362  }
363  }
364  return std::make_tuple(bfield, bfieldp, bfieldint);
365  }
366 
376  void map1 (amrex::ParticleReal const tau,
377  RefPart & refpart,
378  [[maybe_unused]] amrex::ParticleReal & zeval) const
379  {
380  using namespace amrex::literals; // for _rt and _prt
381 
382  // push the reference particle
383  amrex::ParticleReal const t = refpart.t;
384  amrex::ParticleReal const pt = refpart.pt;
385  amrex::ParticleReal const z = zeval;
386 
387  if (pt < -1.0_prt) {
388  refpart.t = t + tau/sqrt(1.0_prt - pow(pt, -2));
389  refpart.pt = pt;
390  }
391  else {
392  refpart.t = t;
393  refpart.pt = pt;
394  }
395 
396  zeval = z + tau;
397 
398  // push the linear map equations
400  amrex::ParticleReal const betgam = refpart.beta_gamma();
401 
402  refpart.map(1,1) = R(1,1) + tau*R(2,1);
403  refpart.map(1,2) = R(1,2) + tau*R(2,2);
404  refpart.map(1,3) = R(1,3) + tau*R(2,3);
405  refpart.map(1,4) = R(1,4) + tau*R(2,4);
406 
407  refpart.map(3,1) = R(3,1) + tau*R(4,1);
408  refpart.map(3,2) = R(3,2) + tau*R(4,2);
409  refpart.map(3,3) = R(3,3) + tau*R(4,3);
410  refpart.map(3,4) = R(3,4) + tau*R(4,4);
411 
412  refpart.map(5,5) = R(5,5) + tau*R(6,5)/pow(betgam,2);
413  refpart.map(5,6) = R(5,6) + tau*R(6,6)/pow(betgam,2);
414 
415  }
416 
426  void map2 (amrex::ParticleReal const tau,
427  RefPart & refpart,
428  amrex::ParticleReal & zeval) const
429  {
430  using namespace amrex::literals; // for _rt and _prt
431 
432  amrex::ParticleReal const t = refpart.t;
433  amrex::ParticleReal const pt = refpart.pt;
434 
435  // Define parameters and intermediate constants
436  amrex::ParticleReal const G0 = m_gscale;
437 
438  // push the reference particle
439  auto [bz, bzp, bzint] = Quad_Bfield(zeval);
440  amrex::ignore_unused(bzp, bzint);
441 
442  refpart.t = t;
443  refpart.pt = pt;
444 
445  // push the linear map equations
447  amrex::ParticleReal const alpha = G0*bz;
448 
449  refpart.map(2,1) = R(2,1) - tau*alpha*R(1,1);
450  refpart.map(2,2) = R(2,2) - tau*alpha*R(1,2);
451  refpart.map(2,3) = R(2,3) - tau*alpha*R(1,3);
452  refpart.map(2,4) = R(2,4) - tau*alpha*R(1,4);
453 
454  refpart.map(4,1) = R(4,1) + tau*alpha*R(3,1);
455  refpart.map(4,2) = R(4,2) + tau*alpha*R(3,2);
456  refpart.map(4,3) = R(4,3) + tau*alpha*R(3,3);
457  refpart.map(4,4) = R(4,4) + tau*alpha*R(3,4);
458 
459  }
460 
463  void
465  {
466  // remove from unique data map
471 
476  }
477 
478  private:
479  amrex::ParticleReal m_gscale;
481  int m_id;
482 
483  int m_ncoef = 0;
484  amrex::ParticleReal* m_cos_h_data = nullptr;
485  amrex::ParticleReal* m_sin_h_data = nullptr;
486  amrex::ParticleReal* m_cos_d_data = nullptr;
487  amrex::ParticleReal* m_sin_d_data = nullptr;
488  };
489 
490 } // namespace impactx
491 
492 #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:24
Definition: ImpactX.cpp:36
@ 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
@ x
position in x [m] (at fixed s OR fixed t)
Definition: ImpactXParticleContainer.H:48
@ y
position in y [m] (at fixed s OR fixed t)
Definition: ImpactXParticleContainer.H:49
@ t
c * time-of-flight [m] (at fixed s)
Definition: ImpactXParticleContainer.H:50
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:487
amrex::ParticleReal * m_sin_h_data
non-owning pointer to host cosine coefficients
Definition: SoftQuad.H:485
amrex::ParticleReal m_gscale
Definition: SoftQuad.H:479
int m_ncoef
unique soft quad id used for data lookup map
Definition: SoftQuad.H:483
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:426
amrex::ParticleReal * m_cos_d_data
non-owning pointer to host sine coefficients
Definition: SoftQuad.H:486
amrex::ParticleReal * m_cos_h_data
number of Fourier coefficients
Definition: SoftQuad.H:484
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:481
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:325
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:192
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:376
void finalize()
Definition: SoftQuad.H:464
int m_mapsteps
scaling factor for quad field gradient
Definition: SoftQuad.H:480
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:135
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