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