Skip to content

Commit

Permalink
[Draft] AMReX SoA Named Components
Browse files Browse the repository at this point in the history
Replace our own named SoA particle components implementation
with the now upstream support in AMReX.
  • Loading branch information
ax3l committed Jan 14, 2025
1 parent f3a7b11 commit a3b79d3
Show file tree
Hide file tree
Showing 4 changed files with 36 additions and 76 deletions.
9 changes: 9 additions & 0 deletions src/Particle/ParticleContainer.H
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,15 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr)
"add a new runtime component with type Int"
)

.def_property_readonly("real_soa_names",
&ParticleContainerType::GetRealSoANames,
"Get the names for the Real SoA components"
)
.def_property_readonly("int_soa_names",
&ParticleContainerType::GetIntSoANames,
"Get the names for the int SoA components"
)

.def_property_readonly("finest_level", &ParticleContainerBase::finestLevel)

// ParticleContainer ( const ParticleContainer &) = delete;
Expand Down
27 changes: 27 additions & 0 deletions src/Particle/StructOfArrays.H
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include <AMReX_GpuAllocators.H>
#include <AMReX_StructOfArrays.H>

#include <stdexcept>
#include <sstream>


Expand Down Expand Up @@ -56,6 +57,32 @@ void make_StructOfArrays(py::module &m, std::string allocstr)
py::arg("index"),
"Get access to a particle Real component Array (compile-time and runtime component)")

// names
.def_property_readonly("real_names",
[](SOAType & self) {
try {
return self.GetRealNames();
}
catch (...)
{
return std::vector<std::string>{};
}
},
"Names for the Real SoA components"
)
.def_property_readonly("int_names",
[](SOAType & self) {
try {
return self.GetIntNames();
}
catch (...)
{
return std::vector<std::string>{};
}
},
"Names for the int SoA components"
)

.def("__len__", &SOAType::size,
"Get the number of particles")
.def_property_readonly("size", &SOAType::size,
Expand Down
7 changes: 0 additions & 7 deletions src/amrex/extensions/ParticleComponentNames.py

This file was deleted.

69 changes: 0 additions & 69 deletions src/amrex/extensions/StructOfArrays.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,71 +9,6 @@
from collections import namedtuple


def soa_real_comps(self, num_comps, spacedim=3, rotate=True):
"""
Name the ParticleReal components in SoA.
Parameters
----------
self : SoA Type
maybe unused, depending on implementation
num_comps : int
number of components to generate names for.
spacedim : int
AMReX dimensionality
rotate : bool = True
start with "x", "y", "z", "a", "b", ...
Returns
-------
A list of length num_comps with values
rotate=True (for pure SoA layout):
- 3D: "x", "y", "z", "a", "b", ... "w", "r0", "r1", ...
- 2D: "x", "y", "a", "b", ... "w", "r0", "r1", ...
- 1D: "x", "a", "b", ... "w", "r0", "r1", ...
rotate=False (for legacy layout):
- 1D-3D: "a", "b", ... "w", "r0", "r1", ...
"""
import string

# x, y, z, a, b, ...
comp_names = list(string.ascii_lowercase)
if rotate:
# rotate x, y, z to be beginning (positions)
comp_names = comp_names[-3:] + comp_names[:-3]
else:
# cut off x, y, z to avoid confusion
comp_names = comp_names[:-3]

num_named = len(comp_names)
if num_comps < num_named:
comp_names = list(comp_names)[0:num_comps]
elif num_comps > num_named:
comp_names.extend(["r" + str(i) for i in range(num_comps - num_named)])

return comp_names


def soa_int_comps(self, num_comps):
"""
Name the int components in SoA.
Parameters
----------
self : SoA Type
maybe unused, depending on implementation
num_comps : int
number of components to generate names for.
Returns
-------
A list of length num_comps with values "i1", "i2", "i3", ...
"""
comp_names = ["i" + str(i) for i in range(num_comps)]

return comp_names


def soa_to_numpy(self, copy=False):
"""
Provide NumPy views into a StructOfArrays.
Expand Down Expand Up @@ -226,10 +161,6 @@ def register_SoA_extension(amr):
and member.__module__ == amr.__name__
and member.__name__.startswith("StructOfArrays_"),
):
# name providers
SoA_type.soa_real_comps = soa_real_comps
SoA_type.soa_int_comps = soa_int_comps

# converters
SoA_type.to_numpy = soa_to_numpy
SoA_type.to_cupy = soa_to_cupy
Expand Down

0 comments on commit a3b79d3

Please sign in to comment.