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 #include "basic_physics.hpp"
18 #include "memory_prediction.hpp"
19 #include "xgc_io.hpp"
20 
21 extern "C" void set_spall_num_and_ptr(int idx, int n_ptl, int n_vecs, VecParticles* ptl);
22 extern "C" void set_min_max_num(int isp, int n_ptl);
23 extern "C" void adjust_n_ptl_for_core_ptl(int* n_ptl);
24 
29  return false;
30 }
31 
36  return false;
37 }
38 
39 // Used for Cabana slices (getting one particle property at a time)
40 namespace PtlSlice{
41 #ifdef ESC_PTL
42 enum{Ph=0,Ct,Gid,Flag};
43 #else
44 enum{Ph=0,Ct,Gid};
45 #endif
46 }
47 
48 struct PtlMvmt{
49  // Options for particle location management when looping over particles
50  enum SendOpt{
51  NoSend=0,
54  };
55 
56  enum ReturnOpt{
60  };
61 
64 
65  PtlMvmt(SendOpt send_opt, ReturnOpt return_opt) : send_opt(send_opt), return_opt(return_opt){}
66 };
67 
71 };
72 
73 // Species class
74 template<class Device>
75 class Species{
76  public:
77 
78  int idx;
79  bool is_electron;
80  bool is_adiabatic;
83  double mass;
84  double charge;
85  double charge_eu;
86  double c_m;
87  double c2_2m;
88 
89  bool is_deltaf;
90 
91  int ncycles;
93 
95  int n_ptl;
96  Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN> particles;
97 
98  // Device particles
99  Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> particles_d;
100 
102 
105 
106  /*** Could be its own class inside species? ***/
110 
111  // phase0 (for ion restoration)
112  Cabana::AoSoA<PhaseDataTypes,HostType,VEC_LEN> phase0;
113  Cabana::AoSoA<PhaseDataTypes,Device,VEC_LEN> phase0_d;
114 
115  // For electron restoration
116  Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN> backup_particles;
117  int n_backup_particles; // Number of particles stored in backup_particles (can't be deduced from its size due to buffer)
118  /****/
119 
121 
123 
124  Eq::Profile<Device> eq_temp; // Equilibrium temperature
125  Eq::Profile<Device> eq_den; // Equilibrium density
126  Eq::Profile<Device> eq_flow; // Equilibrium flow
127  int eq_flow_type; // Type of Equilibirum flow
128 
129  Eq::Profile<Device> eq_fg_temp; // Equilibrium temperature
130  Eq::Profile<Device> eq_fg_den; // Equilibrium density
131  Eq::Profile<Device> eq_fg_flow; // Equilibrium flow
132  int eq_fg_flow_type; // Type of Equilibirum flow
133 
134  Eq::Profile<Device> eq_mk_temp; // Equilibrium temperature
135  Eq::Profile<Device> eq_mk_den; // Equilibrium density
136  Eq::Profile<Device> eq_mk_flow; // Equilibrium flow
137  int eq_mk_flow_type; // Type of Equilibirum flow
138 
139 
141 
142  /*** Constructors ***/
143 
144  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,
145  int ncycles_in);
146 
147  Species(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid, const MagneticField<DeviceType> &magnetic_field, const DomainDecomposition<DeviceType>& pol_decomp, int idx_in, int nonadiabatic_idx_in);
148 
149  // Electron or ion default constructor
150  Species(SpeciesType sp_type, int n_ptl)
151  : idx(sp_type==ELECTRON ? 0 : 1),
152  is_electron(sp_type==ELECTRON),
153  mass(is_electron ? 3.344e-30 : PROTON_MASS),
155  charge_eu(is_electron ? -1.0 : 1.0),
156  is_deltaf(true),
157  is_adiabatic(false),
158  nonadiabatic_idx(idx), // Since is_adiabatic is false above
160  ncycles(is_electron ? 70 : 1),
161  c_m(charge/mass),
162  c2_2m(0.5*charge*charge/mass),
164  n_ptl(n_ptl),
165  backup_particles("backup_particles", 0),
166  particles("particles", add_vec_buffer(n_ptl)),
169  eq_temp(1.0e3,-0.1),
170  eq_den(1.0e19,-0.1),
171  eq_flow_type(2),
172  owns_particles_d(false),
176 
177  // Special constructor for tests that involve tracking particles in memory
178  // The idea is to use these particles to test functions that reorder particles,
179  // e.g. sort, shift, and cleaning
180  Species(int n_ptl_in)
181  : n_ptl(n_ptl_in),
182  is_electron(true),
183  is_adiabatic(false),
184  particles("particles", add_vec_buffer(n_ptl_in)),
186  owns_particles_d(false),
191  {
192 
193  // Slice particle properties
194  auto ph = Cabana::slice<PtlSlice::Ph>(particles);
195  auto ct = Cabana::slice<PtlSlice::Ct>(particles);
196  auto gid = Cabana::slice<PtlSlice::Gid>(particles);
197 #ifdef ESC_PTL
198  auto flag = Cabana::slice<PtlSlice::Flag>(particles);
199 #endif
200 
201  // Offset gid if using MPI
202 #ifdef USE_MPI
203  long long int gid_offset = n_ptl*SML_COMM_RANK;
204 #else
205  long long int gid_offset = 0;
206 #endif
207 
208  // Assign trackable values
209  for (int i=0;i<n_ptl;i++){
210  // Set GID in order
211  gid(i) = gid_offset + i+1; // 1-indexed
212 
213  // Value of properties is gid + 0.1*(property index)
214  // First particle: (1.0, 1.1, ... 1.8)
215  // Second particle: (2.0, 2.1, ... 2.8)
216  for (int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
217  for (int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
218  }
219 
220  // Buffer particles: same but with gid = -1
221  if(n_ptl>0){
222  for (int i=n_ptl;i<add_vec_buffer(n_ptl);i++){
223  gid(i) = -1;
224  for (int j=0;j<6;j++) ph(i, j) = gid(i) + (j)*0.1;
225  for (int j=0;j<3;j++) ct(i, j) = gid(i) + (j+6)*0.1;
226  }
227  }
228  }
229 
230  static std::vector<MemoryPrediction> estimate_memory_usage(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid, const DomainDecomposition<DeviceType>& pol_decomp, int species_idx);
231 
232  static int get_initial_n_ptl(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid, const DomainDecomposition<DeviceType>& pol_decomp, int sml_special, int species_idx, bool verbose);
233 
234  void resize_particles(int new_n_ptl){
235  n_ptl = new_n_ptl;
236 
237 #ifndef USE_GPU
238  // If CPU-only, particles_d points to the same location as particles. If particles is resized, then Cabana will not deallocate the first allocation
239  // since it is still used by particles_d. So, reset particles_d before resize, and point it back to particles only afterwards
240  particles_d = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>();
241 #endif
242 
243  particles.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
244  particles.resize(add_vec_buffer(n_ptl));
245 
246 #ifndef USE_GPU
247  // Point particles_d back to particles after resize
249 #endif
250 
252  }
253 
255 #ifdef USE_GPU
256  particles.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
257  particles.resize(particles_d.size());
258 #else
259  // Point particles back to particles_d after resize
261 #endif
262 
264  }
265 
266  /* If using CPU-only, then "device" particles are a shallow copy of host particles so that
267  * there is no unnecessary duplication. When "device" particles are resized, Cabana will keep the
268  * original allocation if there is a second reference (i.e. host particles).
269  * To resolve this, we free the host particles here so that there is no second reference
270  * */
272  particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>();
273  }
274 
275  /* Resizes device particles if on GPU, or just creates a shallow copy if CPU only
276  * */
278  if(!owns_particles_d) exit_XGC("\nSpecies tried to resize device particles, but doesn't own the device array.");
279 
280 #ifdef USE_GPU
281  particles_d.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
282  // Resize device particles to match n particles
284 #else
285  // If kernels are on CPU, do shallow copy
287 #endif
288  }
289 
290  /* Resizes device particles
291  * */
292  void resize_device_particles(int new_n_ptl){
293  if(!owns_particles_d) exit_XGC("\nSpecies tried to resize device particles, but doesn't own the device array.");
294 
295  n_ptl = new_n_ptl;
296 
297  particles_d.reserve(minimum_ptl_reservation); // Can only raise reservation (no-op if AoSoA is already larger)
298 
299  // Resize device particles to match host particles
301  }
302 
303  /* Copies particles to device - deep copy if using GPU, otherwise shallow copy
304  * Also takes the opportunity to set the buffer particles to realistic values
305  * */
307  if(!owns_particles_d) exit_XGC("\nSpecies tried to copy particles to device, but doesn't own the device array.");
308 
309 #ifdef USE_GPU
310  // Copy to device
311  Cabana::deep_copy(particles_d, particles);
312 #else
313  // No operation required if CPU-only
314 #endif
315 
316  // Copy last particle to fill remainder of trailing vector in AoSoA
318  }
319 
320  /* Copies particles from device - deep copy if using GPU, otherwise no copy is necessary
321  * */
323  if(!owns_particles_d) exit_XGC("\nSpecies tried to copy particles from device, but doesn't own the device array.");
324 
325 #ifdef USE_GPU
326  // Copy particles to host
327  Cabana::deep_copy(particles, particles_d);
328 #else
329  // No operation required if CPU-only
330 #endif
331  }
332 
333  /* Copies particles to device if they are resident on the device
334  * */
337  }
338 
339  /* Copies particles from device if they are resident on the device
340  * */
343  }
344 
345  /* Copies particles to device if they are NOT resident on the device
346  * */
349  }
350 
351  /* Copies particles from device if they are NOT resident on the device
352  * */
355  }
356 
365  if (n_ptl>0){
366  int last_ptl_index = n_ptl - 1;
367  auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
368  auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
369  auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
370 #ifdef ESC_PTL
371  auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
372 #endif
373 
374  Kokkos::parallel_for("set_buffer_particles_d", Kokkos::RangePolicy<ExSpace>( n_ptl, add_vec_buffer(n_ptl) ), KOKKOS_LAMBDA( const int i ){
375  // Buffer particles: same as last particle, gid = -1
376  for (int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
377  for (int j=0;j<3;j++) ct(i, j) = ct(last_ptl_index, j);
378  gid(i) = -1;
379 #ifdef ESC_PTL
380  flag(i) = flag(last_ptl_index);
381 #endif
382  });
383  }
384  }
385 
394  if (n_ptl>0){
395  int last_ptl_index = n_ptl - 1;
396  auto ph = Cabana::slice<PtlSlice::Ph>(phase0_d);
397 
398  Kokkos::parallel_for("set_buffer_phase0", Kokkos::RangePolicy<ExSpace>( n_ptl, add_vec_buffer(n_ptl) ), KOKKOS_LAMBDA( const int i ){
399  // copy final real particle
400  for (int j=0;j<6;j++) ph(i, j) = ph(last_ptl_index, j);
401  });
402  }
403  }
404 
405  // Options for custom launch bounds since kokkos defaults are suboptimal for electron push kernel
406  enum class LaunchBounds{
407  Default,
408  Custom
409  };
410 
416  template<typename F>
417  inline void for_all_particles(const std::string label, F lambda_func) const {
418  Kokkos::RangePolicy<ExSpace> particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
419  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
420  }
421 
422  inline void back_up_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA, int offset, int n) const{
423  auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
424  auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
425  auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
426 #ifdef ESC_PTL
427  auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
428 #endif
429 
430  auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
431  auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
432  auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
433 #ifdef ESC_PTL
434  auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
435 #endif
436 
437  Kokkos::parallel_for("backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA( const int i ){
438  int i_offset = i + offset;
439  // Make backup copy
440  for (int j=0;j<6;j++) ph_b(i, j) = ph(i_offset, j);
441  for (int j=0;j<3;j++) ct_b(i, j) = ct(i_offset, j);
442  gid_b(i) = gid(i_offset);
443 #ifdef ESC_PTL
444  flag_b(i) = flag(i_offset);
445 #endif
446  // Deactivate
447  gid(i_offset) = -1;
448  });
449  }
450 
451  inline void restore_backup_SoA(Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>& backup_SoA, int offset, int n) const{
452  auto ph_b = Cabana::slice<PtlSlice::Ph>(backup_SoA);
453  auto ct_b = Cabana::slice<PtlSlice::Ct>(backup_SoA);
454  auto gid_b = Cabana::slice<PtlSlice::Gid>(backup_SoA);
455 #ifdef ESC_PTL
456  auto flag_b = Cabana::slice<PtlSlice::Flag>(backup_SoA);
457 #endif
458 
459  auto ph = Cabana::slice<PtlSlice::Ph>(particles_d);
460  auto ct = Cabana::slice<PtlSlice::Ct>(particles_d);
461  auto gid = Cabana::slice<PtlSlice::Gid>(particles_d);
462 #ifdef ESC_PTL
463  auto flag = Cabana::slice<PtlSlice::Flag>(particles_d);
464 #endif
465 
466  Kokkos::parallel_for("backup_first_soa", Kokkos::RangePolicy<ExSpace>( 0, n ), KOKKOS_LAMBDA( const int i ){
467  int i_offset = i + offset;
468  // Restore from backup copy
469  for (int j=0;j<6;j++) ph(i_offset, j) = ph_b(i, j);
470  for (int j=0;j<3;j++) ct(i_offset, j) = ct_b(i, j);
471  gid(i_offset) = gid_b(i);
472 #ifdef ESC_PTL
473  flag(i_offset) = flag_b(i);
474 #endif
475  });
476  }
477 
483  template<typename F>
484  inline void for_particle_range(int begin_idx, int end_idx, const std::string label, F lambda_func) const {
485  if(end_idx <= begin_idx) return; // Return if range is 0 or less
486 
487  // Still need the subset to line up with the AoSoA vector length
488  int first_soa = begin_idx/VEC_LEN;
489  int n_other_ptl_in_first_soa = begin_idx - first_soa*VEC_LEN;
490  bool first_soa_is_partial = (n_other_ptl_in_first_soa>0);
491  int last_soa = (end_idx-1)/VEC_LEN;
492  int n_other_ptl_in_last_soa = (last_soa+1)*VEC_LEN - end_idx;
493  bool last_soa_is_partial = (n_other_ptl_in_last_soa>0);
494 
495 #ifdef USE_GPU
496  int first_item_in_shifted_range = first_soa*VEC_LEN;
497 #else
498  int first_item_in_shifted_range = first_soa;
499 #endif
500 
501  Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_first_soa;
502  Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN> ptl_last_soa;
503  if(first_soa_is_partial){
504  // Make a backup of the first SoA and set the particles_d GIDs to -1
505  ptl_first_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>("ptl_first_soa", n_other_ptl_in_first_soa);
506  back_up_SoA(ptl_first_soa, first_soa*VEC_LEN, n_other_ptl_in_first_soa);
507  }
508  if(last_soa_is_partial){
509  // Make a backup of the last SoA and set the particles_d GIDs to -1
510  ptl_last_soa = Cabana::AoSoA<ParticleDataTypes,Device,VEC_LEN>("ptl_last_soa", n_other_ptl_in_last_soa);
511  back_up_SoA(ptl_last_soa, end_idx, n_other_ptl_in_last_soa);
512  }
513 
514  // Finally, do the parallel_for
515  Kokkos::RangePolicy<ExSpace> particle_range_policy( first_item_in_shifted_range, p_range<DeviceType>(end_idx) );
516  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
517 
518  // Restore from backup
519  if(first_soa_is_partial){
520  restore_backup_SoA(ptl_first_soa, first_soa*VEC_LEN, n_other_ptl_in_first_soa);
521  }
522  if(last_soa_is_partial){
523  restore_backup_SoA(ptl_last_soa, end_idx, n_other_ptl_in_last_soa);
524  }
525  }
526 
534  template<typename F>
535  inline void for_all_particles(const std::string label, F lambda_func,
536  const PtlMvmt mvmt, LaunchBounds launch_bounds=LaunchBounds::Default) {
537  if(!owns_particles_d) exit_XGC("\nSpecies tried to loop over particles on device, but doesn't own the device array.");
538 
539  bool use_streaming = stream_particles;
540 #ifndef USE_STREAMS
541  use_streaming = false; // Just to be safe, turn streams off here
542 #endif
543 
544  bool send_ptl = ( mvmt.send_opt==PtlMvmt::Send ||
546  bool return_ptl = ( mvmt.return_opt==PtlMvmt::Return ||
548 
549  // Don't need to stream if particles are already present and don't need to be returned
550  if((!send_ptl) && (!return_ptl)) use_streaming = false;
551 
552  if(use_streaming){
553 #ifdef USE_STREAMS
554  Streamed::Option stream_option = Streamed::Normal; // Send to device and back
555  if(!send_ptl) stream_option = Streamed::NoSend;
556  if(!return_ptl) stream_option = Streamed::NoReturn;
557 
558  // Execute streaming parallel_for
559  Streamed::parallel_for(label, n_ptl, lambda_func, stream_option, particles, particles_d);
560 #endif
561  }else{
562  if(send_ptl) TIMER("copy_ptl_to_device",copy_particles_to_device() );
563 
564  if (launch_bounds==LaunchBounds::Custom) {
565 #ifdef USE_EPUSH_LAUNCH_BOUNDS
566 # if !defined(PUSH_MAX_THREADS_PER_BLOCK) || !defined(PUSH_MIN_WARPS_PER_EU)
567 # error "USE_EPUSH_LAUNCH_BOUNDS requires PUSH_MAX_THREADS_PER_BLOCK and PUSH_MIN_WARPS_PER_EU to be defined"
568 # endif
569  Kokkos::RangePolicy<ExSpace, Kokkos::LaunchBounds<PUSH_MAX_THREADS_PER_BLOCK, PUSH_MIN_WARPS_PER_EU>>
570  particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
571  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
572 #else
573  exit_XGC("\nERROR: LaunchBounds::Custom specified, but USE_EPUSH_LAUNCH_BOUNDS is not defined\n");
574 #endif
575  } else {
576  Kokkos::RangePolicy<ExSpace>
577  particle_range_policy( 0, p_range<DeviceType>(n_ptl) );
578  Kokkos::parallel_for(label, Opt::require(particle_range_policy, Async), lambda_func);
579  }
580 
581  if(return_ptl) TIMER("copy_ptl_from_device", copy_particles_from_device() );
582  }
583  }
584 
585  KOKKOS_INLINE_FUNCTION VecParticles* ptl() const{
586  return (VecParticles*)(&particles_d.access(0));
587  }
588 
589  KOKKOS_INLINE_FUNCTION VecPhase* ph0() const{
590  return (VecPhase*)(&phase0_d.access(0));
591  }
592 
594  for_all_particles("copy_to_phase0", KOKKOS_LAMBDA( const int idx ){
595  AoSoAIndices<DeviceType> inds(idx);
596  VecParticles* ptl_loc = species.ptl();
597  VecPhase* ph0_loc = species.ph0();
598  for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
599  int p_vec = inds.a + i_simd;
600  ph0_loc[inds.s].r[p_vec] = ptl_loc[inds.s].ph.r[p_vec];
601  ph0_loc[inds.s].z[p_vec] = ptl_loc[inds.s].ph.z[p_vec];
602  ph0_loc[inds.s].phi[p_vec] = ptl_loc[inds.s].ph.phi[p_vec];
603  ph0_loc[inds.s].rho[p_vec] = ptl_loc[inds.s].ph.rho[p_vec];
604  ph0_loc[inds.s].w1[p_vec] = ptl_loc[inds.s].ph.w1[p_vec];
605  ph0_loc[inds.s].w2[p_vec] = ptl_loc[inds.s].ph.w2[p_vec];
606  }
607  });
608  }
609 
612  // If particles are resident on the device, then use the host particle allocation as the backup
613  // Otherwise, resize the backup_particle array
616  }else{
617  backup_particles.resize(particles_d.size());
618  }
619 
620  // Copy particle data to backup
621  Cabana::deep_copy(backup_particles, particles_d);
623  } else {
624  // For ions, save phase to phase0
625  phase0_d.resize(particles_d.size());
626  copy_to_phase0(*this); // Separate function to avoid implicit copy into lambda
627  }
629  }
630 
633  // If particles are resident on device, don't need to resize host particles since they are already used as the backup.
634  // If not, resize host particle array to fit backup particles
638  }else{
639  // If particles are not resident on device, resize host particle array to fit backup particles
641  }
642 
643  // Resize device particle array to fit backed up particles
645 
646  // Restore particle data from backup
647  Cabana::deep_copy(particles_d, backup_particles);
648 
650  // If the backup particles are simply pointing to the host particles, then
651  // point them to their own 0-sized allocation when finished using them
652  // so that they don't make a copy when host particles get resized
653  backup_particles = Cabana::AoSoA<ParticleDataTypes,HostType,VEC_LEN>("backup_particles", 0);
654  }
655  } else {
656  // When ipc==2, copy phase0 to device
657  // First, resize device phase0 and reset pointer
658  phase0_d.resize(phase0.size());
659 
660  // Next copy data to phase0 on device
661  Cabana::deep_copy(phase0_d, phase0);
662 
663  // Fill buffer with realistic ptl data
665  }
666  particles_are_backed_up = false;
667  }
668 
669  KOKKOS_INLINE_FUNCTION void restore_phase_from_phase0(const AoSoAIndices<Device>& inds, SimdParticles& part_one) const {
670  VecPhase* ph0_loc = ph0();
671  for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
672  int p_vec = inds.a + i_simd;
673  part_one.ph.r[i_simd] = ph0_loc[inds.s].r[p_vec];
674  part_one.ph.z[i_simd] = ph0_loc[inds.s].z[p_vec];
675  part_one.ph.phi[i_simd] = ph0_loc[inds.s].phi[p_vec];
676  part_one.ph.rho[i_simd] = ph0_loc[inds.s].rho[p_vec];
677  part_one.ph.w1[i_simd] = ph0_loc[inds.s].w1[p_vec];
678  part_one.ph.w2[i_simd] = ph0_loc[inds.s].w2[p_vec];
679  }
680  }
681 
682  long long int get_total_n_ptl(){
683 #ifdef USE_MPI
684  long long int tmp_n_ptl = n_ptl;
685  long long int out_n_ptl = 0;
686  MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_LONG_LONG_INT, MPI_SUM, SML_COMM_WORLD);
687  return out_n_ptl;
688 #else
689  return (long long int)(n_ptl);
690 #endif
691  }
692 
694 #ifdef USE_MPI
695  int tmp_n_ptl = n_ptl;
696  int out_n_ptl = 0;
697  MPI_Allreduce(&tmp_n_ptl, &out_n_ptl, 1, MPI_INT, MPI_MAX, SML_COMM_WORLD);
698  return out_n_ptl;
699 #else
700  return n_ptl;
701 #endif
702  }
703 
704  // Gets the gyro_radius of a species based on equilibrium temperature
705  // inode is the LOCAL (poloidally decomposed) grid node index to get temperature
706  // smu_n is the normalized sqrt(mu)
707  // bfield is the magnetic field at inode
708  KOKKOS_INLINE_FUNCTION double get_f0_eq_gyro_radius(int inode, double smu_n, double bfield) const{
709  // Should replace UNIT_CHARGE*charge_eu with charge(?)
710  return smu_n*sqrt(mass*f0.temp_ev(inode)*EV_2_J) / (UNIT_CHARGE*charge_eu*bfield);
711  }
712 
713  // Gets the equilibrium thermal velocity of a species based on f0 temperature
714  // inode is the GLOBAL node index to get temperature
715  KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity(int inode) const{
716  return thermal_velocity(mass, f0.temp_global(inode));
717  }
718 
719  // Gets the equilibrium thermal velocity of a species based on f0 temperature, on device
720  // inode is the local node index to get temperature
721  KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode(int inode) const{
722  return thermal_velocity(mass, f0.temp_ev(inode));
723  }
724 
725  // Gets the equilibrium thermal velocity of a species based on f0 temperature, on host
726  // inode is the local node index to get temperature
727  KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode_h(int inode) const{
728  return thermal_velocity(mass, f0.temp_ev_h(inode));
729  }
730 
731  // Get species velocity
732  KOKKOS_INLINE_FUNCTION void get_particle_velocity_and_nearest_node(const Grid<DeviceType>& grid, const MagneticField<DeviceType>& magnetic_field, const DomainDecomposition<DeviceType>& pol_decomp, SimdParticles& part, Simd<double>& smu, Simd<double>& vp, Simd<int>& nearest_node, Simd<bool>& not_in_triangle, Simd<bool>& not_in_poloidal_domain) const{
733 
734  // This modulo surely doesnt need to be here (at least, should be elsewhere).
735  // Modulo phi coordinate
736  grid.wedge_modulo_phi(part.ph.phi);
737 
739  grid.get_grid_weights(magnetic_field, part.ph.v(), grid_wts0);
740 
741  // Output argument
742  for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
743  not_in_triangle[i_simd] = !grid_wts0.is_valid(i_simd);
744  }
745 
746  Simd<double> bmag;
747  magnetic_field.bmag_interpol(part.ph.v(), bmag);
748 
749  for (int i_simd = 0; i_simd<SIMD_SIZE; i_simd++){
750  if(!grid_wts0.is_valid(i_simd)) continue;
751 
752  nearest_node[i_simd]=grid_wts0.node[i_simd] - pol_decomp.node_offset;
753  not_in_poloidal_domain[i_simd] = (nearest_node[i_simd]<0 || nearest_node[i_simd]>=pol_decomp.nnodes);
754 
755  double temp_ev_norm = not_in_poloidal_domain[i_simd] ? f0.temp_ev(0) : f0.temp_ev(nearest_node[i_simd]);
756 
757  // get vp and smu
758  const double& B = bmag[i_simd];
759  vp[i_simd] = normalized_v_para(c_m, mass, B, temp_ev_norm, part.ph.rho[i_simd]);
760  smu[i_simd] = normalized_sqrt_mu(B, temp_ev_norm, part.ct.mu[i_simd]);
761  }
762  }
763 
765  const Grid<DeviceType>& grid,
767  const VelocityGrid& vgrid);
768 
771 
772  void write_ptl_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name);
773  void write_f0_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name);
774  void read_f0_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name);
775  void read_ptl_checkpoint_files(const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, std::string sp_name, bool n_ranks_is_same);
776 
777  long long int get_max_gid() const;
778  void get_ptl_write_total_and_offsets(const DomainDecomposition<DeviceType>& pol_decomp, long long int& inum_total, long long int& ioff) const;
779 };
780 
781 #endif
void calculate_global_f0_arrays(const Grid< DeviceType > &grid, const MagneticField< DeviceType > &magnetic_field)
Definition: species.cpp:432
Cabana::AoSoA< PhaseDataTypes, HostType, VEC_LEN > phase0
Definition: species.hpp:112
bool stream_particles
Whether to stream particles between host and device if possible.
Definition: species.hpp:104
Definition: globals.hpp:84
KOKKOS_INLINE_FUNCTION VecPhase * ph0() const
Definition: species.hpp:589
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:101
void back_up_SoA(Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > &backup_SoA, int offset, int n) const
Definition: species.hpp:422
KOKKOS_INLINE_FUNCTION double normalized_v_para(double c_m, double mass, double B, double temp_ev, double rho)
Definition: basic_physics.hpp:80
KOKKOS_INLINE_FUNCTION double get_f0_eq_gyro_radius(int inode, double smu_n, double bfield) const
Definition: species.hpp:708
subroutine adjust_n_ptl_for_core_ptl(n_ptl)
Definition: load.F90:473
void set_spall_num_and_ptr(int idx, int n_ptl, int n_vecs, VecParticles *ptl)
void for_particle_range(int begin_idx, int end_idx, const std::string label, F lambda_func) const
Definition: species.hpp:484
constexpr double PROTON_MASS
Definition: constants.hpp:7
Distribution< Device > f0
Species distribution in velocity space on local mesh nodes.
Definition: species.hpp:120
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:116
bool is_electron
Whether this species is the electrons.
Definition: species.hpp:79
Eq::Profile< Device > eq_mk_temp
Definition: species.hpp:134
void for_all_particles(const std::string label, F lambda_func, const PtlMvmt mvmt, LaunchBounds launch_bounds=LaunchBounds::Default)
Definition: species.hpp:535
void save_backup_particles()
Definition: species.hpp:610
void update_decomposed_f0_calculations(const DomainDecomposition< DeviceType > &pol_decomp, const Grid< DeviceType > &grid, const MagneticField< DeviceType > &magnetic_field, const VelocityGrid &vgrid)
Definition: species.cpp:385
double c2_2m
c2/2m
Definition: species.hpp:87
double rho[VEC_LEN]
Definition: particles.hpp:96
void copy_to_phase0(Species< Device > &species)
Definition: species.hpp:593
Definition: species.hpp:58
Simd< double > w1
Definition: particles.hpp:22
double c_m
c/m
Definition: species.hpp:86
Definition: species.hpp:57
constexpr double EV_2_J
Conversion rate ev to J.
Definition: constants.hpp:5
Definition: velocity_grid.hpp:8
bool default_streaming_option()
Definition: species.hpp:28
Eq::Profile< Device > eq_den
Definition: species.hpp:125
Definition: globals.hpp:89
KOKKOS_INLINE_FUNCTION VecParticles * ptl() const
Definition: species.hpp:585
Definition: grid_weights.hpp:47
KOKKOS_INLINE_FUNCTION double thermal_velocity(double mass, double temp_ev)
Definition: basic_physics.hpp:58
Definition: NamelistReader.hpp:193
KinType kintype
Whether the species is gyrokinetic or drift kinetic.
Definition: species.hpp:82
Eq::Profile< Device > eq_fg_temp
Definition: species.hpp:129
Definition: magnetic_field.hpp:12
void get_ptl_write_total_and_offsets(const DomainDecomposition< DeviceType > &pol_decomp, long long int &inum_total, long long int &ioff) const
Definition: species.cpp:515
int add_vec_buffer(int n_ptl)
Definition: particles.hpp:194
int idx
Index in all_species.
Definition: species.hpp:78
Definition: particles.hpp:92
int a
The index in the inner array of the AoSoA.
Definition: particles.hpp:150
Definition: particles.hpp:109
KOKKOS_INLINE_FUNCTION void bmag_interpol(const SimdVector &v, Simd< double > &bmag) const
Definition: magnetic_field.tpp:239
bool particles_are_backed_up
Whether particles are currently backed up.
Definition: species.hpp:109
int nonadiabatic_idx
Index of species skipping adiabatic species (for compatibility with fortran arrays) ...
Definition: species.hpp:81
bool default_residence_option()
Definition: species.hpp:35
int n_ptl
Number of particles.
Definition: species.hpp:95
Eq::Profile< Device > eq_mk_flow
Definition: species.hpp:136
KOKKOS_INLINE_FUNCTION void get_particle_velocity_and_nearest_node(const Grid< DeviceType > &grid, const MagneticField< DeviceType > &magnetic_field, const DomainDecomposition< DeviceType > &pol_decomp, SimdParticles &part, Simd< double > &smu, Simd< double > &vp, Simd< int > &nearest_node, Simd< bool > &not_in_triangle, Simd< bool > &not_in_poloidal_domain) const
Definition: species.hpp:732
Definition: streamed_parallel_for.hpp:16
int node_offset
Offset of first mesh node belonging to this MPI rank.
Definition: domain_decomposition.hpp:55
Definition: streamed_parallel_for.hpp:14
Eq::Profile< Device > eq_fg_flow
Definition: species.hpp:131
void set_buffer_phase0_d()
Definition: species.hpp:393
long long int get_total_n_ptl()
Definition: species.hpp:682
void set_buffer_particles_d()
Definition: species.hpp:364
Simd< double > rho
Definition: particles.hpp:21
Definition: species.hpp:59
int p_range< DeviceType >(int num_particle)
Definition: particles.hpp:187
int eq_flow_type
Definition: species.hpp:127
double charge_eu
Particle charge in eu.
Definition: species.hpp:85
Definition: species.hpp:51
int nnodes
Number of nodes belonging to this MPI rank.
Definition: domain_decomposition.hpp:56
void resize_particles(int new_n_ptl)
Definition: species.hpp:234
double mass
Particle mass.
Definition: species.hpp:83
KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity(int inode) const
Definition: species.hpp:715
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.cpp:25
void write_ptl_checkpoint_files(const DomainDecomposition< DeviceType > &pol_decomp, const XGC_IO_Stream &stream, std::string sp_name)
Definition: species.cpp:540
void for_all_particles(const std::string label, F lambda_func) const
Definition: species.hpp:417
Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > particles_d
Particles on device.
Definition: species.hpp:99
Definition: species.hpp:48
double w2[VEC_LEN]
Definition: particles.hpp:98
#define TIMER(N, F)
Definition: timer_macro.hpp:24
Eq::Profile< Device > eq_mk_den
Definition: species.hpp:135
RKRestorationMethod
Definition: species.hpp:68
idx
Definition: diag_f0_df_port1.hpp:32
void copy_particles_to_device_if_not_resident()
Definition: species.hpp:347
void read_ptl_checkpoint_files(const DomainDecomposition< DeviceType > &pol_decomp, const XGC_IO_Stream &stream, std::string sp_name, bool n_ranks_is_same)
Definition: species.cpp:590
RKRestorationMethod RK_restoration_method
Currently, electrons must use first method and ions must use second.
Definition: species.hpp:107
Simd< double > r
Definition: particles.hpp:18
void resize_host_particles_to_match_device()
Definition: species.hpp:254
Definition: species.hpp:69
KOKKOS_INLINE_FUNCTION SimdVector & v()
Definition: particles.hpp:39
ReturnOpt return_opt
Definition: species.hpp:63
ReturnOpt
Definition: species.hpp:56
Option
Definition: streamed_parallel_for.hpp:13
void restore_particles_from_backup()
Definition: species.hpp:631
Definition: globals.hpp:90
static std::vector< MemoryPrediction > estimate_memory_usage(NLReader::NamelistReader &nlr, const Grid< DeviceType > &grid, const DomainDecomposition< DeviceType > &pol_decomp, int species_idx)
Definition: species.cpp:56
SendOpt send_opt
Definition: species.hpp:62
double charge
Particle charge.
Definition: species.hpp:84
SimdPhase ph
Definition: particles.hpp:59
void copy_particles_from_device()
Definition: species.hpp:322
KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode(int inode) const
Definition: species.hpp:721
void copy_particles_from_device_if_not_resident()
Definition: species.hpp:353
KOKKOS_INLINE_FUNCTION void wedge_modulo_phi(Simd< double > &phi_mod) const
Definition: grid.tpp:790
void unassign_host_particles()
Definition: species.hpp:271
int ncycles_between_sorts
Number of subcycles between sorts.
Definition: species.hpp:92
Definition: particles.hpp:58
Cabana::AoSoA< PhaseDataTypes, Device, VEC_LEN > phase0_d
Definition: species.hpp:113
int SML_COMM_RANK
Definition: my_mpi.cpp:5
KinType
Definition: globals.hpp:88
Species(SpeciesType sp_type, int n_ptl)
Definition: species.hpp:150
bool is_deltaf
Whether this species is deltaf.
Definition: species.hpp:89
Definition: xgc_io.hpp:19
VecPhase ph
Definition: particles.hpp:110
Definition: species.hpp:44
Definition: species.hpp:70
void set_min_max_num(int isp, int n_ptl)
int minimum_ptl_reservation
The minimum reservation size for particles.
Definition: species.hpp:94
int s
The index in the outer array of the AoSoA.
Definition: particles.hpp:149
KOKKOS_INLINE_FUNCTION double normalized_sqrt_mu(double B, double temp_ev, double mu)
Definition: basic_physics.hpp:90
Simd< double > z
Definition: particles.hpp:19
void copy_particles_to_device_if_resident()
Definition: species.hpp:335
Definition: species.hpp:52
void resize_device_particles(int new_n_ptl)
Definition: species.hpp:292
Definition: species.hpp:53
constexpr double UNIT_CHARGE
Charge of an electron (C)
Definition: constants.hpp:4
long long int get_max_gid() const
Definition: species.cpp:492
KOKKOS_INLINE_FUNCTION double get_f0_eq_thermal_velocity_lnode_h(int inode) const
Definition: species.hpp:727
void exit_XGC(std::string msg)
Definition: globals.hpp:37
void copy_particles_from_device_if_resident()
Definition: species.hpp:341
bool is_adiabatic
Whether this species is adiabatic.
Definition: species.hpp:80
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:82
int n_backup_particles
Definition: species.hpp:117
Eq::Profile< Device > eq_flow
Definition: species.hpp:126
Definition: streamed_parallel_for.hpp:15
SendOpt
Definition: species.hpp:50
void read_f0_checkpoint_files(const DomainDecomposition< DeviceType > &pol_decomp, const XGC_IO_Stream &stream, std::string sp_name)
Definition: species.cpp:685
SimdConstants ct
Definition: particles.hpp:60
int eq_fg_flow_type
Definition: species.hpp:132
void write_f0_checkpoint_files(const DomainDecomposition< DeviceType > &pol_decomp, const XGC_IO_Stream &stream, std::string sp_name)
Definition: species.cpp:649
void copy_particles_to_device()
Definition: species.hpp:306
KOKKOS_INLINE_FUNCTION void restore_phase_from_phase0(const AoSoAIndices< Device > &inds, SimdParticles &part_one) const
Definition: species.hpp:669
Species(int n_ptl_in)
Definition: species.hpp:180
double phi[VEC_LEN]
Definition: particles.hpp:95
Simd< double > w2
Definition: particles.hpp:23
double r[VEC_LEN]
Definition: particles.hpp:93
GyroAverageMatrices< HostType > gyro_avg_matrices
Definition: species.hpp:140
Definition: species.hpp:75
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
KOKKOS_INLINE_FUNCTION void get_grid_weights(const MagneticField< Device > &magnetic_field, const SimdVector &v, const Simd< double > &psi, SimdVector2D &xff, SimdGridWeights< Order::One, PIT > &grid_wts) const
Definition: grid.tpp:705
Definition: species.hpp:44
Eq::Profile< Device > eq_temp
Definition: species.hpp:124
bool particles_resident_on_device
Whether the particles can reside on device.
Definition: species.hpp:103
Simd< double > mu
Definition: particles.hpp:52
Eq::Profile< Device > eq_fg_den
Definition: species.hpp:130
PtlMvmt(SendOpt send_opt, ReturnOpt return_opt)
Definition: species.hpp:65
static int get_initial_n_ptl(NLReader::NamelistReader &nlr, const Grid< DeviceType > &grid, const DomainDecomposition< DeviceType > &pol_decomp, int sml_special, int species_idx, bool verbose)
Definition: species.cpp:113
int ncycles
Number of subcycles.
Definition: species.hpp:91
int eq_mk_flow_type
Definition: species.hpp:137
Definition: profile.hpp:171
int collision_grid_index
Which collision grid to use.
Definition: species.hpp:122
Definition: particles.hpp:148
SpeciesType
Definition: globals.hpp:83
double z[VEC_LEN]
Definition: particles.hpp:94
void restore_backup_SoA(Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > &backup_SoA, int offset, int n) const
Definition: species.hpp:451
Definition: distribution.hpp:10
void resize_device_particles()
Definition: species.hpp:277
int get_max_n_ptl()
Definition: species.hpp:693
LaunchBounds
Definition: species.hpp:406
Cabana::AoSoA< ParticleDataTypes, HostType, VEC_LEN > particles
Particles.
Definition: species.hpp:96
double w1[VEC_LEN]
Definition: particles.hpp:97
Definition: species.hpp:44