XGC1
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
poly_basis.hpp
Go to the documentation of this file.
1 #ifndef POLYNOMIAL_BASIS_HPP
2 #define POLYNOMIAL_BASIS_HPP
3 
4 #include "uniform_range.hpp"
5 #include "vgrid_distribution.hpp"
6 
7 /* Struct to store a polynomial representation of a distribution function */
8 template<class Device>
10  View<double***, CLayout, Device> f;
11 
12  PolynomialBasisDistribution(int n_poly, int n)
13  : f("f", n_poly, n_poly, n){} // Initialize to 0.0
14 };
15 
16 /* Class for storing a set of basis functions for a 2D grid, as well as functions for converting a distribution
17  * function to and from a polynomial representation */
18 template<class Device>
20  using exec_space = typename Device::execution_space;
21 
22  int n_poly;
25  View<double***, CLayout, Device> bases;
26  View<double***, CLayout, HostType> bases_h; // Bases construct on host for now
27  View<double****, CLayout, Device> matrix;
28 
29  // Making these functions public for the moment, due to:
30  // The enclosing parent function ("compute_inner_product") for an extended __host__ __device__ lambda cannot have private or protected access within its class
31  public:
32 
73  View<double***, CLayout, Device> compute_inner_product(const View<double***, CLayout, Device>& dist_in) const{
74  int nnode = dist_in.extent(0);
75  View<double***, CLayout, Device> inner_prod_matrix("inner_prod_matrix", n_poly, n_poly, nnode);
76 
77  const PolynomialBasis pb = *this;
78 
79  Kokkos::parallel_for("inner_product", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
80  for (int iy = 0; iy < pb.yrange.n; iy++) {
81  for (int ix = 0; ix < pb.xrange.n; ix++) {
82  double w = dist_in(inode, ix, iy);
83 
84  for(int ip = 0; ip<pb.n_poly; ip++){
85  for(int jp=ip; jp<pb.n_poly; jp++){
86  inner_prod_matrix(ip, jp, inode) += w*pb.matrix(ip, jp, ix, iy);
87  }
88  }
89  }
90  }
91 
92  // Reflect to create symmetric matrix
93  for(int ip = 0; ip<pb.n_poly; ip++){
94  for(int jp=ip+1; jp<pb.n_poly; jp++){
95  inner_prod_matrix(jp, ip, inode) = inner_prod_matrix(ip, jp, inode);
96  }
97  }
98 
99  // Normalize
100  double den = inner_prod_matrix(0,0,inode);
101  for(int i=0; i<pb.n_poly; i++){
102  for(int j=0; j<pb.n_poly; j++){
103  inner_prod_matrix(i,j,inode) /= den;
104  }
105  }
106  });
107  Kokkos::fence();
108  return inner_prod_matrix;
109  }
110 
115  PolynomialBasisDistribution<Device> get_coefficients(const View<double***, CLayout, Device>& inner_prod_matrix) const{
116  int nnode = inner_prod_matrix.extent(2);
118 
119  View<double**, CLayout, Device> norm("norm", n_poly, nnode); // Initialize to 0.0;
120 
121  const PolynomialBasis pb = *this;
122  Kokkos::parallel_for("get_coeffs", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
123  // The first one is trivial
124  polynom.f(0, 0, inode) = 1.0;
125  for (int j=0; j<pb.n_poly; j++) {
126  double sum_value = 0.0;
127  for (int k = 0; k < pb.n_poly; k++) {
128  sum_value += polynom.f(k, 0, inode) * inner_prod_matrix(k, j, inode);
129  }
130  norm(0, inode) += polynom.f(j, 0, inode) * sum_value;
131  }
132 
133  // Now the rest
134  for (int i = 1; i<pb.n_poly; i++){
135  polynom.f(i, i, inode) = 1.0;
136  for (int j = i-1; j >= 0; j--) {
137  double sum_value = 0.0;
138  for (int k = 0; k < pb.n_poly; ++k) {
139  sum_value += inner_prod_matrix(k, i, inode) * polynom.f(k, j, inode);
140  }
141  for (int l = 0; l < pb.n_poly; ++l) {
142  polynom.f(l, i, inode) -= sum_value / norm(j, inode) * polynom.f(l, j, inode);
143  }
144  }
145  for (int j = 0; j < pb.n_poly; j++) {
146  double sum_value = 0.0;
147  for (int k = 0; k < pb.n_poly; ++k) {
148  sum_value += polynom.f(k, i, inode) * inner_prod_matrix(k, j, inode);
149  }
150  norm(i, inode) += polynom.f(j, i, inode) * sum_value;
151  }
152  }
153 
154  for (int i = 0; i < pb.n_poly; i++) {
155  double norm_sqrt = sqrt(norm(i, inode));
156  for (int j = 0; j < pb.n_poly; j++) {
157  polynom.f(j, i, inode) /= norm_sqrt;
158  }
159  }
160  });
161  Kokkos::fence();
162  return polynom;
163  }
164 
170  template<typename F>
171  void add_basis(int& basis_ind, F func){
172  // This should be small - keep on host for now
173  for(int i=0; i<xrange.n; i++){
174  double x = xrange.get_val(i);
175  for(int j=0; j<yrange.n; j++){
176  double y = yrange.get_val(j);
177  bases_h(basis_ind, i, j) = func(x, y);
178  }
179  }
180 
181  basis_ind += 1; // Move to next basis function
182  }
183 
184  public:
185 
191  template<typename... Fs>
192  PolynomialBasis(const UniformRange& xrange_in, const UniformRange& yrange_in, Fs... funcs)
193  : xrange(xrange_in),
194  yrange(yrange_in),
195  n_poly(sizeof...(Fs)), // sizeof counts number of entries when applied to a parameter pack
196  bases(NoInit("bases"), n_poly, xrange.n, yrange.n),
197  bases_h(NoInit("bases_h"), bases.layout()),
198  matrix(NoInit("matrix"), n_poly, n_poly, xrange.n, yrange.n)
199  {
200  // Step through the parameter pack of functions, populating the individual basis functions
201  int basis_ind = 0;
202  (..., add_basis(basis_ind, funcs));
203 
204  // Compute the matrix - small so keep on host for now
205  View<double****, CLayout, HostType> matrix_h(NoInit("matrix_h"), matrix.layout());
206  for(int ip=0; ip<n_poly; ip++){
207  for(int jp=ip; jp<n_poly; jp++){
208  for(int i=0; i<xrange.n; i++){
209  for(int j=0; j<yrange.n; j++){
210  matrix_h(ip, jp, i, j) = bases_h(ip, i, j)*bases_h(jp, i, j);
211  }
212  }
213  }
214  }
215  Kokkos::deep_copy(bases, bases_h);
216  Kokkos::deep_copy(matrix, matrix_h);
217  }
218 
223  PolynomialBasisDistribution<Device> get_distribution(const View<double***, CLayout, Device>& dist_in) const{
224  auto inner_prod_matrix = compute_inner_product(dist_in);
225 
226  return get_coefficients(inner_prod_matrix);
227  }
228 
234  template<typename F>
235  inline View<double**, CLayout, Device> get_moments(const int nnode, F lambda_func) const{
236  const View<double**, CLayout, Device> g_mom(NoInit("g_mom"), n_poly, nnode);
237 
238  Kokkos::parallel_for("get_moments", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
239  lambda_func(inode, g_mom);
240  });
241  Kokkos::fence();
242  return g_mom;
243  }
244 
245  /* Computes the scaling polynomials with the precomputed basis
246  * polynomials and the new moments (n,v_||,mu,v_||^2)
247  *
248  * Ansatz: g = sum_[i=1]^4 d_i lambda_i f
249  * ==> d_i = int(lambda_i g d^3v) / <lambda_i,lambda_i>
250  * = sum_[j=1]^4 c_[i,j] int(gamma_j g d^3v) / <lambda_i,lambda_i>
251  * After some manipulation --->
252  * g = f sum_[k=1]^4 D_k gamma_k with
253  * D_k = sum_[i,j=1]^4 int(gamma_j g d^3v * c_[i,j]*c_[i,k] / <lambda_i,lambda_i>
254  * @param[in] polynom Polynomial representation of the distribution
255  * @param[in] g_mom Array of new moments
256  * return The scaling
257  */
258  View<double***, CLayout, Device> get_scaling(const PolynomialBasisDistribution<Device>& polynom, const View<double**, CLayout, Device>& g_mom) const{
259  int nnode = polynom.f.extent(2);
260  View<double***, CLayout, Device> scaling("scaling", nnode, xrange.n, yrange.n);
261  View<double**, CLayout, Device> D_k("D_k", n_poly, nnode);
262 
263  const PolynomialBasis pb = *this;
264  Kokkos::parallel_for("get_scaling", Kokkos::RangePolicy<exec_space>(0,nnode), KOKKOS_LAMBDA( const int inode ){
265  // 1) Compute the polynomial coefficients D_k
266  for (int i = 0; i < pb.n_poly; i++) {
267  for (int j = 0; j < pb.n_poly; j++) {
268  for (int k = 0; k < pb.n_poly; k++) {
269  double c2 = polynom.f(j,i,inode)*polynom.f(k,i,inode);
270  D_k(k,inode) += g_mom(j,inode)*c2;
271  }
272  }
273  }
274 
275  // 2) Compute the new distribution function using the final
276  // scaling polynomial P = sum_[k=1]^4 D_k gamma_k
277  for (int iy = 0; iy < pb.yrange.n; iy++) {
278  for (int ix = 0; ix < pb.xrange.n; ix++) {
279  double fac = -1.0;
280  for(int ip=0; ip<pb.n_poly; ip++){
281  fac += D_k(ip,inode)*pb.bases(ip,ix,iy);
282  }
283  scaling(inode, ix, iy) = fac;
284  }
285  }
286  });
287  Kokkos::fence();
288  return scaling;
289  }
290 };
291 
292 #endif
PolynomialBasisDistribution< Device > get_distribution(const View< double ***, CLayout, Device > &dist_in) const
Definition: poly_basis.hpp:223
PolynomialBasis(const UniformRange &xrange_in, const UniformRange &yrange_in, Fs...funcs)
Definition: poly_basis.hpp:192
View< double ***, CLayout, Device > f
Definition: poly_basis.hpp:10
View< double ***, CLayout, Device > get_scaling(const PolynomialBasisDistribution< Device > &polynom, const View< double **, CLayout, Device > &g_mom) const
Definition: poly_basis.hpp:258
typename Device::execution_space exec_space
Definition: poly_basis.hpp:20
View< double ***, CLayout, Device > bases
Definition: poly_basis.hpp:25
KOKKOS_INLINE_FUNCTION double get_val(int i) const
Definition: uniform_range.hpp:21
PolynomialBasisDistribution(int n_poly, int n)
Definition: poly_basis.hpp:12
View< double **, CLayout, Device > get_moments(const int nnode, F lambda_func) const
Definition: poly_basis.hpp:235
int n_poly
Definition: poly_basis.hpp:22
PolynomialBasisDistribution< Device > get_coefficients(const View< double ***, CLayout, Device > &inner_prod_matrix) const
Definition: poly_basis.hpp:115
void add_basis(int &basis_ind, F func)
Definition: poly_basis.hpp:171
View< double ***, CLayout, Device > compute_inner_product(const View< double ***, CLayout, Device > &dist_in) const
Definition: poly_basis.hpp:73
Definition: poly_basis.hpp:19
Definition: poly_basis.hpp:9
UniformRange yrange
Definition: poly_basis.hpp:24
Definition: uniform_range.hpp:6
UniformRange xrange
Definition: poly_basis.hpp:23
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
View< double ****, CLayout, Device > matrix
Definition: poly_basis.hpp:27
Kokkos::ViewAllocateWithoutInitializing NoInit
Definition: space_settings.hpp:68
View< double ***, CLayout, HostType > bases_h
Definition: poly_basis.hpp:26
int n
Definition: uniform_range.hpp:7