XGCa
 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"
7 #include "velocity_grid.hpp"
9 #include "linear_weights.hpp"
10 
11 // Struct containing indices and weights on velocity grid
12 struct VGridWeights{
13  int i_vr;
14  int i_vz;
15  double w_00;
16  double w_01;
17  double w_10;
18  double w_11;
19 
20  KOKKOS_INLINE_FUNCTION VGridWeights(const LinearWeights& vr_wt, const LinearWeights& vz_wt)
21  : i_vr(vr_wt.i),
22  i_vz(vz_wt.i),
23  w_00(vr_wt.w[0]*vz_wt.w[0]),
24  w_01(vr_wt.w[0]*vz_wt.w[1]),
25  w_10(vr_wt.w[1]*vz_wt.w[0]),
26  w_11(vr_wt.w[1]*vz_wt.w[1])
27  {}
28 };
29 
30 template<class Device>
32  public:
33 
34  using exec_space = typename Device::execution_space;
35 
36  View<double****,CLayout, Device> f;
37 
38  // Create a 1D unmanaged view looking at the same address
39  View<double*,CLayout, Device, Kokkos::MemoryTraits<Kokkos::Unmanaged>> f_1D;
40 
41  // Make f private when possible - ALS
42  //public:
43 
44  double vp_max;
45  double dvp;
46  double smu_max;
47  double dsmu;
49 
51 
52  VGridDistribution(int nsp, const VelocityGrid& vgrid, const DomainDecomposition<DeviceType>& pol_decomp)
53  : f("f", nsp, vgrid.nvr, pol_decomp.nnodes, vgrid.nvz),
54  f_1D(f.data(), f.size()),
55  vp_max(vgrid.vp_max),
56  dvp(vgrid.dvp),
57  smu_max(vgrid.smu_max),
58  dsmu(vgrid.dsmu),
60  {}
61 
62  // () operator: behave as if accessing the View directly
63  KOKKOS_INLINE_FUNCTION double& operator()(int isp, int ivr, int inode, int ivz) const { return f(isp, ivr, inode, ivz);}
64 
65  // [] operator: access the data as if it has been squashed to 1 dimension
66  KOKKOS_INLINE_FUNCTION double& operator[](int i) const { return f_1D(i);}
67 
68  double* data(){return f.data();}
69 
75  template<typename F>
76  inline void for_all_elements(const std::string label, F lambda_func) const {
77  Kokkos::parallel_for(label, Kokkos::RangePolicy<exec_space>(0, f.size()), lambda_func);
78  }
79 
85  template<typename F>
86  inline void for_each_element(const std::string label, F lambda_func) const {
87  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);
88  }
89 
90 
91  // Determine vgrid weights from input smu, vp and vgrid_vol
92  KOKKOS_INLINE_FUNCTION VGridWeights get_weights(double smu, double vp, double inv_vgrid_vol) const{
93  // Transformation from mu to vr:
94  const double& vr = smu; // square root mu for sqrt(mu) grid, v_perp/v_th else
95 
96  // Transformation from vp to vz:
97  double vz = vp + vp_max; // Shift so that value is (probably) between 0 and 2*vp_max
98 
99  // Calculate linear weights for velocity grid
100  LinearWeights vr_wt(vr, 1.0/dsmu);
101  LinearWeights vz_wt(vz, 1.0/dvp);
102 
103  //rh For f0_f: conversion from mu-vpar to sqrt(mu)-vpar or vperp-vpar
104  vr_wt.w[0] *= inv_vgrid_vol * (vr_wt.i==0 ? 2.0 : 1.0);
105  vr_wt.w[1] *= inv_vgrid_vol * (vr_wt.i==n_vr()-2 ? 2.0 : 1.0);
106 
107  return VGridWeights(vr_wt, vz_wt);
108  }
109 
110  // Scatter value onto vgrid distribution
111  KOKKOS_INLINE_FUNCTION void scatter(int i_node, const VGridWeights& wt, double value) const{
112  Kokkos::atomic_add(&(f(wt.i_vr+0,i_node,wt.i_vz+0)),value*wt.w_00);
113  Kokkos::atomic_add(&(f(wt.i_vr+0,i_node,wt.i_vz+1)),value*wt.w_01);
114  Kokkos::atomic_add(&(f(wt.i_vr+1,i_node,wt.i_vz+0)),value*wt.w_10);
115  Kokkos::atomic_add(&(f(wt.i_vr+1,i_node,wt.i_vz+1)),value*wt.w_11);
116  }
117 
118  // Gather value from vgrid distribution
119  KOKKOS_INLINE_FUNCTION double gather(int i_node, const VGridWeights& wt) const{
120  return ( f(wt.i_vr+0,i_node,wt.i_vz+0)*wt.w_00 +
121  f(wt.i_vr+0,i_node,wt.i_vz+1)*wt.w_01 +
122  f(wt.i_vr+1,i_node,wt.i_vz+0)*wt.w_10 +
123  f(wt.i_vr+1,i_node,wt.i_vz+1)*wt.w_11 );
124  }
125 
126  // Access dimensions
127  KOKKOS_INLINE_FUNCTION int n_species() const{
128  return f.extent(0);
129  }
130 
131  KOKKOS_INLINE_FUNCTION int n_vr() const {
132  return f.extent(1);
133  }
134 
135  KOKKOS_INLINE_FUNCTION int n_nodes() const{
136  return f.extent(2);
137  }
138 
139  KOKKOS_INLINE_FUNCTION int n_vz() const{
140  return f.extent(3);
141  }
142 
143  KOKKOS_INLINE_FUNCTION int size() const{
144  return f.size();
145  }
146 
147  // Miscellaneous
148  inline double get_smu_n(int imu) const{
149  return (imu>0 ? dsmu*imu : dsmu*inv_mu0_factor);
150  }
151 };
152 
153 #endif
KOKKOS_INLINE_FUNCTION int n_vr() const
Definition: vgrid_distribution.hpp:131
KOKKOS_INLINE_FUNCTION int size() const
Definition: vgrid_distribution.hpp:143
double w_00
Definition: vgrid_distribution.hpp:15
Definition: linear_weights.hpp:7
Definition: velocity_grid.hpp:7
double get_smu_n(int imu) const
Definition: vgrid_distribution.hpp:148
KOKKOS_INLINE_FUNCTION double gather(int i_node, const VGridWeights &wt) const
Definition: vgrid_distribution.hpp:119
double smu_max
max mu
Definition: vgrid_distribution.hpp:46
KOKKOS_INLINE_FUNCTION VGridWeights get_weights(double smu, double vp, double inv_vgrid_vol) const
Definition: vgrid_distribution.hpp:92
double dvp
grid spacing in parallel velocity
Definition: vgrid_distribution.hpp:45
double w_01
Definition: vgrid_distribution.hpp:16
View< double ****, CLayout, Device > f
Definition: vgrid_distribution.hpp:36
KOKKOS_INLINE_FUNCTION int n_species() const
Definition: vgrid_distribution.hpp:127
KOKKOS_INLINE_FUNCTION double & operator()(int isp, int ivr, int inode, int ivz) const
Definition: vgrid_distribution.hpp:63
double inv_mu0_factor
Definition: vgrid_distribution.hpp:48
double w_11
Definition: vgrid_distribution.hpp:18
VGridDistribution()
Definition: vgrid_distribution.hpp:50
KOKKOS_INLINE_FUNCTION int n_vz() const
Definition: vgrid_distribution.hpp:139
View< double *, CLayout, Device, Kokkos::MemoryTraits< Kokkos::Unmanaged > > f_1D
Definition: vgrid_distribution.hpp:39
VGridDistribution(int nsp, const VelocityGrid &vgrid, const DomainDecomposition< DeviceType > &pol_decomp)
Definition: vgrid_distribution.hpp:52
double w[2]
Definition: linear_weights.hpp:9
typename HostType::execution_space exec_space
Definition: vgrid_distribution.hpp:34
Definition: vgrid_distribution.hpp:31
int i
Definition: linear_weights.hpp:8
int i_vr
Definition: vgrid_distribution.hpp:13
int i_vz
Definition: vgrid_distribution.hpp:14
double dsmu
grid spacing in mu
Definition: vgrid_distribution.hpp:47
void for_each_element(const std::string label, F lambda_func) const
Definition: vgrid_distribution.hpp:86
Definition: vgrid_distribution.hpp:12
KOKKOS_INLINE_FUNCTION VGridWeights(const LinearWeights &vr_wt, const LinearWeights &vz_wt)
Definition: vgrid_distribution.hpp:20
double * data()
Definition: vgrid_distribution.hpp:68
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:135
void for_all_elements(const std::string label, F lambda_func) const
Definition: vgrid_distribution.hpp:76
KOKKOS_INLINE_FUNCTION double & operator[](int i) const
Definition: vgrid_distribution.hpp:66
double w_10
Definition: vgrid_distribution.hpp:17
double vp_max
max parallel velocity
Definition: vgrid_distribution.hpp:44
KOKKOS_INLINE_FUNCTION void scatter(int i_node, const VGridWeights &wt, double value) const
Definition: vgrid_distribution.hpp:111