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 
89 namespace data
90 {
94  {
95  amrex::ParticleReal m_gscale;
96  int m_mapsteps;
97 
98  int m_ncoef = 0;
99  amrex::ParticleReal* m_cos_data = nullptr;
100  amrex::ParticleReal* m_sin_data = nullptr;
101 
103  amrex::ParticleReal gscale,
104  int mapsteps = 1
105  ) : m_gscale(gscale), m_mapsteps(mapsteps)
106  {}
107 
112 
113  };
114 } // namespace data
115 
117  : public elements::BeamOptic<SoftQuadrupole>,
118  public elements::Thick,
120  {
121  static constexpr auto name = "SoftQuadrupole";
123 
136  amrex::ParticleReal ds,
137  amrex::ParticleReal gscale,
138  std::vector<amrex::ParticleReal> cos_coef,
139  std::vector<amrex::ParticleReal> sin_coef,
140  int mapsteps = 1,
141  int nslice = 1
142  )
143  : Thick(ds, nslice),
144  SoftQuadrupole_device_copyable(gscale, mapsteps)
145  {
146  m_ncoef = cos_coef.size();
147  if (m_ncoef != int(sin_coef.size()))
148  throw std::runtime_error("SoftQuadrupole: cos and sin coefficients must have same length!");
149 
150  m_cos_coef.resize(m_ncoef);
151  m_sin_coef.resize(m_ncoef);
153  cos_coef.begin(), cos_coef.end(),
154  m_cos_coef.begin());
156  sin_coef.begin(), sin_coef.end(),
157  m_sin_coef.begin());
159 
160  // low-level objects we can use on device
161  m_cos_data = m_cos_coef.data();
162  m_sin_data = m_sin_coef.data();
163  }
164 
165  // copy and move constructors
167  : Thick(other.m_ds, other.m_nslice),
168  SoftQuadrupole_device_copyable(other.m_gscale, other.m_mapsteps)
169  {
170 #if !AMREX_DEVICE_COMPILE
171  // copy the data container if we copy the host element
172  m_cos_coef = other.m_cos_coef;
173  m_sin_coef = other.m_sin_coef;
175 #endif
176  m_ncoef = m_cos_coef.size();
177  m_cos_data = m_cos_coef.data();
178  m_sin_data = m_sin_coef.data();
179  }
181  {
182  if (this == &other)
183  return *this;
184 
185  Thick::operator=(other);
186  SoftQuadrupole_device_copyable::operator=(other);
187 #if !AMREX_DEVICE_COMPILE
188  // copy the data container if we copy the host element
189  m_cos_coef = other.m_cos_coef;
190  m_sin_coef = other.m_sin_coef;
192 #endif
193  m_ncoef = m_cos_coef.size();
194  m_cos_data = m_cos_coef.data();
195  m_sin_data = m_sin_coef.data();
196  return *this;
197  }
198  SoftQuadrupole(SoftQuadrupole && other) = default;
199  SoftQuadrupole& operator=(SoftQuadrupole && other) = default;
200 
202  ~SoftQuadrupole() = default;
203 
205  using BeamOptic::operator();
206 
217  void operator() (
218  PType& AMREX_RESTRICT p,
219  amrex::ParticleReal & AMREX_RESTRICT px,
220  amrex::ParticleReal & AMREX_RESTRICT py,
221  amrex::ParticleReal & AMREX_RESTRICT pt,
222  [[maybe_unused]] RefPart const & refpart
223  ) const
224  {
225  using namespace amrex::literals; // for _rt and _prt
226 
227  // access AoS data such as positions and cpu/id
228  amrex::ParticleReal const x = p.pos(0);
229  amrex::ParticleReal const y = p.pos(1);
230  amrex::ParticleReal const t = p.pos(2);
231 
232  // initialize output values of momenta
233  amrex::ParticleReal pxout = px;
234  amrex::ParticleReal pyout = py;
235  amrex::ParticleReal ptout = pt;
236 
237  // get the linear map
239 
240  // symplectic linear map for a quadrupole is computed using the
241  // Hamiltonian formalism as described in:
242  // https://uspas.fnal.gov/materials/09UNM/ComputationalMethods.pdf .
243  // R denotes the transfer matrix in the basis (x,px,y,py,t,pt),
244  // so that, e.g., R(3,4) = dyf/dpyi.
245 
246  // push particles using the linear map
247  p.pos(0) = R(1,1)*x + R(1,2)*px + R(1,3)*y
248  + R(1,4)*py + R(1,5)*t + R(1,6)*pt;
249  pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y
250  + R(2,4)*py + R(2,5)*t + R(2,6)*pt;
251  p.pos(1) = R(3,1)*x + R(3,2)*px + R(3,3)*y
252  + R(3,4)*py + R(3,5)*t + R(3,6)*pt;
253  pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y
254  + R(4,4)*py + R(4,5)*t + R(4,6)*pt;
255  p.pos(2) = R(5,1)*x + R(5,2)*px + R(5,3)*y
256  + R(5,4)*py + R(5,5)*t + R(5,6)*pt;
257  ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y
258  + R(6,4)*py + R(6,5)*t + R(6,6)*pt;
259 
260  // assign updated momenta
261  px = pxout;
262  py = pyout;
263  pt = ptout;
264  }
265 
271  void operator() (RefPart & AMREX_RESTRICT refpart) const
272  {
273  using namespace amrex::literals; // for _rt and _prt
274 
275  // assign input reference particle values
276  amrex::ParticleReal const x = refpart.x;
277  amrex::ParticleReal const px = refpart.px;
278  amrex::ParticleReal const y = refpart.y;
279  amrex::ParticleReal const py = refpart.py;
280  amrex::ParticleReal const z = refpart.z;
281  amrex::ParticleReal const pz = refpart.pz;
282  amrex::ParticleReal const pt = refpart.pt;
283  amrex::ParticleReal const s = refpart.s;
284  amrex::ParticleReal const sedge = refpart.sedge;
285 
286  // initialize linear map (deviation) values
287  for (int i=1; i<7; i++) {
288  for (int j=1; j<7; j++) {
289  auto const default_value = (i == j) ? 1.0_prt : 0.0_prt;
290  refpart.map(i, j) = default_value;
291  }
292  }
293 
294  // length of the current slice
295  amrex::ParticleReal const slice_ds = m_ds / nslice();
296 
297  // compute intial value of beta*gamma
298  amrex::ParticleReal const bgi = sqrt(pow(pt, 2) - 1.0_prt);
299 
300  // call integrator to advance (t,pt)
301  amrex::ParticleReal const zin = s - sedge;
302  amrex::ParticleReal const zout = zin + slice_ds;
303  int const nsteps = m_mapsteps;
304 
305  integrators::symp2_integrate(refpart,zin,zout,nsteps,*this);
306  amrex::ParticleReal const ptf = refpart.pt;
307 
308  /*
309  // print computed linear map:
310  for(int i=1; i<7; ++i){
311  for(int j=1; j<7; ++j){
312  amrex::PrintToFile("QuadMap.txt") << i << " " <<
313  j << " " << refpart.map(i,j) << "\n";
314  }
315  }
316  //
317  */
318 
319  // advance position (x,y,z)
320  refpart.x = x + slice_ds*px/bgi;
321  refpart.y = y + slice_ds*py/bgi;
322  refpart.z = z + slice_ds*pz/bgi;
323 
324  // compute final value of beta*gamma
325  amrex::ParticleReal const bgf = sqrt(pow(ptf, 2) - 1.0_prt);
326 
327  // advance momentum (px,py,pz)
328  refpart.px = px*bgf/bgi;
329  refpart.py = py*bgf/bgi;
330  refpart.pz = pz*bgf/bgi;
331 
332  // advance integrated path length
333  refpart.s = s + slice_ds;
334  }
335 
342  std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal>
344  Quad_Bfield (amrex::ParticleReal const zeval) const
345  {
346  using namespace amrex::literals; // for _rt and _prt
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 (abs(z)<=zmid)
360  {
361  bfield = 0.5_prt*m_cos_data[0];
362  bfieldint = z*bfield;
363  for (int j=1; j < m_ncoef; ++j)
364  {
365  bfield = bfield + m_cos_data[j]*cos(j*2*pi*z/zlen) +
366  m_sin_data[j]*sin(j*2*pi*z/zlen);
367  bfieldp = bfieldp-j*2*pi*m_cos_data[j]*sin(j*2*pi*z/zlen)/zlen +
368  j*2*pi*m_sin_data[j]*cos(j*2*pi*z/zlen)/zlen;
369  bfieldint = bfieldint + zlen*m_cos_data[j]*sin(j*2*pi*z/zlen)/(j*2*pi) -
370  zlen*m_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 
470 
471  private:
472  // we cannot copy these to device with a memcpy when we copy the element class
475  };
476 
477 } // namespace impactx
478 
479 #endif // IMPACTX_SOFTQUAD_H
amrex::Gpu::DeviceVector< amrex::ParticleReal > m_sin_coef
cosine coefficients in Fourier expansion of on-axis magnetic field Bz
Definition: SoftQuad.H:474
amrex::Vector< amrex::ParticleReal > default_cos_coef
Definition: SoftQuad.H:54
SoftQuadrupole & operator=(SoftQuadrupole const &other)
Definition: SoftQuad.H:180
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
amrex::Vector< amrex::ParticleReal > default_sin_coef
Definition: SoftQuad.H:82
constexpr std::enable_if_t< std::is_floating_point< T >::value, T > pi()
data
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
SoftQuadrupole_device_copyable(amrex::ParticleReal gscale, int mapsteps=1)
non-owning pointer to device sine coefficients
Definition: SoftQuad.H:102
Definition: ImpactX.cpp:31
static constexpr HostToDevice hostToDevice
Definition: beamoptic.H:134
int m_mapsteps
scaling factor for quad field gradient
Definition: SoftQuad.H:96
amrex::ParticleReal pt
energy deviation, normalized by rest energy
Definition: ReferenceParticle.H:39
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void ignore_unused(const Ts &...)
AMREX_GPU_HOST 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:135
#define AMREX_FORCE_INLINE
nslice
Definition: __init__.py:32
amrex::ParticleReal m_gscale
Definition: SoftQuad.H:95
#define AMREX_GPU_HOST_DEVICE
amrex::ParticleReal t
clock time * c in meters
Definition: ReferenceParticle.H:35
int nsteps
i
Definition: thick.H:23
amrex::Array2D< amrex::ParticleReal, 1, 6, 1, 6 > map
linearized map
Definition: ReferenceParticle.H:44
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE T abs(const GpuComplex< T > &a_z) noexcept
SoftQuadrupole(SoftQuadrupole const &other)
Definition: SoftQuad.H:166
Definition: ReferenceParticle.H:29
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 void map2(amrex::ParticleReal const tau, RefPart &refpart, amrex::ParticleReal &zeval) const
Definition: SoftQuad.H:435
string name
amrex::Gpu::DeviceVector< amrex::ParticleReal > m_cos_coef
Definition: SoftQuad.H:473
Definition: SoftQuad.H:116
s
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE amrex::ParticleReal beta_gamma() const
Definition: ReferenceParticle.H:79
static constexpr amrex::Real pi
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE GpuComplex< T > sqrt(const GpuComplex< T > &a_z) noexcept
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:344
#define AMREX_GPU_HOST
void synchronize() noexcept
void copyAsync(HostToDevice, InIter begin, InIter end, OutIter result) noexcept
Definition: SoftQuad.H:52