diff --git a/src/Particle/ParticleContainer.H b/src/Particle/ParticleContainer.H index ef103d3e..e2ac199e 100644 --- a/src/Particle/ParticleContainer.H +++ b/src/Particle/ParticleContainer.H @@ -204,6 +204,9 @@ void make_ParticleContainer_and_Iterators (py::module &m, std::string allocstr) .def_property_readonly("num_position_components", [](const py::object&){ return AMREX_SPACEDIM; }) .def_property_readonly("byte_spread", &ParticleContainerType::ByteSpread) + // compile-time components + .def("set_soa_compile_time_names", &ParticleContainerType::SetSoACompileTimeNames) + // runtime components .def("add_real_comp", py::overload_cast(&ParticleContainerType::AddRealComp), py::arg("communicate")=1, @@ -224,6 +227,33 @@ 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("has_real_comp", + &ParticleContainerType::HasRealComp, + "Check if a container has an ParticleReal component" + ) + .def("has_int_comp", + &ParticleContainerType::HasIntComp, + "Check if a container has an Integer component" + ) + + .def("get_real_comp_index", + &ParticleContainerType::GetRealCompIndex, + "Get the ParticleReal SoA index of a component" + ) + .def("get_int_comp_index", + &ParticleContainerType::GetIntCompIndex, + "Get the Integer SoA index of a component" + ) + .def_property_readonly("finest_level", &ParticleContainerBase::finestLevel) // ParticleContainer ( const ParticleContainer &) = delete; diff --git a/src/Particle/StructOfArrays.H b/src/Particle/StructOfArrays.H index e21d67cf..1c0c3712 100644 --- a/src/Particle/StructOfArrays.H +++ b/src/Particle/StructOfArrays.H @@ -56,6 +56,16 @@ 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::GetRealNames, + "Names for the Real SoA components" + ) + .def_property_readonly("int_names", + &SOAType::GetIntNames, + "Names for the int SoA components" + ) + .def("__len__", &SOAType::size, "Get the number of particles") .def_property_readonly("size", &SOAType::size, diff --git a/src/amrex/extensions/ParticleComponentNames.py b/src/amrex/extensions/ParticleComponentNames.py deleted file mode 100644 index 6922cdce..00000000 --- a/src/amrex/extensions/ParticleComponentNames.py +++ /dev/null @@ -1,7 +0,0 @@ -""" -This file is part of pyAMReX - -Copyright 2024 AMReX community -Authors: Axel Huebl -License: BSD-3-Clause-LBNL -""" diff --git a/src/amrex/extensions/StructOfArrays.py b/src/amrex/extensions/StructOfArrays.py index c23a5716..3f58a415 100644 --- a/src/amrex/extensions/StructOfArrays.py +++ b/src/amrex/extensions/StructOfArrays.py @@ -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. @@ -107,18 +42,17 @@ def soa_to_numpy(self, copy=False): else: soa_view = SoA_np({}, {}, None) - # for the legacy data layout, do not start with x, y, z but with a, b, c, ... - if self.has_idcpu: - real_comp_names = self.soa_real_comps(self.num_real_comps) - else: - real_comp_names = self.soa_real_comps(self.num_real_comps, rotate=False) - + real_comp_names = self.real_names + if len(real_comp_names) != self.num_real_comps: + raise ValueError("Missing names for SoA Real components.") for idx_real in range(self.num_real_comps): soa_view.real[real_comp_names[idx_real]] = self.get_real_data( idx_real ).to_numpy(copy=copy) - int_comp_names = self.soa_int_comps(self.num_int_comps) + int_comp_names = self.int_names + if len(int_comp_names) != self.num_int_comps: + raise ValueError("Missing names for SoA int components.") for idx_int in range(self.num_int_comps): soa_view.int[int_comp_names[idx_int]] = self.get_int_data(idx_int).to_numpy( copy=copy @@ -165,18 +99,17 @@ def soa_to_cupy(self, copy=False): else: soa_view = SoA_cp({}, {}, None) - # for the legacy data layout, do not start with x, y, z but with a, b, c, ... - if self.has_idcpu: - real_comp_names = self.soa_real_comps(self.num_real_comps) - else: - real_comp_names = self.soa_real_comps(self.num_real_comps, rotate=False) - + real_comp_names = self.real_names + if len(real_comp_names) != self.num_real_comps: + raise ValueError("Missing names for SoA Real components.") for idx_real in range(self.num_real_comps): soa_view.real[real_comp_names[idx_real]] = self.get_real_data(idx_real).to_cupy( copy=copy ) - int_comp_names = self.soa_int_comps(self.num_int_comps) + int_comp_names = self.int_names + if len(int_comp_names) != self.num_int_comps: + raise ValueError("Missing names for SoA int components.") for idx_int in range(self.num_int_comps): soa_view.int[int_comp_names[idx_int]] = self.get_int_data(idx_int).to_cupy( copy=copy @@ -226,10 +159,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 diff --git a/tests/test_particleContainer.py b/tests/test_particleContainer.py index 1e359ef5..4174c870 100644 --- a/tests/test_particleContainer.py +++ b/tests/test_particleContainer.py @@ -57,9 +57,9 @@ def particle_container(Npart, std_geometry, distmap, boxarr, std_real_box): pc.init_random(Npart, iseed, myt, False, std_real_box) # add runtime components: 1 real 2 int - pc.add_real_comp(True) - pc.add_int_comp(True) - pc.add_int_comp(True) + pc.add_real_comp("b", True) + pc.add_int_comp("i1", True) + pc.add_int_comp("i2", True) # assign some values to runtime components for lvl in range(pc.finest_level + 1): @@ -79,13 +79,21 @@ def soa_particle_container(Npart, std_geometry, distmap, boxarr, std_real_box): myt.real_array_data = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8] myt.int_array_data = [] + with pytest.raises(Exception): + pc.set_soa_compile_time_names( + ["x", "y", "z", "z", "b", "c", "d", "e"], [] + ) # error: z added twice + pc.set_soa_compile_time_names(["x", "y", "z", "a", "b", "c", "d", "e"], []) + iseed = 1 pc.init_random(Npart, iseed, myt, False, std_real_box) # add runtime components: 1 real 2 int - pc.add_real_comp(True) - pc.add_int_comp(True) - pc.add_int_comp(True) + with pytest.raises(Exception): + pc.add_real_comp("a", True) # already used as a compile-time component + pc.add_real_comp("f", True) + pc.add_int_comp("i1", True) + pc.add_int_comp("i2", True) # assign some values to runtime components for lvl in range(pc.finest_level + 1):