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