XGCa
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
plasma.hpp
Go to the documentation of this file.
1 #ifndef PLASMA_HPP
2 #define PLASMA_HPP
3 #include "my_subview.hpp"
4 #include "species.hpp"
5 #include "vgrid_distribution.hpp"
6 
7 extern "C" void f0_init_decomposed_ptrs(double* f0_T_ev_cpp, double* f0_fg_T_ev_cpp, double* f0_inv_grid_vol_cpp, double* f0_grid_vol_cpp,
8  double* f0_grid_vol_vonly_cpp, double* f0_den_cpp,
9  double* f0_flow_cpp);
10 extern "C" void f0_set_ptrs(int nnode, double* f0_delta_n_cpp, double* f0_delta_u_cpp, double* f0_delta_T_cpp);
11 extern "C" void set_f0_f0g_ptr(int f0_inode1, int f0_inode2, double* f0_f0g_loc);
12 
13 class Plasma{
14 
15  // Variables for managing the device particle memory allocation
19 
20 
21 #ifdef DELTAF_CONV
22  static constexpr bool reduced_deltaf = true;
23 #else
24  static constexpr bool reduced_deltaf = false;
25 #endif
27 
28  public:
29 
30  std::vector<Species<DeviceType>> all_species;
31 
32  // Poloidally decomposed f0 values
33  // These will probably become a member of species later, but the fortran code expects
34  // the species to be an array index
36  View<double*,CLayout,HostType> f0_node_cost;
37 
38  private:
39 
40  /* Contains f0 values that are poloidally decomposed but don't need to be transferred between ranks during load rebalancing
41  * */
43  View<double**,CLayout, HostType> f0_T_ev;
44  View<double**,CLayout, HostType> f0_fg_T_ev;
45  View<double**,CLayout, HostType> f0_fg_vth_inv;
46  View<double**,CLayout, HostType> f0_inv_grid_vol;
47  View<double**,CLayout, HostType> f0_grid_vol;
48  View<double**,CLayout, HostType> f0_grid_vol_vonly;
49  View<double**,CLayout, HostType> f0_den;
50  View<double**,CLayout, HostType> f0_flow;
51 
52 
54 
56  : f0_T_ev("f0_T_ev", nsp, pol_decomp.nnodes),
57  f0_fg_T_ev("f0_fg_T_ev", nsp, pol_decomp.nnodes),
58  f0_fg_vth_inv("f0_fg_vth_inv", nsp, pol_decomp.nnodes),
59  f0_inv_grid_vol("f0_inv_grid_vol", nsp, pol_decomp.nnodes),
60  f0_grid_vol("f0_grid_vol", nsp, pol_decomp.nnodes),
61  f0_grid_vol_vonly("f0_grid_vol_vonly", nsp, pol_decomp.nnodes),
62  f0_den("f0_den", nsp, pol_decomp.nnodes),
63  f0_flow("f0_flow", nsp, pol_decomp.nnodes)
64  {
65 #ifndef NO_FORTRAN_MODULES
66  // Fortran arrays are set to point to these Views
68  f0_fg_T_ev.data(),
69  f0_inv_grid_vol.data(),
70  f0_grid_vol.data(),
71  f0_grid_vol_vonly.data(),
72  f0_den.data(),
73  f0_flow.data());
74 #endif
75  }
76  };
77 
78 
79  public:
80 
82 
83  // Not poloidally decomposed
84  View<double**,CLayout, HostType> f0_delta_n;
85  View<double**,CLayout, HostType> f0_delta_u;
86  View<double**,CLayout, HostType> f0_delta_T;
87  View<double*,CLayout, HostType> spitzer_resistivity; // Pre-calculated for equilibrium density and temperature
88 
89  int nspecies;
91 
92  std::vector<std::string> sp_names;
93 
94  // Constructors
96  : particles_d_has_owner(false),
98  nspecies(0), // Initialize with 0 species
99  n_nonadiabatic_species(0), // Initialize with 0 species
100  sp_names{"e", "i", "i2", "i3", "i4", "i5", "i6"}
101  {}
102 
104 
108  };
109 
113  };
114 
115  // Loop over all species
116  template<typename F>
117  inline void for_all_species(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
118  for(int isp = 0; isp<all_species.size(); isp++){
119  manage_particle_ownership(isp, device_ptl_opt);
120  func(all_species[isp]);
121  }
122  }
123 
124  // Loop over all non-adiabatic species
125  template<typename F>
126  inline void for_all_nonadiabatic_species(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
127  for(int isp = 0; isp<all_species.size(); isp++){
128  if(isp>10) { // ISP ERROR
129  printf("ISP ERROR in for_all_nonadiabatic_sepcies: isp=%d",isp);
130  fflush(stdout);
131  }
132  if(!all_species[isp].is_adiabatic){
133  manage_particle_ownership(isp, device_ptl_opt);
134  func(all_species[isp]);
135  }
136  }
137  }
138 
139  // Loop over electrons
140  template<typename F>
141  inline void for_electrons(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
142  for(int isp = 0; isp<all_species.size(); isp++){
143  if(all_species[isp].is_electron){
144  manage_particle_ownership(isp, device_ptl_opt);
145  func(all_species[isp]);
146  }
147  }
148  }
149 
150  // Loop over ions
151  template<typename F>
152  inline void for_all_ions(F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
153  for(int isp = 0; isp<all_species.size(); isp++){
154  if(!all_species[isp].is_electron){
155  manage_particle_ownership(isp, device_ptl_opt);
156  func(all_species[isp]);
157  }
158  }
159  }
160 
161  // Loop over electrons or ions as specified
162  template<typename F>
163  inline void for_all(ParticleType particle_type, F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
164  for(int isp = 0; isp<all_species.size(); isp++){
165  if((particle_type==Electrons && all_species[isp].is_electron) ||
166  (particle_type==Ions && !all_species[isp].is_electron)){
167  manage_particle_ownership(isp, device_ptl_opt);
168  func(all_species[isp]);
169  }
170  }
171  }
172 
173  // Operate on one species
174  template<typename F>
175  inline void for_one_species(int isp, F func, DevicePtlOpt device_ptl_opt = UseDevicePtl){
176  manage_particle_ownership(isp, device_ptl_opt);
177  func(all_species[isp]);
178  }
179 
180  // Loop over all species, return true if the input condition is met for any of them
181  template<typename F>
182  inline bool true_for_some_species(F func) const{
183  for(int isp = 0; isp<all_species.size(); isp++){
184  if(func(all_species[isp])) return true;
185  }
186  return false;
187  }
188 
189  template<typename F>
190  inline bool true_for_all_species(F func) const{
191  for(int isp = 0; isp<all_species.size(); isp++){
192  if(!func(all_species[isp])) return false;
193  }
194  return true;
195  }
196 
197  // Replace !sml_electron_on
198  inline bool electrons_are_adiabatic() const{
199  return all_species[ELECTRON].is_adiabatic;
200  }
201 
202  inline bool f0_grid() const{
203  return true_for_some_species([&](const Species<DeviceType>& s){return s.dynamic_f0;});
204  }
205 
206  int largest_n_ptl(bool check_backup){
207  int max_n_ptl = 0;
208 
209  // Loop over all species
211  if(species.is_electron && check_backup){
212  // Check on backup particles instead, if that's what's getting used
213  max_n_ptl = std::max(max_n_ptl,species.n_backup_particles);
214  } else {
215  // For now, need to access fortran object for n_ptl
216  max_n_ptl = std::max(max_n_ptl,species.n_ptl);
217  }
218  }, NoDevicePtl);
219 
220  return max_n_ptl;
221  }
222 
223  /* Deallocates the device particles and resets ownership tracking variables
224  * */
226  for_all_nonadiabatic_species([&](Species<DeviceType>& species){
227  if(species.owns_particles_d){
228  species.particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>();
229  species.owns_particles_d = false;
230  }
231  });
232  particles_d_has_owner = false;
233  }
234 
235  static std::vector<MemoryPrediction> estimate_memory_usage(NLReader::NamelistReader& nlr, const Grid<DeviceType> &grid, const DomainDecomposition<DeviceType>& pol_decomp);
236 
237  /* Reallocate and recalculate arrays that are decomposed but recalculated rather than
238  * transferred by the load balance (f0_T_ev etc) */
239  void update_decomposed_f0_calculations(const DomainDecomposition<DeviceType>& pol_decomp,
240  const Grid<DeviceType>& grid,
242  const VelocityGrid& vgrid);
243 
244  private:
245 
246  void initialize_spitzer_res();
247 
248  /* Allocate and calculate global f0 arrays. Should be part of a constructor */
250  // Allocate for all species (even though most simulations don't need it for adiabatic species)
251  f0_delta_n = View<double**,CLayout, HostType>("f0_delta_n", nspecies, grid.nnode);
252  f0_delta_u = View<double**,CLayout, HostType>("f0_delta_u", nspecies, grid.nnode);
253  f0_delta_T = View<double**,CLayout, HostType>("f0_delta_T", nspecies, grid.nnode);
254  // Set all_species unmanaged views
255  for_all_species([&](Species<DeviceType>& species){
256  set_unmanaged_f0_species_view(f0_delta_n, species.idx, species.f0.delta_n_h);
257  set_unmanaged_f0_species_view(f0_delta_u, species.idx, species.f0.delta_u_h);
258  set_unmanaged_f0_species_view(f0_delta_T, species.idx, species.f0.delta_T_h);
259  }, NoDevicePtl);
260 #ifndef NO_FORTRAN_MODULES
261  // Set fortran pointers
262  f0_set_ptrs(grid.nnode, f0_delta_n.data(), f0_delta_u.data(), f0_delta_T.data());
263 #endif
264 
265  // These are constant for the simulation:
266  for_all_species([&](Species<DeviceType>& species){
267  species.initialize_global_f0_arrays(grid, magnetic_field);
268  }, NoDevicePtl);
269 
270  initialize_spitzer_res();
271  }
272 
273  // Creates an unmanaged view pointing to the subview specified by isp
274  template<typename T_in, typename T_out>
275  void set_unmanaged_f0_species_view(const T_in& view_in, int isp, T_out& view_out){
276  auto view_in_subview = my_subview(view_in, isp);
277  view_out = T_out(view_in_subview.data(), view_in_subview.layout());
278  }
279 
280  public:
281 
282  void resize_f0_f0g(const DomainDecomposition<DeviceType>& pol_decomp, const VelocityGrid& vgrid){
283  // Create new object rather than resizing since we don't need to preserve data
284  f0_f0g = VGridDistribution<HostType>(n_nonadiabatic_species, vgrid, pol_decomp);
285 #ifndef NO_FORTRAN_MODULES
286  int f0_inode1 = pol_decomp.node_offset + 1; // 1-indexed
287  int f0_inode2 = f0_inode1 + pol_decomp.nnodes - 1; // 1-indexed
288  set_f0_f0g_ptr(f0_inode1, f0_inode2, f0_f0g.data());
289 #endif
290 
291  // Set all_species unmanaged views
292  int f0_species_cnt=0;
293  for_all_nonadiabatic_species([&](Species<DeviceType>& species){
294  set_unmanaged_f0_species_view(f0_f0g.f, f0_species_cnt, species.f0.f0g_h);
295  f0_species_cnt++;
296  }, NoDevicePtl);
297  }
298 
299  private:
300 
302  if(!particles_d_has_owner){
303  // Set new owner and initialize with 0 particles
304  particles_d_owner = isp;
305 
306  all_species[particles_d_owner].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>("particles_d", 0);
307  all_species[particles_d_owner].owns_particles_d = true;
308  particles_d_has_owner = true;
309  }else{
310  // No-op if species is handing off to itself
311  if(particles_d_owner == isp) return;
312 
313  // Shallow copy to new owner
314  all_species[isp].particles_d = all_species[particles_d_owner].particles_d;
315 
316  // Delete original
317  all_species[particles_d_owner].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>();
318  all_species[particles_d_owner].owns_particles_d = false;
319 
320  // Set new owner
321  particles_d_owner = isp;
322  all_species[particles_d_owner].owns_particles_d = true;
323  }
324  }
325 
326  void manage_particle_ownership(int isp, DevicePtlOpt device_ptl_opt){
327  if((!all_species[isp].is_adiabatic) && device_ptl_opt==UseDevicePtl){
328  if(species_share_particles_d_ownership){
329  transfer_particles_d_ownership(isp);
330  }else{
331  if(!all_species[isp].owns_particles_d){
332  all_species[isp].particles_d = Cabana::AoSoA<ParticleDataTypes,DeviceType,VEC_LEN>("particles_d", 0);
333  all_species[isp].owns_particles_d = true;
334  }
335  }
336 
337  // Resize device particles if they are used
338  all_species[isp].resize_device_particles();
339  }
340  }
341 
342  public:
343  void validate_f0_checkpoint_file_dims(const XGC_IO_Stream& stream, const Grid<DeviceType>& grid, const VelocityGrid& vgrid);
344  void write_checkpoint_files(const Grid<DeviceType>& grid, const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream);
345  void read_checkpoint_files(const Grid<DeviceType>& grid, const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp, const XGC_IO_Stream& stream, const XGC_IO_Stream& f0_stream, bool n_ranks_is_same, int version);
346 
348  double main_ion_mass = all_species[MAIN_ION].mass;
349  double axis_length = TWOPI * magnetic_field.equil.axis.r;
350  double characteristic_velocity = sqrt(2*main_ion_characteristic_energy/main_ion_mass);
351  return axis_length / characteristic_velocity;
352  }
353 };
354 
355 #endif
View< double *, CLayout, HostType > spitzer_resistivity
Definition: plasma.hpp:87
Definition: globals.hpp:84
void for_all_ions(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:152
bool owns_particles_d
Whether the species owns the device particle allocation right now.
Definition: species.hpp:108
double * data() const
Definition: vgrid_distribution.hpp:73
void init_global_f0_arrays(const Grid< DeviceType > &grid, const MagneticField< DeviceType > &magnetic_field)
Definition: plasma.hpp:249
Distribution< Device > f0
Species distribution in velocity space on local mesh nodes.
Definition: species.hpp:127
void manage_particle_ownership(int isp, DevicePtlOpt device_ptl_opt)
Definition: plasma.hpp:326
bool is_electron
Whether this species is the electrons.
Definition: species.hpp:79
void f0_set_ptrs(int nnode, double *f0_delta_n_cpp, double *f0_delta_u_cpp, double *f0_delta_T_cpp)
static constexpr bool reduced_deltaf
Equivalent to the preprocessor flag for now.
Definition: plasma.hpp:24
ParticleType
Definition: plasma.hpp:110
bool dynamic_f0
Whether f0 can evolve in time.
Definition: species.hpp:96
void for_one_species(int isp, F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:175
bool f0_grid() const
Definition: plasma.hpp:202
Definition: velocity_grid.hpp:8
bool particles_d_has_owner
Whether a species owns the device particles allocation.
Definition: plasma.hpp:17
void deallocate_device_ptl()
Definition: plasma.hpp:225
Definition: plasma.hpp:107
View< double **, CLayout, HostType > f0_fg_vth_inv
thermal velocity from f_grid norm. temp. at nodes (to avoid sqrt ops)
Definition: plasma.hpp:45
Definition: NamelistReader.hpp:193
View< double **, CLayout, HostType > f0_fg_T_ev
f_Grid normalization temperature at nodes
Definition: plasma.hpp:44
Definition: magnetic_field.hpp:14
DecomposedRecalculableF0Arrays decomposed_recalculable_f0_arrays
Contains f0 values that are poloidally decomposed but don&#39;t need to be transferred between ranks duri...
Definition: plasma.hpp:81
int idx
Index in all_species.
Definition: species.hpp:78
VGridDistribution< HostType > f0_f0g
Definition: plasma.hpp:35
View< double **, CLayout, HostType > f0_flow
Equilibrium flow at nodes.
Definition: plasma.hpp:50
void for_all_nonadiabatic_species(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:126
View< double **, CLayout, HostType > f0_grid_vol_vonly
Grid volume (v only) at nodes.
Definition: plasma.hpp:48
Equilibrium equil
The object containing information about the magnetic equilibrium.
Definition: magnetic_field.hpp:46
View< double **, CLayout, HostType > f0_den
Equilibrium density at nodes.
Definition: plasma.hpp:49
bool default_residence_option()
Definition: species.hpp:35
int n_ptl
Number of particles.
Definition: species.hpp:102
int node_offset
Offset of first mesh node belonging to this MPI rank.
Definition: domain_decomposition.hpp:59
std::vector< Species< DeviceType > > all_species
Every particle species in the simulation.
Definition: plasma.hpp:30
void resize_f0_f0g(const DomainDecomposition< DeviceType > &pol_decomp, const VelocityGrid &vgrid)
Definition: plasma.hpp:282
int nnodes
Number of nodes belonging to this MPI rank.
Definition: domain_decomposition.hpp:60
void initialize_global_f0_arrays(const Grid< DeviceType > &grid, const MagneticField< DeviceType > &magnetic_field)
Definition: species.cpp:494
Definition: plasma.hpp:111
Cabana::AoSoA< ParticleDataTypes, Device, VEC_LEN > particles_d
Particles on device.
Definition: species.hpp:106
View< double **, CLayout, HostType > f0_T_ev
Equilibrium temperature at nodes.
Definition: plasma.hpp:43
bool true_for_some_species(F func) const
Definition: plasma.hpp:182
void set_unmanaged_f0_species_view(const T_in &view_in, int isp, T_out &view_out)
Definition: plasma.hpp:275
int particles_d_owner
Which species, if any, owns the device particles allocation.
Definition: plasma.hpp:18
void f0_init_decomposed_ptrs(double *f0_T_ev_cpp, double *f0_fg_T_ev_cpp, double *f0_inv_grid_vol_cpp, double *f0_grid_vol_cpp, double *f0_grid_vol_vonly_cpp, double *f0_den_cpp, double *f0_flow_cpp)
double get_main_ion_toroidal_transit_time(const MagneticField< DeviceType > &magnetic_field) const
Definition: plasma.hpp:347
int largest_n_ptl(bool check_backup)
Definition: plasma.hpp:206
DevicePtlOpt
Definition: plasma.hpp:105
View< double **, CLayout, HostType > f0_delta_T
Flux-surface averaged change of temperature.
Definition: plasma.hpp:86
int nspecies
Number of species including electrons.
Definition: plasma.hpp:89
void for_all_species(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:117
int n_nonadiabatic_species
Number of nonadiabatic species.
Definition: plasma.hpp:90
View< double **, CLayout, HostType > f0_inv_grid_vol
Inverse grid volume at nodes.
Definition: plasma.hpp:46
void set_f0_f0g_ptr(int f0_inode1, int f0_inode2, double *f0_f0g_loc)
Definition: xgc_io.hpp:24
Definition: plasma.hpp:106
DecomposedRecalculableF0Arrays()
Definition: plasma.hpp:53
double r
Definition: grid_structs.hpp:29
Definition: globals.hpp:85
Kokkos::View< T *, Kokkos::LayoutRight, Device > my_subview(const Kokkos::View< T ****, Kokkos::LayoutRight, Device > &view, int i, int j, int k)
Definition: my_subview.hpp:8
Definition: magnetic_field.F90:1
int n_backup_particles
Definition: species.hpp:124
void for_electrons(F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:141
View< double **, CLayout, HostType > f0_delta_n
Flux-surface averaged change of density.
Definition: plasma.hpp:84
Definition: plasma.hpp:13
Definition: plasma.hpp:112
void for_all(ParticleType particle_type, F func, DevicePtlOpt device_ptl_opt=UseDevicePtl)
Definition: plasma.hpp:163
bool species_share_particles_d_ownership
Whether to use the device particles sharing scheme.
Definition: plasma.hpp:16
View< double *, CLayout, HostType > f0_node_cost
Definition: plasma.hpp:36
Plasma()
Definition: plasma.hpp:95
Definition: species.hpp:75
void transfer_particles_d_ownership(int isp)
Definition: plasma.hpp:301
View< double **, CLayout, HostType > f0_grid_vol
Grid volume at nodes.
Definition: plasma.hpp:47
View< double **, CLayout, HostType > f0_delta_u
Flux-surface averaged change of parallel flow.
Definition: plasma.hpp:85
std::vector< std::string > sp_names
Definition: plasma.hpp:92
double main_ion_characteristic_energy
Definition: plasma.hpp:26
int nnode
Number of grid nodes.
Definition: grid.hpp:156
RZPair axis
Definition: equil.hpp:86
DecomposedRecalculableF0Arrays(int nsp, const DomainDecomposition< DeviceType > &pol_decomp)
Definition: plasma.hpp:55
bool electrons_are_adiabatic() const
Definition: plasma.hpp:198
bool true_for_all_species(F func) const
Definition: plasma.hpp:190
constexpr double TWOPI
Definition: constants.hpp:9