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