ImpactX
SoftSol.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_SOFTSOL_H
11 #define IMPACTX_SOFTSOL_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 {
47  {
49  0.350807812299706,
50  0.323554693720069,
51  0.260320578919415,
52  0.182848575294969,
53  0.106921016050403,
54  4.409581845710694E-002,
55  -9.416427163897508E-006,
56  -2.459452716865687E-002,
57  -3.272762575737291E-002,
58  -2.936414401076162E-002,
59  -1.995780078926890E-002,
60  -9.102893342953847E-003,
61  -2.456410658713271E-006,
62  5.788233017324325E-003,
63  8.040408292420691E-003,
64  7.480064552867431E-003,
65  5.230254569468851E-003,
66  2.447614547094685E-003,
67  -1.095525090532255E-006,
68  -1.614586867387170E-003,
69  -2.281365457438345E-003,
70  -2.148709081338292E-003,
71  -1.522541739363011E-003,
72  -7.185505862719508E-004,
73  -6.171194824600157E-007,
74  4.842109305036943E-004,
75  6.874508102002901E-004,
76  6.535550288205728E-004,
77  4.648795813759210E-004,
78  2.216564722797528E-004,
79  -4.100982995210341E-007,
80  -1.499332112463395E-004,
81  -2.151538438342482E-004,
82  -2.044590946652016E-004,
83  -1.468242784844341E-004
84  };
85 
87  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
88  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
89  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
90  0, 0, 0, 0, 0
91  };
92  };
93 
100 namespace SoftSolenoidData
101 {
103  static inline int next_id = 0;
104 
106  static inline std::map<int, std::vector<amrex::ParticleReal>> h_cos_coef = {};
108  static inline std::map<int, std::vector<amrex::ParticleReal>> h_sin_coef = {};
109 
111  static inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_cos_coef = {};
113  static inline std::map<int, amrex::Gpu::DeviceVector<amrex::ParticleReal>> d_sin_coef = {};
114 
115 } // namespace SoftSolenoidData
116 
118  : public elements::BeamOptic<SoftSolenoid>,
119  public elements::Thick
120  {
121  static constexpr auto name = "SoftSolenoid";
123 
135  amrex::ParticleReal ds,
136  amrex::ParticleReal bscale,
137  std::vector<amrex::ParticleReal> cos_coef,
138  std::vector<amrex::ParticleReal> sin_coef,
139  int mapsteps = 1,
140  int nslice = 1
141  )
142  : Thick(ds, nslice),
143  m_bscale(bscale), m_mapsteps(mapsteps), m_id(SoftSolenoidData::next_id)
144  {
145  // next created soft solenoid has another id for its data
147 
148  // validate sin and cos coefficients are the same length
149  m_ncoef = int(cos_coef.size());
150  if (m_ncoef != int(sin_coef.size()))
151  throw std::runtime_error("SoftSolenoid: cos and sin coefficients must have same length!");
152 
153  // host data
158 
159  // device data
163  cos_coef.begin(), cos_coef.end(),
166  sin_coef.begin(), sin_coef.end(),
169 
170  // low-level objects we can use on device
173  }
174 
176  using BeamOptic::operator();
177 
189  PType& AMREX_RESTRICT p,
190  amrex::ParticleReal & AMREX_RESTRICT px,
191  amrex::ParticleReal & AMREX_RESTRICT py,
192  amrex::ParticleReal & AMREX_RESTRICT pt,
193  [[maybe_unused]] RefPart const & refpart
194  ) const
195  {
196  using namespace amrex::literals; // for _rt and _prt
197 
198  // access AoS data such as positions and cpu/id
199  amrex::ParticleReal const x = p.pos(RealAoS::x);
200  amrex::ParticleReal const y = p.pos(RealAoS::y);
201  amrex::ParticleReal const t = p.pos(RealAoS::t);
202 
203  // initialize output values of momenta
204  amrex::ParticleReal pxout = px;
205  amrex::ParticleReal pyout = py;
206  amrex::ParticleReal ptout = pt;
207 
208  // get the linear map
210 
211  // symplectic linear map for a solenoid is computed using the
212  // Hamiltonian formalism as described in:
213  // https://uspas.fnal.gov/materials/09UNM/ComputationalMethods.pdf.
214  // R denotes the transfer matrix in the basis (x,px,y,py,t,pt),
215  // so that, e.g., R(3,4) = dyf/dpyi.
216 
217  // push particles using the linear map
218  p.pos(RealAoS::x) = R(1,1)*x + R(1,2)*px + R(1,3)*y
219  + R(1,4)*py + R(1,5)*t + R(1,6)*pt;
220  pxout = R(2,1)*x + R(2,2)*px + R(2,3)*y
221  + R(2,4)*py + R(2,5)*t + R(2,6)*pt;
222  p.pos(RealAoS::y) = R(3,1)*x + R(3,2)*px + R(3,3)*y
223  + R(3,4)*py + R(3,5)*t + R(3,6)*pt;
224  pyout = R(4,1)*x + R(4,2)*px + R(4,3)*y
225  + R(4,4)*py + R(4,5)*t + R(4,6)*pt;
226  p.pos(RealAoS::t) = R(5,1)*x + R(5,2)*px + R(5,3)*y
227  + R(5,4)*py + R(5,5)*t + R(5,6)*pt;
228  ptout = R(6,1)*x + R(6,2)*px + R(6,3)*y
229  + R(6,4)*py + R(6,5)*t + R(6,6)*pt;
230 
231  // assign updated momenta
232  px = pxout;
233  py = pyout;
234  pt = ptout;
235  }
236 
242  void operator() (RefPart & AMREX_RESTRICT refpart) const
243  {
244  using namespace amrex::literals; // for _rt and _prt
245 
246  // assign input reference particle values
247  amrex::ParticleReal const x = refpart.x;
248  amrex::ParticleReal const px = refpart.px;
249  amrex::ParticleReal const y = refpart.y;
250  amrex::ParticleReal const py = refpart.py;
251  amrex::ParticleReal const z = refpart.z;
252  amrex::ParticleReal const pz = refpart.pz;
253  amrex::ParticleReal const pt = refpart.pt;
254  amrex::ParticleReal const s = refpart.s;
255  amrex::ParticleReal const sedge = refpart.sedge;
256 
257  // initialize linear map (deviation) values
258  for (int i=1; i<7; i++) {
259  for (int j=1; j<7; j++) {
260  auto const default_value = (i == j) ? 1.0_prt : 0.0_prt;
261  refpart.map(i, j) = default_value;
262  }
263  }
264 
265  // length of the current slice
266  amrex::ParticleReal const slice_ds = m_ds / nslice();
267 
268  // compute intial value of beta*gamma
269  amrex::ParticleReal const bgi = sqrt(pow(pt, 2) - 1.0_prt);
270 
271  // call integrator to advance (t,pt)
272  amrex::ParticleReal const zin = s - sedge;
273  amrex::ParticleReal const zout = zin + slice_ds;
274  int const nsteps = m_mapsteps;
275 
276  integrators::symp2_integrate_split3(refpart,zin,zout,nsteps,*this);
277  amrex::ParticleReal const ptf = refpart.pt;
278 
279  /* print computed linear map:
280  for(int i=1; i<7; ++i){
281  for(int j=1; j<7; ++j){
282  amrex::PrintToFile("SolMap.txt") << i << " " <<
283  j << " " << refpart.map(i,j) << "\n";
284  }
285  }
286  */
287 
288  // advance position (x,y,z)
289  refpart.x = x + slice_ds*px/bgi;
290  refpart.y = y + slice_ds*py/bgi;
291  refpart.z = z + slice_ds*pz/bgi;
292 
293  // compute final value of beta*gamma
294  amrex::ParticleReal const bgf = sqrt(pow(ptf, 2) - 1.0_prt);
295 
296  // advance momentum (px,py,pz)
297  refpart.px = px*bgf/bgi;
298  refpart.py = py*bgf/bgi;
299  refpart.pz = pz*bgf/bgi;
300 
301  // advance integrated path length
302  refpart.s = s + slice_ds;
303  }
304 
311  std::tuple<amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal>
313  Sol_Bfield (amrex::ParticleReal const zeval) const
314  {
315  using namespace amrex::literals; // for _rt and _prt
316 
317  // pick the right data depending if we are on the host side
318  // (reference particle push) or device side (particles):
319 #if AMREX_DEVICE_COMPILE
320  amrex::ParticleReal* cos_data = m_cos_d_data;
321  amrex::ParticleReal* sin_data = m_sin_d_data;
322 #else
323  amrex::ParticleReal* cos_data = m_cos_h_data;
324  amrex::ParticleReal* sin_data = m_sin_h_data;
325 #endif
326 
327  // specify constants
329  amrex::ParticleReal const zlen = m_ds;
330  amrex::ParticleReal const zmid = zlen / 2.0_prt;
331 
332  // compute on-axis magnetic field (z is relative to solenoid midpoint)
333  amrex::ParticleReal bfield = 0.0;
334  amrex::ParticleReal bfieldp = 0.0;
335  amrex::ParticleReal bfieldint = 0.0;
336  amrex::ParticleReal const z = zeval - zmid;
337 
338  if (std::abs(z) <= zmid)
339  {
340  bfield = 0.5_prt*cos_data[0];
341  bfieldint = z*bfield;
342  for (int j=1; j < m_ncoef; ++j)
343  {
344  bfield = bfield + cos_data[j]*cos(j*2*pi*z/zlen) +
345  sin_data[j]*sin(j*2*pi*z/zlen);
346  bfieldp = bfieldp-j*2*pi*cos_data[j]*sin(j*2*pi*z/zlen)/zlen +
347  j*2*pi*sin_data[j]*cos(j*2*pi*z/zlen)/zlen;
348  bfieldint = bfieldint + zlen*cos_data[j]*sin(j*2*pi*z/zlen)/(j*2*pi) -
349  zlen*sin_data[j]*cos(j*2*pi*z/zlen)/(j*2*pi);
350  }
351  }
352  return std::make_tuple(bfield, bfieldp, bfieldint);
353  }
354 
364  void map1 (amrex::ParticleReal const tau,
365  RefPart & refpart,
366  [[maybe_unused]] amrex::ParticleReal & zeval) const
367  {
368  using namespace amrex::literals; // for _rt and _prt
369 
370  // push the reference particle
371  amrex::ParticleReal const t = refpart.t;
372  amrex::ParticleReal const pt = refpart.pt;
373  amrex::ParticleReal const z = zeval;
374 
375  if (pt < -1.0_prt) {
376  refpart.t = t + tau/sqrt(1.0_prt - pow(pt, -2));
377  refpart.pt = pt;
378  }
379  else {
380  refpart.t = t;
381  refpart.pt = pt;
382  }
383 
384  zeval = z + tau;
385 
386  // push the linear map equations
388  amrex::ParticleReal const betgam = refpart.beta_gamma();
389 
390  refpart.map(1,1) = R(1,1) + tau*R(2,1);
391  refpart.map(1,2) = R(1,2) + tau*R(2,2);
392  refpart.map(1,3) = R(1,3) + tau*R(2,3);
393  refpart.map(1,4) = R(1,4) + tau*R(2,4);
394 
395  refpart.map(3,1) = R(3,1) + tau*R(4,1);
396  refpart.map(3,2) = R(3,2) + tau*R(4,2);
397  refpart.map(3,3) = R(3,3) + tau*R(4,3);
398  refpart.map(3,4) = R(3,4) + tau*R(4,4);
399 
400  refpart.map(5,5) = R(5,5) + tau*R(6,5)/pow(betgam,2);
401  refpart.map(5,6) = R(5,6) + tau*R(6,6)/pow(betgam,2);
402 
403  }
404 
414  void map2 (amrex::ParticleReal const tau,
415  RefPart & refpart,
416  amrex::ParticleReal & zeval) const
417  {
418  using namespace amrex::literals; // for _rt and _prt
419 
420  amrex::ParticleReal const t = refpart.t;
421  amrex::ParticleReal const pt = refpart.pt;
422 
423  // Define parameters and intermediate constants
424  amrex::ParticleReal const B0 = m_bscale;
425 
426  // push the reference particle
427  auto [bz, bzp, bzint] = Sol_Bfield(zeval);
428  amrex::ignore_unused(bzp, bzint);
429 
430  refpart.t = t;
431  refpart.pt = pt;
432 
433  // push the linear map equations
435  amrex::ParticleReal const alpha = B0*bz/2.0_prt;
436  amrex::ParticleReal const alpha2 = pow(alpha,2);
437 
438  refpart.map(2,1) = R(2,1) - tau*alpha2*R(1,1);
439  refpart.map(2,2) = R(2,2) - tau*alpha2*R(1,2);
440  refpart.map(2,3) = R(2,3) - tau*alpha2*R(1,3);
441  refpart.map(2,4) = R(2,4) - tau*alpha2*R(1,4);
442 
443  refpart.map(4,1) = R(4,1) - tau*alpha2*R(3,1);
444  refpart.map(4,2) = R(4,2) - tau*alpha2*R(3,2);
445  refpart.map(4,3) = R(4,3) - tau*alpha2*R(3,3);
446  refpart.map(4,4) = R(4,4) - tau*alpha2*R(3,4);
447 
448  }
449 
459  void map3 (amrex::ParticleReal const tau,
460  RefPart & refpart,
461  amrex::ParticleReal & zeval) const
462  {
463  using namespace amrex::literals; // for _rt and _prt
464 
465  amrex::ParticleReal const t = refpart.t;
466  amrex::ParticleReal const pt = refpart.pt;
467  amrex::ParticleReal const z = zeval;
468 
469  // Define parameters and intermediate constants
470  amrex::ParticleReal const B0 = m_bscale;
471 
472  // push the reference particle
473  auto [bz, bzp, bzint] = Sol_Bfield(z);
474  amrex::ignore_unused(bzp, bzint);
475 
476  refpart.t = t;
477  refpart.pt = pt;
478 
479  // push the linear map equations
481  amrex::ParticleReal const theta = tau*B0*bz/2.0_prt;
482  amrex::ParticleReal const cs = cos(theta);
483  amrex::ParticleReal const sn = sin(theta);
484 
485  refpart.map(1,1) = R(1,1)*cs + R(3,1)*sn;
486  refpart.map(1,2) = R(1,2)*cs + R(3,2)*sn;
487  refpart.map(1,3) = R(1,3)*cs + R(3,3)*sn;
488  refpart.map(1,4) = R(1,4)*cs + R(3,4)*sn;
489 
490  refpart.map(2,1) = R(2,1)*cs + R(4,1)*sn;
491  refpart.map(2,2) = R(2,2)*cs + R(4,2)*sn;
492  refpart.map(2,3) = R(2,3)*cs + R(4,3)*sn;
493  refpart.map(2,4) = R(2,4)*cs + R(4,4)*sn;
494 
495  refpart.map(3,1) = R(3,1)*cs - R(1,1)*sn;
496  refpart.map(3,2) = R(3,2)*cs - R(1,2)*sn;
497  refpart.map(3,3) = R(3,3)*cs - R(1,3)*sn;
498  refpart.map(3,4) = R(3,4)*cs - R(1,4)*sn;
499 
500  refpart.map(4,1) = R(4,1)*cs - R(2,1)*sn;
501  refpart.map(4,2) = R(4,2)*cs - R(2,2)*sn;
502  refpart.map(4,3) = R(4,3)*cs - R(2,3)*sn;
503  refpart.map(4,4) = R(4,4)*cs - R(2,4)*sn;
504 
505  }
506 
509  void
511  {
512  // remove from unique data map
517 
522  }
523 
524  private:
525  amrex::ParticleReal m_bscale;
527  int m_id;
528 
529  int m_ncoef = 0;
530  amrex::ParticleReal* m_cos_h_data = nullptr;
531  amrex::ParticleReal* m_sin_h_data = nullptr;
532  amrex::ParticleReal* m_cos_d_data = nullptr;
533  amrex::ParticleReal* m_sin_d_data = nullptr;
534  };
535 
536 } // namespace impactx
537 
538 #endif // IMPACTX_SOFTSOL_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, std::vector< amrex::ParticleReal > > h_cos_coef
host: cosine coefficients in Fourier expansion of on-axis magnetic field Bz
Definition: SoftSol.H:106
static int next_id
last used id for a created soft solenoid
Definition: SoftSol.H:103
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: SoftSol.H:113
static std::map< int, std::vector< amrex::ParticleReal > > h_sin_coef
host: sine coefficients in Fourier expansion of on-axis magnetic field Bz
Definition: SoftSol.H:108
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: SoftSol.H:111
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void symp2_integrate_split3(RefPart &refpart, amrex::ParticleReal const zin, amrex::ParticleReal const zout, int const nsteps, T_Element const &element)
Definition: Integrators.H:53
Definition: ImpactX.cpp:34
s
int nsteps
int count
@ 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, 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: SoftSol.H:120
void finalize()
Definition: SoftSol.H:510
static constexpr auto name
Definition: SoftSol.H:121
amrex::ParticleReal * m_cos_d_data
non-owning pointer to host sine coefficients
Definition: SoftSol.H:532
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void map2(amrex::ParticleReal const tau, RefPart &refpart, amrex::ParticleReal &zeval) const
Definition: SoftSol.H:414
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void map1(amrex::ParticleReal const tau, RefPart &refpart, [[maybe_unused]] amrex::ParticleReal &zeval) const
Definition: SoftSol.H:364
SoftSolenoid(amrex::ParticleReal ds, amrex::ParticleReal bscale, std::vector< amrex::ParticleReal > cos_coef, std::vector< amrex::ParticleReal > sin_coef, int mapsteps=1, int nslice=1)
Definition: SoftSol.H:134
std::tuple< amrex::ParticleReal, amrex::ParticleReal, amrex::ParticleReal > AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE Sol_Bfield(amrex::ParticleReal const zeval) const
Definition: SoftSol.H:313
int m_ncoef
unique soft solenoid id used for data lookup map
Definition: SoftSol.H:529
amrex::ParticleReal * m_cos_h_data
number of Fourier coefficients
Definition: SoftSol.H:530
int m_mapsteps
scaling factor for solenoid Bz field
Definition: SoftSol.H:526
AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE void map3(amrex::ParticleReal const tau, RefPart &refpart, amrex::ParticleReal &zeval) const
Definition: SoftSol.H:459
int m_id
number of map integration steps per slice
Definition: SoftSol.H:527
amrex::ParticleReal * m_sin_h_data
non-owning pointer to host cosine coefficients
Definition: SoftSol.H:531
ImpactXParticleContainer::ParticleType PType
Definition: SoftSol.H:122
amrex::ParticleReal m_bscale
Definition: SoftSol.H:525
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: SoftSol.H:188
amrex::ParticleReal * m_sin_d_data
non-owning pointer to device cosine coefficients
Definition: SoftSol.H:533
Definition: SoftSol.H:47
amrex::Vector< amrex::ParticleReal > default_cos_coef
Definition: SoftSol.H:48
amrex::Vector< amrex::ParticleReal > default_sin_coef
Definition: SoftSol.H:86
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