XGCa
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
species.hpp
Go to the documentation of this file.
1 #ifndef SPECIES_HPP
2 #define SPECIES_HPP
3 #include <Cabana_AoSoA.hpp>
4 #include <Cabana_DeepCopy.hpp>
5 #include <Kokkos_Core.hpp>
6 #include "NamelistReader.hpp"
7 #include "timer_macro.hpp"
8 #include "magnetic_field.hpp"
9 #include "grid.hpp"
10 #include "domain_decomposition.hpp"
11 #include "particles.hpp"
12 #include "space_settings.hpp"
13 #include "distribution.hpp"
14 #include "profile.hpp"
15 #include "gyro_avg_mat.hpp"
17 
18 extern "C" void set_spall_num_and_ptr(int idx, int n_ptl, int n_vecs, VecParticles* ptl);
19 extern "C" void set_min_max_num(int isp, int n_ptl);
20 extern "C" void adjust_n_ptl_for_core_ptl(int* n_ptl);
21 
26  return false;
27 }
28 
33  return false;
34 }
35 
36 // Used for Cabana slices (getting one particle property at a time)
37 namespace PtlSlice{
38 #ifdef ESC_PTL
39 enum{Ph=0,Ct,Gid,Flag};
40 #else
41 enum{Ph=0,Ct,Gid};
42 #endif
43 }
44 
45 struct PtlMvmt{
46  // Options for particle location management when looping over particles
47  enum SendOpt{
48  NoSend=0,
51  };
52 
53  enum ReturnOpt{
57  };
58 
61 
62  PtlMvmt(SendOpt send_opt, ReturnOpt return_opt) : send_opt(send_opt), return_opt(return_opt){}
63 };
64 
68 };
69 
70 // Species class
71 template<class Device>
72 class Species{
73  public:
74 
75  int idx;
76  bool is_electron;
77  bool is_adiabatic;
80  double mass;
81  double charge;
82  double charge_eu;
83  double c_m;
84  double c2_2m;
85 
86  bool is_deltaf;
87 
88  int ncycles;
90 
92  int n_ptl;
93  Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN> particles;
94 
95  // Device particles
96  Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> particles_d;
97 
99 
102 
103  /*** Could be its own class inside species? ***/
107 
108  // phase0 (for ion restoration)
109  Cabana::AoSoA<PhaseDataTypes,HostType,VEC_LEN> phase0;
110  Cabana::AoSoA<PhaseDataTypes,Device,VEC_LEN> phase0_d;
111 
112  // For electron restoration
113  Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN> backup_particles;
114  int n_backup_particles; // Number of particles stored in backup_particles (can't be deduced from its size due to buffer)
115  /****/
116 
118 
120 
121  Eq::Profile<Device> eq_temp; // Equilibrium temperature
122  Eq::Profile<Device> eq_den; // Equilibrium density
123  Eq::Profile<Device> eq_flow; // Equilibrium flow
124  int eq_flow_type; // Type of Equilibirum flow
125 
127 
128  /*** Constructors ***/
129 
130  Species(int idx_in, int nonadiabatic_idx_in, bool is_electron_in, bool is_adiabatic_in, KinType kintype_in, double mass_in, double charge_in, double charge_eu_in, bool is_deltaf_in,
131  int ncycles_in);
132 
133  template<class GridDevice>
134  Species(NLReader::NamelistReader& nlr, const Grid<GridDevice> &grid, const MagneticField<DeviceType> &magnetic_field, const DomainDecomposition<DeviceType>& pol_decomp, int idx_in, int nonadiabatic_idx_in);
135 
136  // Electron or ion default constructor
137  Species(SpeciesType sp_type, int n_ptl)
138  : idx(sp_type==ELECTRON ? 0 : 1),
139  is_electron(sp_type==ELECTRON),
140  mass(is_electron ? 3.344e-30 : PROTON_MASS),
142  charge_eu(is_electron ? -1.0 : 1.0),
143  is_deltaf(true),
144  is_adiabatic(false),
145  nonadiabatic_idx(idx), // Since is_adiabatic is false above
147  ncycles(is_electron ? 70 : 1),
148  c_m(charge/mass),
149  c2_2m(0.5*charge*charge/mass),
151  n_ptl(n_ptl),
152  backup_particles("backup_particles", 0),
153  particles("particles", add_vec_buffer(n_ptl)),
156  eq_temp(1.0e3,-0.1),
157  eq_den(1.0e19,-0.1),
158  eq_flow_type(2),
159  owns_particles_d(false),
163 
164  // Special constructor for tests that involve tracking particles in memory
165  // The idea is to use these particles to test functions that reorder particles,
166  // e.g. sort, shift, and cleaning
167  Species(int n_ptl_in)
168  : n_ptl(n_ptl_in),
169  is_electron(true),
170  is_adiabatic(false),
171  particles("particles", add_vec_buffer(n_ptl_in)),
173  owns_particles_d(false),
178  {
179 
180  // Slice particle properties
181  auto ph = Cabana::slice<PtlSlice::Ph>(particles);
182  auto ct = Cabana::slice<PtlSlice::Ct>(particles);
183  auto gid = Cabana::slice<PtlSlice::Gid>(particles);
184 #ifdef ESC_PTL
185  auto flag = Cabana::slice<PtlSlice::Flag>(particles);
186 #endif
187 
188  // Offset gid if using MPI
189 #ifdef USE_MPI
190  long long int gid_offset = n_ptl*SML_COMM_RANK;
191 #else
192  long long int gid_offset = 0;
193 #endif
194 
195  // Assign trackable values
196  for (int i=0;i<n_ptl;i++){
197  // Set GID in order
198  gid(i) = gid_offset + i+1; // 1-indexed
199 
200  // Value of properties is gid + 0.1*(property index)
201  // First particle: (1.0, 1.1, ... 1.8)
202  // Second particle: (2.0, 2.1, ... 2.8)
203  for (int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
204  for (int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
205  }
206 
207  // Buffer particles: same but with gid = -1
208  if(n_ptl>0){
209  for (int i=n_ptl;i<add_vec_buffer(n_ptl);i++){
210  gid(i) = -1;
211  for (int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
212  for (int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
213  }
214  }
215  }
216 
217  void resize_particles(int new_n_ptl){
218  n_ptl = new_n_ptl;
219 
220 #ifndef USE_GPU
221  // If CPU-only, particles_d points to the same location as particles. If particles is resized, then Cabana will not deallocate the first allocation
222  // since it is still used by particles_d. So, reset particles_d before resize, and point it back to particles only afterwards
223  particles_d = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>();
224 #endif
225 
226  particles.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
227  particles.resize(add_vec_buffer(n_ptl));
228 
229 #ifndef USE_GPU
230  // Point particles_d back to particles after resize
232 #endif
233 
235  }
236 
238 #ifdef USE_GPU
239  particles.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
240  particles.resize(particles_d.size());
241 #else
242  // Point particles back to particles_d after resize
244 #endif
245 
247  }
248 
249  /* If using CPU-only, then "device" particles are a shallow copy of host particles so that
250  * there is no unnecessary duplication. When "device" particles are resized, Cabana will keep the
251  * original allocation if there is a second reference (i.e. host particles).
252  * To resolve this, we free the host particles here so that there is no second reference
253  * */
255  particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>();
256  }
257 
258  /* Resizes device particles if on GPU, or just creates a shallow copy if CPU only
259  * */
261  if(!owns_particles_d) exit_XGC("\nSpecies tried to resize device particles, but doesn't own the device array.");
262 
263 #ifdef USE_GPU
264  particles_d.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
265  // Resize device particles to match n particles
267 #else
268  // If kernels are on CPU, do shallow copy
270 #endif
271  }
272 
273  /* Resizes device particles
274  * */
275  void resize_device_particles(int new_n_ptl){
276  if(!owns_particles_d) exit_XGC("\nSpecies tried to resize device particles, but doesn't own the device array.");
277 
278  n_ptl = new_n_ptl;
279 
280  particles_d.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
281 
282  // Resize device particles to match host particles
284  }
285 
286  /* Copies particles to device - deep copy if using GPU, otherwise shallow copy
287  * Also takes the opportunity to set the buffer particles to realistic values
288  * */
290  if(!owns_particles_d) exit_XGC("\nSpecies tried to copy particles to device, but doesn't own the device array.");
291 
292 #ifdef PRINT_COPA_ARTIFACTS
293  if(is_rank_zero()) printf("\nCopying Cabana AoSoA particle data to GPU with Cabana::deep_copy\n");
294 #endif
295 
296 #ifdef USE_GPU
297  // Copy to device
298  Cabana::deep_copy(particles_d, particles);
299 #else
300  // No operation required if CPU-only
301 #endif
302 
303  // Copy last particle to fill remainder of trailing vector in AoSoA
305  }
306 
307  /* Copies particles from device - deep copy if using GPU, otherwise no copy is necessary
308  * */
310  if(!owns_particles_d) exit_XGC("\nSpecies tried to copy particles from device, but doesn't own the device array.");
311 
312 #ifdef PRINT_COPA_ARTIFACTS
313  if(is_rank_zero()) printf("\nCopying Cabana AoSoA particle data from GPU with Cabana::deep_copy\n");
314 #endif
315 
316 #ifdef USE_GPU
317  // Copy particles to host
318  Cabana::deep_copy(particles, particles_d);
319 #else
320  // No operation required if CPU-only
321 #endif
322  }
323 
324  /* Copies particles to device if they are resident on the device
325  * */
328  }
329 
330  /* Copies particles from device if they are resident on the device
331  * */
334  }
335 
336  /* Copies particles to device if they are NOT resident on the device
337  * */
340  }
341 
342  /* Copies particles from device if they are NOT resident on the device
343  * */
346  }
347 
356  if (n_ptl>0){
357  int last_ptl_index = n_ptl - 1;
358  auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
359  auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
360  auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
361 #ifdef ESC_PTL
362  auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
363 #endif
364 
365  Kokkos::parallel_for("set_buffer_particles_d", Kokkos::RangePolicy<ExSpace>( n_ptl, add_vec_buffer(n_ptl) ), KOKKOS_LAMBDA( const int i ){
366  // Buffer particles: same as last particle, gid = -1
367  for (int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
368  for (int j=0;j<3;j++) ct(i, j) = ct(last_ptl_index, j);
369  gid(i) = -1;
370 #ifdef ESC_PTL
371  flag(i) = flag(last_ptl_index);
372 #endif
373  });
374  }
375  }
376 
385  if (n_ptl>0){
386  int last_ptl_index = n_ptl - 1;
387  auto ph = Cabana::slice<PtlSlice::Ph>(phase0_d);
388 
389  Kokkos::parallel_for("set_buffer_phase0", Kokkos::RangePolicy<ExSpace>( n_ptl, add_vec_buffer(n_ptl) ), KOKKOS_LAMBDA( const int i ){
390  // copy final real particle
391  for (int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
392  });
393  }
394  }
395 
396  // Options for custom launch bounds since kokkos defaults are suboptimal for electron push kernel
397  enum class LaunchBounds{
398  Default,
399  Custom
400  };
401 
407  template<typename F>
408  inline void for_all_particles(const std::string label, F lambda_func) const {
409 #ifdef PRINT_COPA_ARTIFACTS
410  if(is_rank_zero()) printf("\nLaunching GPU kernel '%s' on Cabana AoSoA of particles\n", label.c_str());
411 #endif
412  Kokkos::RangePolicy<ExSpace> particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
413  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
414  }
415 
416  inline void back_up_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA, int offset, int n) const{
417  auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
418  auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
419  auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
420 #ifdef ESC_PTL
421  auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
422 #endif
423 
424  auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
425  auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
426  auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
427 #ifdef ESC_PTL
428  auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
429 #endif
430 
431  Kokkos::parallel_for("backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA( const int i ){
432  int i_offset = i + offset;
433  // Make backup copy
434  for (int j=0;j<6;j++) ph_b(i, j) = ph(i_offset, j);
435  for (int j=0;j<3;j++) ct_b(i, j) = ct(i_offset, j);
436  gid_b(i) = gid(i_offset);
437 #ifdef ESC_PTL
438  flag_b(i) = flag(i_offset);
439 #endif
440  // Deactivate
441  gid(i_offset) = -1;
442  });
443  }
444 
445  inline void restore_backup_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA, int offset, int n) const{
446  auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
447  auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
448  auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
449 #ifdef ESC_PTL
450  auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
451 #endif
452 
453  auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
454  auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
455  auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
456 #ifdef ESC_PTL
457  auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
458 #endif
459 
460  Kokkos::parallel_for("backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA( const int i ){
461  int i_offset = i + offset;
462  // Restore from backup copy
463  for (int j=0;j<6;j++) ph(i_offset, j) = ph_b(i, j);
464  for (int j=0;j<3;j++) ct(i_offset, j) = ct_b(i, j);
465  gid(i_offset) = gid_b(i);
466 #ifdef ESC_PTL
467  flag(i_offset) = flag_b(i);
468 #endif
469  });
470  }
471 
477  template<typename F>
478  inline void for_particle_range(int begin_idx, int end_idx, const std::string label, F lambda_func) const {
479  if(end_idx <= begin_idx) return; // Return if range is 0 or less
480 
481  // Still need the subset to line up with the AoSoA vector length
482  int first_soa = begin_idx/VEC_LEN;
483  int n_other_ptl_in_first_soa = begin_idx - first_soa*VEC_LEN;
484  bool first_soa_is_partial = (n_other_ptl_in_first_soa>0);
485  int last_soa = (end_idx-1)/VEC_LEN;
486  int n_other_ptl_in_last_soa = (last_soa+1)*VEC_LEN - end_idx;
487  bool last_soa_is_partial = (n_other_ptl_in_last_soa>0);
488 
489 #ifdef USE_GPU
490  int first_item_in_shifted_range = first_soa*VEC_LEN;
491 #else
492  int first_item_in_shifted_range = first_soa;
493 #endif
494 
495  Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_first_soa;
496  Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_last_soa;
497  if(first_soa_is_partial){
498  // Make a backup of the first SoA and set the particles_d GIDs to -1
499  ptl_first_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>("ptl_first_soa", n_other_ptl_in_first_soa);
500  back_up_SoA(ptl_first_soa, first_soa*VEC_LEN, n_other_ptl_in_first_soa);
501  }
502  if(last_soa_is_partial){
503  // Make a backup of the last SoA and set the particles_d GIDs to -1
504  ptl_last_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>("ptl_last_soa", n_other_ptl_in_last_soa);
505  back_up_SoA(ptl_last_soa, end_idx, n_other_ptl_in_last_soa);
506  }
507 
508  // Finally, do the parallel_for
509  Kokkos::RangePolicy<ExSpace> particle_range_policy( first_item_in_shifted_range, p_range<DeviceType>(end_idx) );
510  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
511 
512  // Restore from backup
513  if(first_soa_is_partial){
514  restore_backup_SoA(ptl_first_soa, first_soa*VEC_LEN, n_other_ptl_in_first_soa);
515  }
516  if(last_soa_is_partial){
517  restore_backup_SoA(ptl_last_soa, end_idx, n_other_ptl_in_last_soa);
518  }
519  }
520 
528  template<typename F>
529  inline void for_all_particles(const std::string label, F lambda_func,
530  const PtlMvmt mvmt, LaunchBounds launch_bounds=LaunchBounds::Default) {
531  if(!owns_particles_d) exit_XGC("\nSpecies tried to loop over particles on device, but doesn't own the device array.");
532 
533 #ifdef PRINT_COPA_ARTIFACTS
534  if(is_rank_zero()) printf("\nLaunching GPU kernel '%s' on Cabana AoSoA of particles\n", label.c_str());
535 #endif
536 
537  bool use_streaming = stream_particles;
538 #ifndef USE_STREAMS
539  use_streaming = false; // Just to be safe, turn streams off here
540 #endif
541 
542  bool send_ptl = ( mvmt.send_opt==PtlMvmt::Send ||
544  bool return_ptl = ( mvmt.return_opt==PtlMvmt::Return ||
546 
547  // Don't need to stream if particles are already present and don't need to be returned
548  if((!send_ptl) && (!return_ptl)) use_streaming = false;
549 
550  if(use_streaming){
551 #ifdef USE_STREAMS
552  Streamed::Option stream_option = Streamed::Normal; // Send to device and back
553  if(!send_ptl) stream_option = Streamed::NoSend;
554  if(!return_ptl) stream_option = Streamed::NoReturn;
555 
556  // Execute streaming parallel_for
557  Streamed::parallel_for(label, n_ptl, lambda_func, stream_option, particles, particles_d);
558 #endif
559  }else{
560  if(send_ptl) TIMER("copy_ptl_to_device",copy_particles_to_device() );
561 
562  if (launch_bounds==LaunchBounds::Custom) {
563 #ifdef USE_EPUSH_LAUNCH_BOUNDS
564 # if !defined(PUSH_MAX_THREADS_PER_BLOCK) || !defined(PUSH_MIN_WARPS_PER_EU)
565 # error "USE_EPUSH_LAUNCH_BOUNDS requires PUSH_MAX_THREADS_PER_BLOCK and PUSH_MIN_WARPS_PER_EU to be defined"
566 # endif
567  Kokkos::RangePolicy<ExSpace, Kokkos::LaunchBounds<PUSH_MAX_THREADS_PER_BLOCK, PUSH_MIN_WARPS_PER_EU>>
568  particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
569  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
570 #else
571  exit_XGC("\nERROR: LaunchBounds::Custom specified, but USE_EPUSH_LAUNCH_BOUNDS is not defined\n");
572 #endif
573  } else {
574  Kokkos::RangePolicy<ExSpace>
575  particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
576  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
577  }
578 
579  if(return_ptl) TIMER("copy_ptl_from_device", copy_particles_from_device() );
580  }
581  }
582 
583  KOKKOS_INLINE_FUNCTION VecParticles* ptl() const{
584  return (VecParticles*)(&particles_d.access(0));
585  }
586 
587  KOKKOS_INLINE_FUNCTION VecPhase* ph0() const{
588  return (VecPhase*)(&phase0_d.access(0));
589  }
590 
592  for_all_particles("copy_to_phase0", KOKKOS_LAMBDA( const int idx ){
593  AoSoAIndices<DeviceType> inds(idx);
594  VecParticles* ptl_loc = species.ptl();
595  VecPhase* ph0_loc = species.ph0();
596  for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
597  int p_vec = inds.a + i_simd;
598  ph0_loc[inds.s].r[p_vec] = ptl_loc[inds.s].ph.r[p_vec];
599  ph0_loc[inds.s].z[p_vec] = ptl_loc[inds.s].ph.z[p_vec];
600  ph0_loc[inds.s].phi[p_vec] = ptl_loc[inds.s].ph.phi[p_vec];
601  ph0_loc[inds.s].rho[p_vec] = ptl_loc[inds.s].ph.rho[p_vec];
602  ph0_loc[inds.s].w1[p_vec] = ptl_loc[inds.s].ph.w1[p_vec];
603  ph0_loc[inds.s].w2[p_vec] = ptl_loc[inds.s].ph.w2[p_vec];
604  }
605  });
606  }
607 
610  // If particles are resident on the device, then use the host particle allocation as the backup
611  // Otherwise, resize the backup_particle array
614  }else{
615  backup_particles.resize(particles_d.size());
616  }
617 
618  // Copy particle data to backup
619  Cabana::deep_copy(backup_particles, particles_d);
621  } else {
622  // For ions, save phase to phase0
623  phase0_d.resize(particles_d.size());
624  copy_to_phase0(*this); // Separate function to avoid implicit copy into lambda
625  }
627  }
628 
631  // If particles are resident on device, don't need to resize host particles since they are already used as the backup.
632  // If not, resize host particle array to fit backup particles
636  }else{
637  // If particles are not resident on device, resize host particle array to fit backup particles
639  }
640 
641  // Resize device particle array to fit backed up particles
643 
644  // Restore particle data from backup
645  Cabana::deep_copy(particles_d, backup_particles);
646 
648  // If the backup particles are simply pointing to the host particles, then
649  // point them to their own 0-sized allocation when finished using them
650  // so that they don't make a copy when host particles get resized
651  backup_particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>("backup_particles", 0);
652  }
653  } else {
654  // When ipc==2, copy phase0 to device
655  // First, resize device phase0 and reset pointer
656  phase0_d.resize(phase0.size());
657 
658  // Next copy data to phase0 on device
659  Cabana::deep_copy(phase0_d, phase0);
660 
661  // Fill buffer with realistic ptl data
663  }
664  particles_are_backed_up = false;
665  }
666 
667  KOKKOS_INLINE_FUNCTION void restore_phase_from_phase0(const AoSoAIndices<Device>& inds, SimdParticles& part_one) const {
668  VecPhase* ph0_loc = ph0();
669  for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
670  int p_vec = inds.a + i_simd;
671  part_one.ph.r[i_simd] = ph0_loc[inds.s].r[p_vec];
672  part_one.ph.z[i_simd] = ph0_loc[inds.s].z[p_vec];
673  part_one.ph.phi[i_simd] = ph0_loc[inds.s].phi[p_vec];
674  part_one.ph.rho[i_simd] = ph0_loc[inds.s].rho[p_vec];
675  part_one.ph.w1[i_simd] = ph0_loc[inds.s].w1[p_vec];
676  part_one.ph.w2[i_simd] = ph0_loc[inds.s].w2[p_vec];
677  }
678  }
679 
680  long long int get_total_n_ptl(){
681 #ifdef USE_MPI
682  long long int tmp_n_ptl = n_ptl;
683  long long int out_n_ptl = 0;
684  MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_LONG_LONG_INT, MPI_SUM, SML_COMM_WORLD);
685  return out_n_ptl;
686 #else
687  return (long long int)(n_ptl);
688 #endif
689  }
690 
692 #ifdef USE_MPI
693  int tmp_n_ptl = n_ptl;
694  int out_n_ptl = 0;
695  MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_INT, MPI_MAX, SML_COMM_WORLD);
696  return out_n_ptl;
697 #else
698  return n_ptl;
699 #endif
700  }
701 
702  // Gets the gyro_radius of a species based on equilibrium temperature
703  // inode is the LOCAL (poloidally decomposed) grid node index to get temperature
704  // smu_n is the normalized sqrt(mu)
705  // bfield is the magnetic field at inode
706  KOKKOS_INLINE_FUNCTION double get_f0_eq_gyro_radius(int inode, double smu_n, double bfield) const{
707  // Should replace UNIT_CHARGE*charge_eu with charge(?)
708  return smu_n*sqrt(mass*f0.temp_ev(inode)*EV_2_J) / (UNIT_CHARGE*charge_eu*bfield);
709  }
710 
711  // Gets the equilibrium thermal velocity of a species based on f0 temperature
712  // inode is the GLOBAL node index to get temperature
713  KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity(int inode) const{
714  return sqrt(EV_2_J*f0.temp_global(inode)/mass);
715  }
716 
717  // Gets the equilibrium thermal velocity of a species based on f0 temperature, on host
718  // inode is the local node index to get temperature
719  KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode_h(int inode) const{
720  return sqrt(EV_2_J*f0.temp_ev_h(inode)/mass);
721  }
722 };
723 
724 #include "species.tpp"
725 #endif
constexpr double EV_2_J
Conversion rate ev to J.
Definition: globals.hpp:172
Cabana::AoSoA< PhaseDataTypes, HostType, VEC_LEN > phase0
Definition: species.hpp:109
bool stream_particles
Whether to stream particles between host and device if possible.
Definition: species.hpp:101
Definition: globals.hpp:77
KOKKOS_INLINE_FUNCTION VecPhase * ph0() const
Definition: species.hpp:587
KOKKOS_INLINE_FUNCTION int divide_and_round_up(int a, int b)
Definition: globals.hpp:179
bool owns_particles_d
Whether the species owns the device particle allocation right now.
Definition: species.hpp:98
void back_up_SoA(Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > &backup_SoA, int offset, int n) const
Definition: species.hpp:416
KOKKOS_INLINE_FUNCTION double get_f0_eq_gyro_radius(int inode, double smu_n, double bfield) const
Definition: species.hpp:706
subroutine adjust_n_ptl_for_core_ptl(n_ptl)
Definition: load.F90:442
void set_spall_num_and_ptr(int idx, int n_ptl, int n_vecs, VecParticles *ptl)
bool is_rank_zero()
Definition: globals.hpp:26
void for_particle_range(int begin_idx, int end_idx, const std::string label, F lambda_func) const
Definition: species.hpp:478
Distribution< Device > f0
Species distribution in velocity space on local mesh nodes.
Definition: species.hpp:117
MPI_Comm SML_COMM_WORLD
Definition: my_mpi.cpp:4
Cabana::AoSoA< ParticleDataTypes, HostType, VEC_LEN > backup_particles
Copy of particles to be restored for RK2.
Definition: species.hpp:113
bool is_electron
Whether this species is the electrons.
Definition: species.hpp:76
void for_all_particles(const std::string label, F lambda_func, const PtlMvmt mvmt, LaunchBounds launch_bounds=LaunchBounds::Default)
Definition: species.hpp:529
void save_backup_particles()
Definition: species.hpp:608
double c2_2m
c2/2m
Definition: species.hpp:84
double rho[VEC_LEN]
Definition: particles.hpp:72
void copy_to_phase0(Species< Device > &species)
Definition: species.hpp:591
Definition: species.hpp:55
Simd< double > w1
Definition: particles.hpp:22
double c_m
c/m
Definition: species.hpp:83
Definition: species.hpp:54
bool default_streaming_option()
Definition: species.hpp:25
Eq::Profile< Device > eq_den
Definition: species.hpp:122
Definition: globals.hpp:82
KOKKOS_INLINE_FUNCTION VecParticles * ptl() const
Definition: species.hpp:583
Definition: NamelistReader.hpp:193
KinType kintype
Whether the species is gyrokinetic or drift kinetic.
Definition: species.hpp:79
Definition: magnetic_field.hpp:12
int add_vec_buffer(int n_ptl)
Definition: particles.hpp:164
int idx
Index in all_species.
Definition: species.hpp:75
Definition: particles.hpp:68
int a
The index in the inner array of the AoSoA.
Definition: particles.hpp:120
Definition: grid.hpp:53
Definition: particles.hpp:85
bool particles_are_backed_up
Whether particles are currently backed up.
Definition: species.hpp:106
int nonadiabatic_idx
Index of species skipping adiabatic species (for compatibility with fortran arrays) ...
Definition: species.hpp:78
bool default_residence_option()
Definition: species.hpp:32
int n_ptl
Number of particles.
Definition: species.hpp:92
Definition: streamed_parallel_for.hpp:16
Definition: streamed_parallel_for.hpp:14
void set_buffer_phase0_d()
Definition: species.hpp:384
long long int get_total_n_ptl()
Definition: species.hpp:680
void set_buffer_particles_d()
Definition: species.hpp:355
Simd< double > rho
Definition: particles.hpp:21
Definition: species.hpp:56
int p_range< DeviceType >(int num_particle)
Definition: particles.hpp:157
int eq_flow_type
Definition: species.hpp:124
double charge_eu
Particle charge in eu.
Definition: species.hpp:82
Definition: species.hpp:48
void resize_particles(int new_n_ptl)
Definition: species.hpp:217
double mass
Particle mass.
Definition: species.hpp:80
KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity(int inode) const
Definition: species.hpp:713
Species(int idx_in, int nonadiabatic_idx_in, bool is_electron_in, bool is_adiabatic_in, KinType kintype_in, double mass_in, double charge_in, double charge_eu_in, bool is_deltaf_in, int ncycles_in)
Definition: species.tpp:22
void for_all_particles(const std::string label, F lambda_func) const
Definition: species.hpp:408
Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > particles_d
Particles on device.
Definition: species.hpp:96
Definition: species.hpp:45
double w2[VEC_LEN]
Definition: particles.hpp:74
#define TIMER(N, F)
Definition: timer_macro.hpp:24
RKRestorationMethod
Definition: species.hpp:65
idx
Definition: diag_f0_df_port1.hpp:32
void copy_particles_to_device_if_not_resident()
Definition: species.hpp:338
RKRestorationMethod RK_restoration_method
Currently, electrons must use first method and ions must use second.
Definition: species.hpp:104
Simd< double > r
Definition: particles.hpp:18
void resize_host_particles_to_match_device()
Definition: species.hpp:237
Definition: species.hpp:66
ReturnOpt return_opt
Definition: species.hpp:60
ReturnOpt
Definition: species.hpp:53
Option
Definition: streamed_parallel_for.hpp:13
void restore_particles_from_backup()
Definition: species.hpp:629
Definition: globals.hpp:83
SendOpt send_opt
Definition: species.hpp:59
double charge
Particle charge.
Definition: species.hpp:81
SimdPhase ph
Definition: particles.hpp:59
void copy_particles_from_device()
Definition: species.hpp:309
void copy_particles_from_device_if_not_resident()
Definition: species.hpp:344
void unassign_host_particles()
Definition: species.hpp:254
int ncycles_between_sorts
Number of subcycles between sorts.
Definition: species.hpp:89
Definition: particles.hpp:58
Cabana::AoSoA< PhaseDataTypes, Device, VEC_LEN > phase0_d
Definition: species.hpp:110
int SML_COMM_RANK
Definition: my_mpi.cpp:5
KinType
Definition: globals.hpp:81
Species(SpeciesType sp_type, int n_ptl)
Definition: species.hpp:137
bool is_deltaf
Whether this species is deltaf.
Definition: species.hpp:86
VecPhase ph
Definition: particles.hpp:86
Definition: species.hpp:41
Definition: species.hpp:67
constexpr double PROTON_MASS
Definition: globals.hpp:174
void set_min_max_num(int isp, int n_ptl)
int minimum_ptl_reservation
The minimum reservation size for particles.
Definition: species.hpp:91
int s
The index in the outer array of the AoSoA.
Definition: particles.hpp:119
Simd< double > z
Definition: particles.hpp:19
void copy_particles_to_device_if_resident()
Definition: species.hpp:326
Definition: species.hpp:49
void resize_device_particles(int new_n_ptl)
Definition: species.hpp:275
Definition: species.hpp:50
KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode_h(int inode) const
Definition: species.hpp:719
void exit_XGC(std::string msg)
Definition: globals.hpp:36
void copy_particles_from_device_if_resident()
Definition: species.hpp:332
bool is_adiabatic
Whether this species is adiabatic.
Definition: species.hpp:77
Simd< double > phi
Definition: particles.hpp:20
Definition: magnetic_field.F90:1
static constexpr const Kokkos::Experimental::WorkItemProperty::HintLightWeight_t Async
Definition: space_settings.hpp:81
int n_backup_particles
Definition: species.hpp:114
Eq::Profile< Device > eq_flow
Definition: species.hpp:123
Definition: streamed_parallel_for.hpp:15
SendOpt
Definition: species.hpp:47
void copy_particles_to_device()
Definition: species.hpp:289
KOKKOS_INLINE_FUNCTION void restore_phase_from_phase0(const AoSoAIndices< Device > &inds, SimdParticles &part_one) const
Definition: species.hpp:667
Species(int n_ptl_in)
Definition: species.hpp:167
double phi[VEC_LEN]
Definition: particles.hpp:71
Simd< double > w2
Definition: particles.hpp:23
double r[VEC_LEN]
Definition: particles.hpp:69
GyroAverageMatrices< HostType > gyro_avg_matrices
Definition: species.hpp:126
Definition: species.hpp:72
void parallel_for(const std::string name, int n_ptl, Function func, Option option, HostAoSoA aosoa_h, DeviceAoSoA aosoa_d)
Definition: streamed_parallel_for.hpp:252
Definition: species.hpp:41
Eq::Profile< Device > eq_temp
Definition: species.hpp:121
bool particles_resident_on_device
Whether the particles can reside on device.
Definition: species.hpp:100
PtlMvmt(SendOpt send_opt, ReturnOpt return_opt)
Definition: species.hpp:62
int ncycles
Number of subcycles.
Definition: species.hpp:88
Definition: profile.hpp:65
int collision_grid_index
Which collision grid to use.
Definition: species.hpp:119
Definition: particles.hpp:118
SpeciesType
Definition: globals.hpp:76
double z[VEC_LEN]
Definition: particles.hpp:70
constexpr double UNIT_CHARGE
Charge of an electron (C)
Definition: globals.hpp:171
void restore_backup_SoA(Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > &backup_SoA, int offset, int n) const
Definition: species.hpp:445
Definition: distribution.hpp:10
void resize_device_particles()
Definition: species.hpp:260
int get_max_n_ptl()
Definition: species.hpp:691
LaunchBounds
Definition: species.hpp:397
Cabana::AoSoA< ParticleDataTypes, HostType, VEC_LEN > particles
Particles.
Definition: species.hpp:93
double w1[VEC_LEN]
Definition: particles.hpp:73
Definition: species.hpp:41