XGC1
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
vgrid_distribution.hpp
Go to the documentation of this file.
1 #ifndef VGRID_DISTRIBUTION_HPP
2 #define VGRID_DISTRIBUTION_HPP
3 
4 #include "space_settings.hpp"
5 #include "velocity_grid.hpp"
6 #include "view_arithmetic.hpp"
8 #include "maxwellian.hpp"
9 
10 template<class Device>
12  public:
13 
14  using exec_space = typename Device::execution_space;
15 
16  View<double****,CLayout, Device> f;
17 
18  // Create a 1D unmanaged view looking at the same address
19  View<double*,CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_1D;
20 
21  // Make f private when possible - ALS
22  //public:
23 
24  double vp_max;
25  double dvp;
26  double smu_max;
27  double dsmu;
29 
31 
32  // Analytical distribution for testing
33  VGridDistribution(const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp, const std::vector<Maxwellian>& maxwellians)
34  : f("f", maxwellians.size(), vgrid.nvr, pol_decomp.nnodes, vgrid.nvz),
35  f_1D(f.data(), f.size()),
36  vp_max(vgrid.vp_max),
37  dvp(vgrid.dvp),
38  smu_max(vgrid.smu_max),
39  dsmu(vgrid.dsmu),
41  {
42  for (int i=0; i<f.extent(0); i++){
43  for (int imu=0; imu<f.extent(1); imu++){
44  double vperp = imu*dsmu;
45  for (int inode=0; inode<f.extent(2); inode++){
46  for (int ivp=0; ivp<f.extent(3); ivp++){
47  double vpara = ivp*dvp - vp_max;
48  f(i,imu,inode,ivp) = maxwellians[i].get_f(vpara, vperp);
49  }
50  }
51  }
52  }
53  }
54 
55  VGridDistribution(int nsp, const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp)
56  : f("f", nsp, vgrid.nvr, pol_decomp.nnodes, vgrid.nvz),
57  f_1D(f.data(), f.size()),
58  vp_max(vgrid.vp_max),
59  dvp(vgrid.dvp),
60  smu_max(vgrid.smu_max),
61  dsmu(vgrid.dsmu),
63  {}
64 
65  // () operator: behave as if accessing the View directly
66  KOKKOS_INLINE_FUNCTION double& operator()(int isp, int ivr, int inode, int ivz) const { return f(isp, ivr, inode, ivz);}
67 
68  // [] operator: access the data as if it has been squashed to 1 dimension
69  KOKKOS_INLINE_FUNCTION double& operator[](int i) const { return f_1D(i);}
70 
71  double* data() const {return f.data();}
72 
78  template<typename F>
79  inline void for_all_elements(const std::string label, F lambda_func) const {
80  Kokkos::parallel_for(label, Kokkos::RangePolicy<exec_space>(0, f.size()), lambda_func);
81  }
82 
88  template<typename F>
89  inline void for_each_element(const std::string label, F lambda_func) const {
90  Kokkos::parallel_for(label, Kokkos::MDRangePolicy<Kokkos::Rank<4>, exec_space>({0, 0, 0, 0}, {f.extent(0), f.extent(1), f.extent(2), f.extent(3)}), lambda_func);
91  }
92 
93 
94  // Scatter value onto vgrid distribution
95  KOKKOS_INLINE_FUNCTION void scatter(int i_node, const VGridWeights& wt, double value) const{
96  constexpr int ISP_ZERO = 0;
97  Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
98  Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
99  Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
100  Kokkos::atomic_add(&(f(ISP_ZERO, wt.i_vr+1,i_node,wt.i_vz+1)),value*wt.w_11);
101  }
102 
103  // Temporary scatter for views not yet converted to VGridDistribution
104  KOKKOS_INLINE_FUNCTION static void scatter(const View<double***,CLayout,Device>& view, int i_node, const VGridWeights& wt, double value){
105  Kokkos::atomic_add(&(view(wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
106  Kokkos::atomic_add(&(view(wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
107  Kokkos::atomic_add(&(view(wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
108  Kokkos::atomic_add(&(view(wt.i_vr+1,i_node,wt.i_vz+1)),value*wt.w_11);
109  }
110 
111  // Temporary scatter for views not yet converted to VGridDistribution
112  KOKKOS_INLINE_FUNCTION static void scatter(const View<double****,CLayout,Device>& view, int i_node, const VGridWeights& wt, double value){
113  constexpr int ISP_ZERO = 0;
114  Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
115  Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
116  Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
117  Kokkos::atomic_add(&(view(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+1)),value*wt.w_11);
118  }
119 
120  // Gather value from vgrid distribution
121  KOKKOS_INLINE_FUNCTION double gather(int i_node, const VGridWeights& wt) const{
122  constexpr int ISP_ZERO = 0;
123  return ( f(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00 +
124  f(ISP_ZERO,wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01 +
125  f(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10 +
126  f(ISP_ZERO,wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11 );
127  }
128 
129  // Temporary gather for views not yet converted to VGridDistribution
130  template<class T>
131  KOKKOS_INLINE_FUNCTION static double gather(const T& view, int i_node, const VGridWeights& wt){
132  return ( view(wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00 +
133  view(wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01 +
134  view(wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10 +
135  view(wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11 );
136  }
137 
138  // Temporary gather for views not yet converted to VGridDistribution
139  // Divides by normalization view
140  template<class T>
141  KOKKOS_INLINE_FUNCTION static double normed_gather(const T& view, int iw, int i_node, const VGridWeights& wt, const T& norm_view){
142  return (view(iw,wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00/norm_view(iw,wt.i_vr+0,i_node,wt.i_vz+0) +
143  view(iw,wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01/norm_view(iw,wt.i_vr+0,i_node,wt.i_vz+1) +
144  view(iw,wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10/norm_view(iw,wt.i_vr+1,i_node,wt.i_vz+0) +
145  view(iw,wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11/norm_view(iw,wt.i_vr+1,i_node,wt.i_vz+1));
146  }
147 
148  // Access dimensions
149  KOKKOS_INLINE_FUNCTION int n_species() const{
150  return f.extent(0);
151  }
152 
153  KOKKOS_INLINE_FUNCTION int n_vr() const {
154  return f.extent(1);
155  }
156 
157  KOKKOS_INLINE_FUNCTION int n_nodes() const{
158  return f.extent(2);
159  }
160 
161  KOKKOS_INLINE_FUNCTION int n_vz() const{
162  return f.extent(3);
163  }
164 
165  KOKKOS_INLINE_FUNCTION int size() const{
166  return f.size();
167  }
168 
169  // Miscellaneous
170  inline double get_smu_n(int imu) const{
171  return (imu>0 ? dsmu*imu : dsmu*inv_mu0_factor);
172  }
173 
175  return UniformRange(n_vr(), 0.0, smu_max);
176  }
177 
179  return UniformRange(n_vz(), -vp_max, vp_max);
180  }
181 };
182 
183 #endif
KOKKOS_INLINE_FUNCTION int n_vr() const
Definition: vgrid_distribution.hpp:153
double * data() const
Definition: vgrid_distribution.hpp:71
KOKKOS_INLINE_FUNCTION int size() const
Definition: vgrid_distribution.hpp:165
UniformRange vz_range() const
Definition: vgrid_distribution.hpp:178
double w_00
Definition: vgrid_weights.hpp:10
Definition: velocity_grid.hpp:8
double get_smu_n(int imu) const
Definition: vgrid_distribution.hpp:170
UniformRange vr_range() const
Definition: vgrid_distribution.hpp:174
VGridDistribution(const VelocityGrid &vgrid, const DomainDecomposition< DeviceType > &pol_decomp, const std::vector< Maxwellian > &maxwellians)
Definition: vgrid_distribution.hpp:33
KOKKOS_INLINE_FUNCTION double gather(int i_node, const VGridWeights &wt) const
Definition: vgrid_distribution.hpp:121
double smu_max
max mu
Definition: vgrid_distribution.hpp:26
static KOKKOS_INLINE_FUNCTION double gather(const T &view, int i_node, const VGridWeights &wt)
Definition: vgrid_distribution.hpp:131
double dvp
grid spacing in parallel velocity
Definition: vgrid_distribution.hpp:25
double w_01
Definition: vgrid_weights.hpp:11
View< double ****, CLayout, Device > f
Definition: vgrid_distribution.hpp:16
KOKKOS_INLINE_FUNCTION int n_species() const
Definition: vgrid_distribution.hpp:149
KOKKOS_INLINE_FUNCTION double & operator()(int isp, int ivr, int inode, int ivz) const
Definition: vgrid_distribution.hpp:66
static KOKKOS_INLINE_FUNCTION void scatter(const View< double ***, CLayout, Device > &view, int i_node, const VGridWeights &wt, double value)
Definition: vgrid_distribution.hpp:104
double inv_mu0_factor
Definition: vgrid_distribution.hpp:28
double w_11
Definition: vgrid_weights.hpp:13
VGridDistribution()
Definition: vgrid_distribution.hpp:30
KOKKOS_INLINE_FUNCTION int n_vz() const
Definition: vgrid_distribution.hpp:161
View< double *, CLayout, Device, Kokkos::MemoryTraits< Kokkos::Unmanaged > > f_1D
Definition: vgrid_distribution.hpp:19
VGridDistribution(int nsp, const VelocityGrid &vgrid, const DomainDecomposition< DeviceType > &pol_decomp)
Definition: vgrid_distribution.hpp:55
static KOKKOS_INLINE_FUNCTION void scatter(const View< double ****, CLayout, Device > &view, int i_node, const VGridWeights &wt, double value)
Definition: vgrid_distribution.hpp:112
typename HostType::execution_space exec_space
Definition: vgrid_distribution.hpp:14
Definition: vgrid_distribution.hpp:11
int i_vr
Definition: vgrid_weights.hpp:8
int i_vz
Definition: vgrid_weights.hpp:9
double dsmu
grid spacing in mu
Definition: vgrid_distribution.hpp:27
void for_each_element(const std::string label, F lambda_func) const
Definition: vgrid_distribution.hpp:89
Definition: vgrid_weights.hpp:7
Definition: uniform_range.hpp:6
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 int n_nodes() const
Definition: vgrid_distribution.hpp:157
void for_all_elements(const std::string label, F lambda_func) const
Definition: vgrid_distribution.hpp:79
KOKKOS_INLINE_FUNCTION double & operator[](int i) const
Definition: vgrid_distribution.hpp:69
double w_10
Definition: vgrid_weights.hpp:12
static KOKKOS_INLINE_FUNCTION double normed_gather(const T &view, int iw, int i_node, const VGridWeights &wt, const T &norm_view)
Definition: vgrid_distribution.hpp:141
double vp_max
max parallel velocity
Definition: vgrid_distribution.hpp:24
KOKKOS_INLINE_FUNCTION void scatter(int i_node, const VGridWeights &wt, double value) const
Definition: vgrid_distribution.hpp:95