Skip to content

Commit 5dbe653

Browse files
StaticObserverStaticObserver
committed
adding BZjet pgen
1 parent 8488b46 commit 5dbe653

1 file changed

Lines changed: 364 additions & 0 deletions

File tree

‎pgens/BZjet/pgen.hpp‎

Lines changed: 364 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,364 @@
1+
#ifndef PROBLEM_GENERATOR_H
2+
#define PROBLEM_GENERATOR_H
3+
4+
#include "enums.h"
5+
#include "global.h"
6+
7+
#include "arch/kokkos_aliases.h"
8+
#include "arch/traits.h"
9+
#include "utils/comparators.h"
10+
#include "utils/formatting.h"
11+
#include "utils/log.h"
12+
#include "utils/numeric.h"
13+
14+
#include "archetypes/energy_dist.h"
15+
#include "archetypes/particle_injector.h"
16+
#include "archetypes/problem_generator.h"
17+
#include "framework/domain/domain.h"
18+
#include "framework/domain/metadomain.h"
19+
20+
#include <vector>
21+
22+
namespace user {
23+
using namespace ntt;
24+
25+
template <class M, Dimension D>
26+
struct InitFields {
27+
InitFields(M metric_) : metric { metric_ } {}
28+
29+
Inline auto A_3(const coord_t<D>& x_Cd) const -> real_t {
30+
return HALF * (metric.template h_<3, 3>(x_Cd) +
31+
TWO * metric.spin() * metric.template h_<1, 3>(x_Cd) *
32+
metric.beta1(x_Cd));
33+
}
34+
35+
Inline auto A_1(const coord_t<D>& x_Cd) const -> real_t {
36+
return HALF * (metric.template h_<1, 3>(x_Cd) +
37+
TWO * metric.spin() * metric.template h_<1, 1>(x_Cd) *
38+
metric.beta1(x_Cd));
39+
}
40+
41+
Inline auto A_0(const coord_t<D>& x_Cd) const -> real_t {
42+
real_t g_00 { -metric.alpha(x_Cd) * metric.alpha(x_Cd) +
43+
metric.template h_<1, 1>(x_Cd) * metric.beta1(x_Cd) *
44+
metric.beta1(x_Cd) };
45+
return HALF * (metric.template h_<1, 3>(x_Cd) * metric.beta1(x_Cd) +
46+
TWO * metric.spin() * g_00);
47+
}
48+
49+
Inline auto bx1(const coord_t<D>& x_Ph) const -> real_t { // at ( i , j + HALF )
50+
coord_t<D> xi { ZERO }, x0m { ZERO }, x0p { ZERO };
51+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
52+
53+
x0m[0] = xi[0];
54+
x0m[1] = xi[1] - HALF;
55+
x0p[0] = xi[0];
56+
x0p[1] = xi[1] + HALF;
57+
58+
real_t inv_sqrt_detH_ijP { ONE / metric.sqrt_det_h({ xi[0], xi[1] }) };
59+
60+
if (cmp::AlmostZero(x_Ph[1])) {
61+
return ONE;
62+
} else {
63+
return (A_3(x0p) - A_3(x0m)) * inv_sqrt_detH_ijP;
64+
}
65+
}
66+
67+
Inline auto bx2(const coord_t<D>& x_Ph) const -> real_t { // at ( i + HALF , j )
68+
coord_t<D> xi { ZERO }, x0m { ZERO }, x0p { ZERO };
69+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
70+
71+
x0m[0] = xi[0] - HALF;
72+
x0m[1] = xi[1];
73+
x0p[0] = xi[0] + HALF;
74+
x0p[1] = xi[1];
75+
76+
real_t inv_sqrt_detH_ijP { ONE / metric.sqrt_det_h({ xi[0], xi[1] }) };
77+
if (cmp::AlmostZero(x_Ph[1])) {
78+
return ZERO;
79+
} else {
80+
return -(A_3(x0p) - A_3(x0m)) * inv_sqrt_detH_ijP;
81+
}
82+
}
83+
84+
Inline auto bx3(
85+
const coord_t<D>& x_Ph) const -> real_t { // at ( i + HALF , j + HALF )
86+
coord_t<D> xi { ZERO }, x0m { ZERO }, x0p { ZERO };
87+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
88+
89+
x0m[0] = xi[0];
90+
x0m[1] = xi[1] - HALF;
91+
x0p[0] = xi[0];
92+
x0p[1] = xi[1] + HALF;
93+
94+
real_t inv_sqrt_detH_iPjP { ONE / metric.sqrt_det_h({ xi[0], xi[1] }) };
95+
return -(A_1(x0p) - A_1(x0m)) * inv_sqrt_detH_iPjP;
96+
}
97+
98+
Inline auto dx1(const coord_t<D>& x_Ph) const -> real_t { // at ( i + HALF , j )
99+
coord_t<D> xi { ZERO }, x0m { ZERO }, x0p { ZERO };
100+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
101+
102+
real_t alpha_iPj { metric.alpha({ xi[0], xi[1] }) };
103+
real_t inv_sqrt_detH_ij { ONE / metric.sqrt_det_h({ xi[0] - HALF, xi[1] }) };
104+
real_t sqrt_detH_ij { metric.sqrt_det_h({ xi[0] - HALF, xi[1] }) };
105+
real_t beta_ij { metric.beta1({ xi[0] - HALF, xi[1] }) };
106+
real_t alpha_ij { metric.alpha({ xi[0] - HALF, xi[1] }) };
107+
108+
// D1 at ( i + HALF , j )
109+
x0m[0] = xi[0] - HALF;
110+
x0m[1] = xi[1];
111+
x0p[0] = xi[0] + HALF;
112+
x0p[1] = xi[1];
113+
real_t E1d { (A_0(x0p) - A_0(x0m)) };
114+
real_t D1d { E1d / alpha_iPj };
115+
116+
// D3 at ( i , j )
117+
x0m[0] = xi[0] - HALF - HALF;
118+
x0m[1] = xi[1];
119+
x0p[0] = xi[0] - HALF + HALF;
120+
x0p[1] = xi[1];
121+
real_t D3d { (A_3(x0p) - A_3(x0m)) * beta_ij / alpha_ij };
122+
123+
real_t D1u { metric.template h<1, 1>({ xi[0], xi[1] }) * D1d +
124+
metric.template h<1, 3>({ xi[0], xi[1] }) * D3d };
125+
126+
return D1u;
127+
}
128+
129+
Inline auto dx2(const coord_t<D>& x_Ph) const -> real_t { // at ( i , j + HALF )
130+
coord_t<D> xi { ZERO }, x0m { ZERO }, x0p { ZERO };
131+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
132+
x0m[0] = xi[0];
133+
x0m[1] = xi[1] - HALF;
134+
x0p[0] = xi[0];
135+
x0p[1] = xi[1] + HALF;
136+
real_t inv_sqrt_detH_ijP { ONE / metric.sqrt_det_h({ xi[0], xi[1] }) };
137+
real_t sqrt_detH_ijP { metric.sqrt_det_h({ xi[0], xi[1] }) };
138+
real_t alpha_ijP { metric.alpha({ xi[0], xi[1] }) };
139+
real_t beta_ijP { metric.beta1({ xi[0], xi[1] }) };
140+
141+
real_t E2d { (A_0(x0p) - A_0(x0m)) };
142+
real_t D2d { E2d / alpha_ijP - (A_1(x0p) - A_1(x0m)) * beta_ijP / alpha_ijP };
143+
real_t D2u { metric.template h<2, 2>({ xi[0], xi[1] }) * D2d };
144+
145+
return D2u;
146+
}
147+
148+
Inline auto dx3(const coord_t<D>& x_Ph) const -> real_t { // at ( i , j )
149+
coord_t<D> xi { ZERO }, x0m { ZERO }, x0p { ZERO };
150+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
151+
real_t inv_sqrt_detH_ij { ONE / metric.sqrt_det_h({ xi[0], xi[1] }) };
152+
real_t sqrt_detH_ij { metric.sqrt_det_h({ xi[0], xi[1] }) };
153+
real_t beta_ij { metric.beta1({ xi[0], xi[1] }) };
154+
real_t alpha_ij { metric.alpha({ xi[0], xi[1] }) };
155+
real_t alpha_iPj { metric.alpha({ xi[0] + HALF, xi[1] }) };
156+
157+
// D3 at ( i , j )
158+
x0m[0] = xi[0] - HALF;
159+
x0m[1] = xi[1];
160+
x0p[0] = xi[0] + HALF;
161+
x0p[1] = xi[1];
162+
real_t D3d { (A_3(x0p) - A_3(x0m)) * beta_ij / alpha_ij };
163+
164+
// D1 at ( i + HALF , j )
165+
x0m[0] = xi[0] + HALF - HALF;
166+
x0m[1] = xi[1];
167+
x0p[0] = xi[0] + HALF + HALF;
168+
x0p[1] = xi[1];
169+
real_t E1d { (A_0(x0p) - A_0(x0m)) };
170+
real_t D1d { E1d / alpha_iPj };
171+
172+
if (cmp::AlmostZero(x_Ph[1])) {
173+
return metric.template h<1, 3>({ xi[0], xi[1] }) * D1d;
174+
} else {
175+
return metric.template h<3, 3>({ xi[0], xi[1] }) * D3d +
176+
metric.template h<1, 3>({ xi[0], xi[1] }) * D1d;
177+
}
178+
}
179+
180+
private:
181+
const M metric;
182+
};
183+
184+
template <SimEngine::type S, class M>
185+
struct PointDistribution : public arch::SpatialDistribution<S, M> {
186+
PointDistribution(const std::vector<real_t>& xi_min,
187+
const std::vector<real_t>& xi_max,
188+
const real_t d0,
189+
const real_t rho0,
190+
const real_t sigma_thr,
191+
const real_t db_thr,
192+
const real_t dens_thr,
193+
const SimulationParams& params,
194+
Domain<S, M>* domain_ptr)
195+
: arch::SpatialDistribution<S, M> { domain_ptr->mesh.metric }
196+
, metric { domain_ptr->mesh.metric }
197+
, EM { domain_ptr->fields.em }
198+
, density { domain_ptr->fields.buff }
199+
, sigma_thr { sigma_thr }
200+
, db_thr { db_thr }
201+
, d0 { params.template get<real_t>("scales.skindepth0") }
202+
, rho0 { params.template get<real_t>("scales.larmor0") }
203+
, inv_n0 { ONE / params.template get<real_t>("scales.n0") }
204+
, dens_thr { dens_thr } {
205+
std::copy(xi_min.begin(), xi_min.end(), x_min);
206+
std::copy(xi_max.begin(), xi_max.end(), x_max);
207+
208+
std::vector<unsigned short> specs {};
209+
for (auto& sp : domain_ptr->species) {
210+
if (sp.mass() > 0) {
211+
specs.push_back(sp.index());
212+
}
213+
}
214+
215+
Kokkos::deep_copy(density, ZERO);
216+
auto scatter_buff = Kokkos::Experimental::create_scatter_view(density);
217+
// some parameters
218+
auto& mesh = domain_ptr->mesh;
219+
const auto use_weights = params.template get<bool>(
220+
"particles.use_weights");
221+
const auto ni2 = mesh.n_active(in::x2);
222+
223+
for (const auto& sp : specs) {
224+
auto& prtl_spec = domain_ptr->species[sp - 1];
225+
// clang-format off
226+
Kokkos::parallel_for(
227+
"ComputeMoments",
228+
prtl_spec.rangeActiveParticles(),
229+
kernel::ParticleMoments_kernel<S, M, FldsID::Rho, 3>({}, scatter_buff, 0u,
230+
prtl_spec.i1, prtl_spec.i2, prtl_spec.i3,
231+
prtl_spec.dx1, prtl_spec.dx2, prtl_spec.dx3,
232+
prtl_spec.ux1, prtl_spec.ux2, prtl_spec.ux3,
233+
prtl_spec.phi, prtl_spec.weight, prtl_spec.tag,
234+
prtl_spec.mass(), prtl_spec.charge(),
235+
use_weights,
236+
metric, mesh.flds_bc(),
237+
ni2, inv_n0, ZERO));
238+
// clang-format on
239+
}
240+
Kokkos::Experimental::contribute(density, scatter_buff);
241+
}
242+
243+
Inline auto sigma_crit(const coord_t<M::Dim>& x_Ph) const -> bool {
244+
coord_t<M::Dim> xi { ZERO };
245+
if constexpr (M::Dim == Dim::_2D) {
246+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
247+
const auto i1 = static_cast<int>(xi[0]) + static_cast<int>(N_GHOSTS);
248+
const auto i2 = static_cast<int>(xi[1]) + static_cast<int>(N_GHOSTS);
249+
const vec_t<Dim::_3D> B_cntrv { EM(i1, i2, em::bx1),
250+
EM(i1, i2, em::bx2),
251+
EM(i1, i2, em::bx3) };
252+
const vec_t<Dim::_3D> D_cntrv { EM(i1, i2, em::dx1),
253+
EM(i1, i2, em::dx2),
254+
EM(i1, i2, em::dx3) };
255+
vec_t<Dim::_3D> B_cov { ZERO };
256+
metric.template transform<Idx::U, Idx::D>(xi, B_cntrv, B_cov);
257+
const auto bsqr =
258+
DOT(B_cntrv[0], B_cntrv[1], B_cntrv[2], B_cov[0], B_cov[1], B_cov[2]);
259+
const auto db = DOT(D_cntrv[0], D_cntrv[1], D_cntrv[2], B_cov[0], B_cov[1], B_cov[2]);
260+
const auto dens = density(i1, i2, 0);
261+
return (bsqr > sigma_thr * dens) && (db > db_thr * bsqr);
262+
}
263+
return false;
264+
}
265+
266+
Inline auto operator()(const coord_t<M::Dim>& x_Ph) const -> real_t override {
267+
auto fill = false;
268+
for (auto d = 0u; d < M::Dim; ++d) {
269+
fill &= x_Ph[d] > x_min[d] and x_Ph[d] < x_max[d] and sigma_crit(x_Ph);
270+
}
271+
metric.template convert<Crd::Ph, Crd::Cd>(x_Ph, xi);
272+
const auto i1 = static_cast<int>(xi[0]) + static_cast<int>(N_GHOSTS);
273+
const auto i2 = static_cast<int>(xi[1]) + static_cast<int>(N_GHOSTS);
274+
const vec_t<Dim::_3D> B_cntrv { EM(i1, i2, em::bx1),
275+
EM(i1, i2, em::bx2),
276+
EM(i1, i2, em::bx3) };
277+
const vec_t<Dim::_3D> D_cntrv { EM(i1, i2, em::dx1),
278+
EM(i1, i2, em::dx2),
279+
EM(i1, i2, em::dx3) };
280+
vec_t<Dim::_3D> B_cov { ZERO };
281+
metric.template transform<Idx::U, Idx::D>(xi, B_cntrv, B_cov);
282+
const auto bsqr =
283+
DOT(B_cntrv[0], B_cntrv[1], B_cntrv[2], B_cov[0], B_cov[1], B_cov[2]);
284+
const auto db = DOT(D_cntrv[0], D_cntrv[1], D_cntrv[2], B_cov[0], B_cov[1], B_cov[2]);
285+
const real_t inj_n = dens_thr * db / SQRT(bsqr) * SQR(d0) / rho0;
286+
return fill ? inj_n : ZERO;
287+
}
288+
289+
private:
290+
tuple_t<real_t, M::Dim> x_min;
291+
tuple_t<real_t, M::Dim> x_max;
292+
const real_t sigma_thr;
293+
const real_t db_thr;
294+
const real_t dens_thr;
295+
const real_t inv_n0;
296+
const real_t d0;
297+
const real_t rho0;
298+
Domain<S, M>* domain_ptr;
299+
ndfield_t<M::Dim, 3> density;
300+
ndfield_t<M::Dim, 6> EM;
301+
const M metric;
302+
};
303+
304+
305+
template <SimEngine::type S, class M>
306+
struct PGen : public arch::ProblemGenerator<S, M> {
307+
// compatibility traits for the problem generator
308+
static constexpr auto engines { traits::compatible_with<SimEngine::GRPIC>::value };
309+
static constexpr auto metrics {
310+
traits::compatible_with<Metric::Kerr_Schild, Metric::QKerr_Schild, Metric::Kerr_Schild_0>::value
311+
};
312+
static constexpr auto dimensions { traits::compatible_with<Dim::_2D>::value };
313+
314+
// for easy access to variables in the child class
315+
using arch::ProblemGenerator<S, M>::D;
316+
using arch::ProblemGenerator<S, M>::C;
317+
using arch::ProblemGenerator<S, M>::params;
318+
319+
const std::vector<real_t> xi_min;
320+
const std::vector<real_t> xi_max;
321+
const real_t sigma0, sigma_max, multiplicity, temperature, m_eps;
322+
323+
InitFields<M, D> init_flds;
324+
const Metadomain<S, M>* metadomain;
325+
326+
inline PGen(SimulationParams& p, const Metadomain<S, M>& m)
327+
: arch::ProblemGenerator<S, M>(p)
328+
, xi_min { p.template get<std::vector<real_t>>("setup.xi_min") }
329+
, xi_max { p.template get<std::vector<real_t>>("setup.xi_max") }
330+
, sigma_max { p.template get<real_t>("setup.sigma_max") }
331+
, sigma0 { p.template get<real_t>("scales.sigma0") }
332+
, multiplicity { p.template get<real_t>("setup.multiplicity") }
333+
, temperature { p.template get<real_t>("setup.temperature") }
334+
, m_eps { p.template get<real_t>("setup.m_eps") }
335+
, init_flds { m.mesh().metric, m_eps }
336+
, metadomain { &m } {}
337+
338+
void CustomPostStep(std::size_t, long double time, Domain<S, M>& local_domain) {
339+
const auto energy_dist = arch::Maxwellian<S, M>(local_domain.mesh.metric,
340+
local_domain.random_pool,
341+
temperature);
342+
const auto spatial_dist = PointDistribution<S, M>(xi_min,
343+
xi_max,
344+
sigma_max / sigma0,
345+
multiplicity,
346+
params,
347+
&local_domain);
348+
349+
const auto injector =
350+
arch::NonUniformInjector<S, M, arch::Maxwellian, PointDistribution>(
351+
energy_dist,
352+
spatial_dist,
353+
{ 1, 2 });
354+
arch::InjectNonUniform<S, M, decltype(injector)>(params,
355+
local_domain,
356+
injector,
357+
1.0,
358+
true);
359+
}
360+
};
361+
362+
} // namespace user
363+
364+
#endif

0 commit comments

Comments
 (0)