Skip to content

Commit f0a4736

Browse files
StaticObserverStaticObserver
committed
update custom output to pgen accretion
1 parent 0f3b6b4 commit f0a4736

1 file changed

Lines changed: 358 additions & 0 deletions

File tree

‎pgens/accretion/pgen.hpp‎

Lines changed: 358 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,180 @@
1919
namespace user {
2020
using namespace ntt;
2121

22+
enum CustomField : uint8_t {
23+
DB = 0,
24+
Gamma = 1,
25+
V = 2,
26+
Ut = 3,
27+
};
28+
29+
template <class M, CustomField F>
30+
class CustomMoments_kernel{
31+
static_assert(M::is_metric, "M must be a metric class");
32+
static constexpr auto D = M::Dim;
33+
34+
scatter_ndfield_t<D, 6> Buff;
35+
const idx_t buff_idx;
36+
const array_t<int*> i1, i2, i3;
37+
const array_t<prtldx_t*> dx1, dx2, dx3;
38+
const array_t<real_t*> ux1, ux2, ux3;
39+
const array_t<real_t*> phi;
40+
const array_t<real_t*> weight;
41+
const array_t<short*> tag;
42+
const float mass;
43+
const float charge;
44+
const M metric;
45+
const int ni2;
46+
const unsigned short window;
47+
const unsigned short comp;
48+
49+
const real_t smooth;
50+
bool is_axis_i2min { false }, is_axis_i2max { false };
51+
52+
public:
53+
CustomMoments_kernel( const unsigned short comp,
54+
const scatter_ndfield_t<D, 6>& scatter_buff,
55+
idx_t buff_idx,
56+
const array_t<int*>& i1,
57+
const array_t<int*>& i2,
58+
const array_t<int*>& i3,
59+
const array_t<prtldx_t*>& dx1,
60+
const array_t<prtldx_t*>& dx2,
61+
const array_t<prtldx_t*>& dx3,
62+
const array_t<real_t*>& ux1,
63+
const array_t<real_t*>& ux2,
64+
const array_t<real_t*>& ux3,
65+
const array_t<real_t*>& phi,
66+
const array_t<real_t*>& weight,
67+
const array_t<short*>& tag,
68+
float mass,
69+
float charge,
70+
const M& metric,
71+
const boundaries_t<FldsBC>& boundaries,
72+
ncells_t ni2,
73+
real_t inv_n0,
74+
unsigned short window)
75+
: comp { comp }
76+
, Buff { scatter_buff }
77+
, buff_idx { buff_idx }
78+
, i1 { i1 }
79+
, i2 { i2 }
80+
, i3 { i3 }
81+
, dx1 { dx1 }
82+
, dx2 { dx2 }
83+
, dx3 { dx3 }
84+
, ux1 { ux1 }
85+
, ux2 { ux2 }
86+
, ux3 { ux3 }
87+
, phi { phi }
88+
, weight { weight }
89+
, tag { tag }
90+
, mass { mass }
91+
, charge { charge }
92+
, metric { metric }
93+
, ni2 { static_cast<int>(ni2) }
94+
, window { window }
95+
, smooth { inv_n0 / (real_t)(math::pow(TWO * (real_t)window + ONE,
96+
static_cast<int>(D))) } {
97+
raise::ErrorIf(buff_idx >= 6, "Invalid buffer index", HERE);
98+
raise::ErrorIf(window > N_GHOSTS, "Window size too large", HERE);
99+
100+
raise::ErrorIf(comp > 2 || comp < 0, "Invalid component index", HERE);
101+
102+
raise::ErrorIf(D != Dim::_2D, "CustomMoments_kernel only supports 2D", HERE);
103+
raise::ErrorIf(M::CoordType != Coord::Qsph, "CustomMoments_kernel only supports Qspherical coordinates", HERE);
104+
105+
raise::ErrorIf(boundaries.size() < 2, "boundaries defined incorrectly", HERE);
106+
is_axis_i2min = (boundaries[1].first == FldsBC::AXIS);
107+
is_axis_i2max = (boundaries[1].second == FldsBC::AXIS);
108+
}
109+
110+
Inline void operator()(index_t p) const {
111+
if (tag(p) == ParticleTag::dead) {
112+
return;
113+
}
114+
real_t coeff { ZERO };
115+
if constexpr (F == CustomField::Gamma){
116+
117+
coord_t<D> x_Code { ZERO };
118+
x_Code[0] = static_cast<real_t>(i1(p)) + static_cast<real_t>(dx1(p));
119+
x_Code[1] = static_cast<real_t>(i2(p)) + static_cast<real_t>(dx2(p));
120+
vec_t<Dim::_3D> u_Cntrv { ZERO };
121+
122+
metric.template transform<Idx::D, Idx::U>(x_Code,
123+
{ ux1(p), ux2(p), ux3(p) },
124+
u_Cntrv);
125+
coeff = math::sqrt(ONE + u_Cntrv[0] * ux1(p) + u_Cntrv[1] * ux2(p) + u_Cntrv[2] * ux3(p));
126+
}
127+
128+
if constexpr (F == CustomField::V){
129+
coord_t<D> x_Code { ZERO };
130+
real_t gamma { ZERO };
131+
x_Code[0] = static_cast<real_t>(i1(p)) + static_cast<real_t>(dx1(p));
132+
x_Code[1] = static_cast<real_t>(i2(p)) + static_cast<real_t>(dx2(p));
133+
vec_t<Dim::_3D> u_Cntrv { ZERO };
134+
vec_t<Dim::_3D> u_Phys { ZERO };
135+
metric.template transform<Idx::D, Idx::U>(x_Code,
136+
{ ux1(p), ux2(p), ux3(p) },
137+
u_Cntrv);
138+
gamma = math::sqrt(ONE + u_Cntrv[0] * ux1(p) + u_Cntrv[1] * ux2(p) + u_Cntrv[2] * ux3(p));
139+
140+
metric.template transform<Idx::U, Idx::PU>(x_Code, u_Cntrv, u_Phys);
141+
vec_t<Dim::_3D> beta_Phys { ZERO };
142+
metric.template transform<Idx::U, Idx::PU>(x_Code,
143+
{ metric.beta1(x_Code), 0, 0 },
144+
beta_Phys);
145+
coeff = u_Phys[comp] / gamma * metric.alpha(x_Code)
146+
- beta_Phys[0] * static_cast<real_t>(comp == 0);
147+
}
148+
149+
if constexpr (F == CustomField::Ut){
150+
coord_t<D> x_Code { ZERO };
151+
real_t gamma { ZERO };
152+
x_Code[0] = static_cast<real_t>(i1(p)) + static_cast<real_t>(dx1(p));
153+
x_Code[1] = static_cast<real_t>(i2(p)) + static_cast<real_t>(dx2(p));
154+
vec_t<Dim::_3D> u_Cntrv { ZERO };
155+
metric.template transform<Idx::D, Idx::U>(x_Code,
156+
{ ux1(p), ux2(p), ux3(p) },
157+
u_Cntrv);
158+
gamma = math::sqrt(ONE + u_Cntrv[0] * ux1(p) + u_Cntrv[1] * ux2(p) + u_Cntrv[2] * ux3(p));
159+
coeff = -metric.alpha(x_Code) * gamma + metric.beta1(x_Code) * ux1(p);
160+
}
161+
162+
coeff *= weight(p);
163+
coeff *= smooth / metric.sqrt_det_h({ static_cast<real_t>(i1(p)) + HALF,
164+
static_cast<real_t>(i2(p)) + HALF });
165+
166+
auto buff_access = Buff.access();
167+
168+
if constexpr (D == Dim::_2D) {
169+
for (auto di2 { -window }; di2 <= window; ++di2) {
170+
for (auto di1 { -window }; di1 <= window; ++di1) {
171+
172+
// reflect contribution at axes
173+
if (is_axis_i2min && (i2(p) + di2 < 0)) {
174+
buff_access(i1(p) + di1 + N_GHOSTS,
175+
N_GHOSTS - (i2(p) + di2),
176+
buff_idx) += coeff;
177+
} else if (is_axis_i2max && (i2(p) + di2 >= ni2)) {
178+
buff_access(i1(p) + di1 + N_GHOSTS,
179+
2 * ni2 - (i2(p) + di2) + N_GHOSTS,
180+
buff_idx) += coeff;
181+
} else {
182+
buff_access(i1(p) + di1 + N_GHOSTS,
183+
i2(p) + di2 + N_GHOSTS,
184+
buff_idx) += coeff;
185+
}
186+
}
187+
}
188+
}
189+
190+
} // operator()
191+
192+
}; // class CustomMoments_kernel
193+
194+
195+
22196
template <class M, Dimension D>
23197
struct InitFields {
24198
InitFields(M metric_, real_t m_eps) : metric { metric_ }, m_eps { m_eps } {}
@@ -227,6 +401,7 @@ namespace user {
227401
const std::vector<real_t> xi_min;
228402
const std::vector<real_t> xi_max;
229403
const real_t sigma0, sigma_max, multiplicity, nGJ, temperature, m_eps;
404+
const real_t inv_n0;
230405

231406
InitFields<M, D> init_flds;
232407
const Metadomain<S, M>* metadomain;
@@ -242,6 +417,7 @@ namespace user {
242417
SQR(p.template get<real_t>("scales.skindepth0")) }
243418
, temperature { p.template get<real_t>("setup.temperature") }
244419
, m_eps { p.template get<real_t>("setup.m_eps") }
420+
, inv_n0 { ONE / p.template get<real_t>("scales.n0") }
245421
, init_flds { m.mesh().metric, m_eps }
246422
, metadomain { &m } {}
247423

@@ -290,6 +466,188 @@ namespace user {
290466
1.0,
291467
true);
292468
}
469+
470+
void CustomFieldOutput(const std::string& name,
471+
ndfield_t<M::Dim, 6> buffer,
472+
index_t index,
473+
timestep_t,
474+
simtime_t,
475+
const Domain<S, M>& domain) {
476+
477+
if (name == "DB") {
478+
if constexpr (M::Dim == Dim::_2D) {
479+
const auto& EM = domain.fields.em;
480+
const auto& metric = domain.mesh.metric;
481+
Kokkos::parallel_for(
482+
"DB",
483+
domain.mesh.rangeActiveCells(),
484+
Lambda(index_t i1, index_t i2) {
485+
coord_t<M::Dim> xi { static_cast<real_t>(i1 - N_GHOSTS), static_cast<real_t>(i2 - N_GHOSTS) };
486+
const vec_t<Dim::_3D> B_cntrv { EM(i1, i2, em::bx1),
487+
EM(i1, i2, em::bx2),
488+
EM(i1, i2, em::bx3) };
489+
vec_t<Dim::_3D> B_cov { ZERO };
490+
const vec_t<Dim::_3D> D_cntrv { EM(i1, i2, em::dx1),
491+
EM(i1, i2, em::dx2),
492+
EM(i1, i2, em::dx3) };
493+
metric.template transform<Idx::U, Idx::D>(xi, B_cntrv, B_cov);
494+
buffer(i1, i2, index) = DOT(B_cov[0], B_cov[1], B_cov[2], D_cntrv[0], D_cntrv[1], D_cntrv[2])
495+
/ DOT(B_cntrv[0], B_cntrv[1], B_cntrv[2], B_cntrv[0], B_cntrv[1], B_cntrv[2]);
496+
});
497+
}
498+
} else if (name == "Gamma_1" || name == "Gamma_2"){
499+
const auto sp_idx = (name == "Gamma_1") ? 0 : 1;
500+
auto& sp = domain.species[sp_idx];
501+
if constexpr (M::Dim == Dim::_2D){
502+
auto scatter_buff = Kokkos::Experimental::create_scatter_view(buffer);
503+
const auto& metric = domain.mesh.metric;
504+
auto& mesh = domain.mesh;
505+
const auto ni2 = mesh.n_active(in::x2);
506+
// clang-format off
507+
Kokkos::parallel_for(
508+
name,
509+
sp.rangeActiveParticles(),
510+
CustomMoments_kernel<M, CustomField::Gamma>(0, scatter_buff, index,
511+
sp.i1, sp.i2, sp.i3,
512+
sp.dx1, sp.dx2, sp.dx3,
513+
sp.ux1, sp.ux2, sp.ux3, sp.phi,
514+
sp.weight, sp.tag, sp.mass(), sp.charge(),
515+
metric, mesh.flds_bc(), ni2, inv_n0, ZERO));
516+
Kokkos::Experimental::contribute(buffer, scatter_buff);
517+
518+
auto n_buffer = domain.fields.buff;
519+
auto scatter_buff_n = Kokkos::Experimental::create_scatter_view(n_buffer);
520+
Kokkos::parallel_for(
521+
"ComputeMoments",
522+
sp.rangeActiveParticles(),
523+
kernel::ParticleMoments_kernel<S, M, FldsID::N, 3>({}, scatter_buff_n, 0u,
524+
sp.i1, sp.i2, sp.i3,
525+
sp.dx1, sp.dx2, sp.dx3,
526+
sp.ux1, sp.ux2, sp.ux3,
527+
sp.phi, sp.weight, sp.tag,
528+
sp.mass(), sp.charge(),
529+
true,
530+
metric, mesh.flds_bc(),
531+
ni2, inv_n0, ZERO));
532+
Kokkos::Experimental::contribute(n_buffer, scatter_buff_n);
533+
534+
Kokkos::parallel_for(
535+
"ComputeMoments",
536+
mesh.rangeActiveCells(),
537+
Lambda(index_t i1, index_t i2) {
538+
if (cmp::AlmostZero(n_buffer(i1, i2, 0))) {
539+
buffer(i1, i2, index) = ZERO;
540+
} else {
541+
buffer(i1, i2, index) /= n_buffer(i1, i2, 0);
542+
}
543+
});
544+
// clang-format on
545+
}
546+
} else if (name == "Vr_1" || name == "Vth_1" || name == "Vph_1"
547+
|| name == "Vr_2" || name == "Vth_2" || name == "Vph_2"){
548+
const auto comp = (name == "Vr_1" || name == "Vr_2") ? 0 : (name == "Vth_1" || name == "Vth_2") ? 1 : 2;
549+
const auto sp_idx = (name == "Vr_1" || name == "Vth_1" || name == "Vph_1") ? 0 : 1;
550+
auto& sp = domain.species[sp_idx];
551+
552+
553+
if constexpr (M::Dim == Dim::_2D){
554+
auto scatter_buff = Kokkos::Experimental::create_scatter_view(buffer);
555+
const auto& metric = domain.mesh.metric;
556+
auto& mesh = domain.mesh;
557+
const auto ni2 = mesh.n_active(in::x2);
558+
// clang-format off
559+
Kokkos::parallel_for(
560+
name,
561+
sp.rangeActiveParticles(),
562+
CustomMoments_kernel<M, CustomField::V>(comp, scatter_buff, index,
563+
sp.i1, sp.i2, sp.i3,
564+
sp.dx1, sp.dx2, sp.dx3,
565+
sp.ux1, sp.ux2, sp.ux3, sp.phi,
566+
sp.weight, sp.tag, sp.mass(), sp.charge(),
567+
metric, mesh.flds_bc(), ni2, inv_n0, ZERO));
568+
569+
Kokkos::Experimental::contribute(buffer, scatter_buff);
570+
571+
auto n_buffer = domain.fields.buff;
572+
auto scatter_buff_n = Kokkos::Experimental::create_scatter_view(n_buffer);
573+
Kokkos::parallel_for(
574+
"ComputeMoments",
575+
sp.rangeActiveParticles(),
576+
kernel::ParticleMoments_kernel<S, M, FldsID::N, 3>({}, scatter_buff_n, 0u,
577+
sp.i1, sp.i2, sp.i3,
578+
sp.dx1, sp.dx2, sp.dx3,
579+
sp.ux1, sp.ux2, sp.ux3,
580+
sp.phi, sp.weight, sp.tag,
581+
sp.mass(), sp.charge(),
582+
true,
583+
metric, mesh.flds_bc(),
584+
ni2, inv_n0, ZERO));
585+
Kokkos::Experimental::contribute(n_buffer, scatter_buff_n);
586+
587+
Kokkos::parallel_for(
588+
"ComputeMoments",
589+
mesh.rangeActiveCells(),
590+
Lambda(index_t i1, index_t i2) {
591+
if (cmp::AlmostZero(n_buffer(i1, i2, 0))) {
592+
buffer(i1, i2, index) = ZERO;
593+
} else {
594+
buffer(i1, i2, index) /= n_buffer(i1, i2, 0);
595+
}
596+
});
597+
// clang-format on
598+
}
599+
} else if (name == "Ut_1" || name == "Ut_2"){
600+
const auto sp_idx = (name == "Ut_1") ? 0 : 1;
601+
auto& sp = domain.species[sp_idx];
602+
auto scatter_buff = Kokkos::Experimental::create_scatter_view(buffer);
603+
const auto& metric = domain.mesh.metric;
604+
auto& mesh = domain.mesh;
605+
const auto ni2 = mesh.n_active(in::x2);
606+
// clang-format off
607+
Kokkos::parallel_for(
608+
name,
609+
sp.rangeActiveParticles(),
610+
CustomMoments_kernel<M, CustomField::Ut>(0, scatter_buff, index,
611+
sp.i1, sp.i2, sp.i3,
612+
sp.dx1, sp.dx2, sp.dx3,
613+
sp.ux1, sp.ux2, sp.ux3, sp.phi,
614+
sp.weight, sp.tag, sp.mass(), sp.charge(),
615+
metric, mesh.flds_bc(), ni2, inv_n0, ZERO));
616+
Kokkos::Experimental::contribute(buffer, scatter_buff);
617+
618+
auto n_buffer = domain.fields.buff;
619+
auto scatter_buff_n = Kokkos::Experimental::create_scatter_view(n_buffer);
620+
Kokkos::parallel_for(
621+
"ComputeMoments",
622+
sp.rangeActiveParticles(),
623+
kernel::ParticleMoments_kernel<S, M, FldsID::N, 3>({}, scatter_buff_n, 0u,
624+
sp.i1, sp.i2, sp.i3,
625+
sp.dx1, sp.dx2, sp.dx3,
626+
sp.ux1, sp.ux2, sp.ux3,
627+
sp.phi, sp.weight, sp.tag,
628+
sp.mass(), sp.charge(),
629+
true,
630+
metric, mesh.flds_bc(),
631+
ni2, inv_n0, ZERO));
632+
Kokkos::Experimental::contribute(n_buffer, scatter_buff_n);
633+
634+
Kokkos::parallel_for(
635+
"ComputeMoments",
636+
mesh.rangeActiveCells(),
637+
Lambda(index_t i1, index_t i2) {
638+
if (cmp::AlmostZero(n_buffer(i1, i2, 0))) {
639+
buffer(i1, i2, index) = ZERO;
640+
} else {
641+
buffer(i1, i2, index) /= n_buffer(i1, i2, 0);
642+
}
643+
});
644+
// clang-format on
645+
} else {
646+
raise::Error("Custom output not provided", HERE);
647+
}
648+
}
649+
650+
293651
};
294652

295653
} // namespace user

0 commit comments

Comments
 (0)