1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
|
/* Copyright 2021-2022 The WarpX Community
*
* Authors: Axel Huebl, Remi Lehe
* License: BSD-3-Clause-LBNL
*/
#include "Python/pyWarpX.H"
#include <Particles/WarpXParticleContainer.H>
void init_WarpXParIter (py::module& m)
{
py::class_<
WarpXParIter, amrex::ParIter<0,0,PIdx::nattribs>
>(m, "WarpXParIter")
.def(py::init<amrex::ParIter<0,0,PIdx::nattribs>::ContainerType&, int>(),
py::arg("particle_container"), py::arg("level"))
.def(py::init<amrex::ParIter<0,0,PIdx::nattribs>::ContainerType&, int, amrex::MFItInfo&>(),
py::arg("particle_container"), py::arg("level"),
py::arg("info"))
;
}
void init_WarpXParticleContainer (py::module& m)
{
py::class_<
WarpXParticleContainer,
amrex::ParticleContainer<0, 0, PIdx::nattribs, 0>
> wpc (m, "WarpXParticleContainer");
wpc
.def("add_real_comp",
[](WarpXParticleContainer& pc, const std::string& name, bool const comm) { pc.AddRealComp(name, comm); },
py::arg("name"), py::arg("comm")
)
.def("add_n_particles",
[](WarpXParticleContainer& pc, int lev,
int n, py::array_t<double> &x,
py::array_t<double> &y,
py::array_t<double> &z,
py::array_t<double> &ux,
py::array_t<double> &uy,
py::array_t<double> &uz,
const int nattr_real, py::array_t<double> &attr_real,
const int nattr_int, py::array_t<int> &attr_int,
int uniqueparticles, int id
) {
amrex::Vector<amrex::ParticleReal> xp(x.data(), x.data() + n);
amrex::Vector<amrex::ParticleReal> yp(y.data(), y.data() + n);
amrex::Vector<amrex::ParticleReal> zp(z.data(), z.data() + n);
amrex::Vector<amrex::ParticleReal> uxp(ux.data(), ux.data() + n);
amrex::Vector<amrex::ParticleReal> uyp(uy.data(), uy.data() + n);
amrex::Vector<amrex::ParticleReal> uzp(uz.data(), uz.data() + n);
// create 2d arrays of real and in attributes
amrex::Vector<amrex::Vector<amrex::ParticleReal>> attr;
const double *attr_data = attr_real.data();
for (int ii=0; ii<nattr_real; ii++) {
amrex::Vector<amrex::ParticleReal> attr_ii(n);
for (int jj=0; jj<n; jj++) {
attr_ii[jj] = attr_data[ii + jj*nattr_real];
}
attr.push_back(attr_ii);
}
amrex::Vector<amrex::Vector<int>> iattr;
const int *iattr_data = attr_int.data();
for (int ii=0; ii<nattr_int; ii++) {
amrex::Vector<int> attr_ii(n);
for (int jj=0; jj<n; jj++) {
attr_ii[jj] = iattr_data[ii + jj*nattr_int];
}
iattr.push_back(attr_ii);
}
pc.AddNParticles(
lev, n, xp, yp, zp, uxp, uyp, uzp, nattr_real, attr,
nattr_int, iattr, uniqueparticles, id
);
},
py::arg("lev"), py::arg("n"),
py::arg("x"), py::arg("y"), py::arg("z"),
py::arg("ux"), py::arg("uy"), py::arg("uz"),
py::arg("nattr_real"), py::arg("attr_real"),
py::arg("nattr_int"), py::arg("attr_int"),
py::arg("uniqueparticles"), py::arg("id")=-1
)
.def("num_real_comps", &WarpXParticleContainer::NumRealComps)
.def("get_comp_index",
[](WarpXParticleContainer& pc, std::string comp_name)
{
auto particle_comps = pc.getParticleComps();
return particle_comps.at(comp_name);
},
py::arg("comp_name")
)
.def("num_local_tiles_at_level",
&WarpXParticleContainer::numLocalTilesAtLevel,
py::arg("level")
)
.def("total_number_of_particles",
&WarpXParticleContainer::TotalNumberOfParticles,
py::arg("valid_particles_only"), py::arg("local")
)
.def("deposit_charge",
[](WarpXParticleContainer& pc,
amrex::MultiFab* rho, const int lev)
{
for (WarpXParIter pti(pc, lev); pti.isValid(); ++pti)
{
const long np = pti.numParticles();
auto& wp = pti.GetAttribs(PIdx::w);
pc.DepositCharge(pti, wp, nullptr, rho, 0, 0, np, 0, lev, lev);
}
},
py::arg("rho"), py::arg("lev")
)
;
}
|
