Skip to content

Commit

Permalink
Breaking Change: IntVectND
Browse files Browse the repository at this point in the history
  • Loading branch information
ax3l committed Jun 18, 2024
1 parent c36971f commit 5070edc
Showing 1 changed file with 159 additions and 136 deletions.
295 changes: 159 additions & 136 deletions src/Base/IntVect.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -15,146 +15,169 @@
#include <string>


namespace
{
template<int dim>
void init_IntVectND(py::module &m)
{
using namespace amrex;

auto const iv_name = std::string("IntVect").append(std::to_string(dim)).append("D");
using iv_type = IntVectND<dim>;

py::class_< iv_type > py_iv(m, iv_name.c_str());
py_iv
.def("__repr__",
[iv_name](const iv_type&) {
return "<amrex." + iv_name + ">";
}
)
.def("__str",
[](const iv_type& iv) {
std::stringstream s;
s << iv;
return s.str();
})
;

if constexpr (dim == 2) {
py_iv.def(py::init<int, int>());
} else if constexpr (dim == 3) {
py_iv.def(py::init<int, int, int>());
}

py_iv
.def(py::init<>())
.def(py::init<int>())
.def(py::init<const std::array<int, dim>&>())

.def_property_readonly("sum", &iv_type::sum)
.def_property_readonly("max",
py::overload_cast<>(&iv_type::max, py::const_))
.def_property_readonly("min",
py::overload_cast<>(&iv_type::min, py::const_))
.def_static("zero_vector", &iv_type::TheZeroVector)
.def_static("unit_vector", &iv_type::TheUnitVector)
.def_static("node_vector", &iv_type::TheNodeVector)
.def_static("cell_vector", &iv_type::TheCellVector)
.def_static("max_vector", &iv_type::TheMaxVector)
.def_static("min_vector", &iv_type::TheMinVector)
;

if constexpr (dim >= 1 && dim <=3) {
py_iv.def("dim3", &iv_type::template dim3<dim>);
}

py_iv
.def("__getitem__",
[](const iv_type& v, const int i) {
const int ii = (i >= 0) ? i : dim + i;
if ((ii < 0) || (ii >= dim))
throw py::index_error(
"Index must be between 0 and " +
std::to_string(dim));
return v[ii];
})
.def("__setitem__",
[](iv_type& v, const int i, const int& val) {
const int ii = (i >= 0) ? i : dim + i;
if ((ii < 0) || (ii >= dim))
throw py::index_error(
"Index must be between 0 and " +
std::to_string(dim));
return v[ii] = val;
})

.def("__len__", [](iv_type const &) { return dim; })
.def("__iter__", [](iv_type const & v) {
return py::make_iterator(v.begin(), v.end());
}, py::keep_alive<0, 1>()) /* Keep vector alive while iterator is used */

.def("__eq__",
py::overload_cast<int>(&iv_type::operator==, py::const_))
.def("__eq__",
py::overload_cast<const iv_type&>(&iv_type::operator==, py::const_))
.def("__ne__",
py::overload_cast<int>(&iv_type::operator!=, py::const_))
.def("__ne__",
py::overload_cast<const iv_type&>(&iv_type::operator!=, py::const_))
.def("__lt__", &iv_type::operator<)
.def("__le__", &iv_type::operator<=)
.def("__gt__", &iv_type::operator>)
.def("__ge__", &iv_type::operator>=)

.def("__add__",
py::overload_cast<int>(&iv_type::operator+, py::const_))
.def("__add__",
py::overload_cast<const iv_type&>(&iv_type::operator+, py::const_))
.def("__sub__",
py::overload_cast<int>(&iv_type::operator-, py::const_))
.def("__sub__",
py::overload_cast<const iv_type&>(&iv_type::operator-, py::const_))
.def("__mul__",
py::overload_cast<int>(&iv_type::operator*, py::const_))
.def("__mul__",
py::overload_cast<const iv_type&>(&iv_type::operator*, py::const_))
.def("__truediv__",
py::overload_cast<int>(&iv_type::operator/, py::const_))
.def("__truediv__",
py::overload_cast<const iv_type&>(&iv_type::operator/, py::const_))
.def("__iadd__",
py::overload_cast<int>(&iv_type::operator+=))
.def("__iadd__",
py::overload_cast<const iv_type&>(&iv_type::operator+=))
.def("__isub__",
py::overload_cast<int>(&iv_type::operator-=))
.def("__isub__",
py::overload_cast<const iv_type&>(&iv_type::operator-=))
.def("__imul__",
py::overload_cast<int>(&iv_type::operator*=))
.def("__imul__",
py::overload_cast<const iv_type&>(&iv_type::operator*=))
.def("__itruediv__",
py::overload_cast<int>(&iv_type::operator/=))
.def("__itruediv__",
py::overload_cast<const iv_type&>(&iv_type::operator/=))

.def("numpy",
[](const iv_type& iv) {
auto result = py::array(
py::buffer_info(
nullptr,
sizeof(int),
py::format_descriptor<int>::value,
1,
{ dim },
{ sizeof(int) }
));
auto buf = result.request();
int* ptr = static_cast<int*>(buf.ptr);
for (int i=0; i < dim; ++i)
ptr[i] = iv[0];

return result;
})
;

m.def("coarsen",
py::overload_cast<const iv_type&, const iv_type&>(&coarsen<dim>));
m.def("coarsen",
py::overload_cast<const Dim3&, const iv_type&>(&coarsen<dim>));
m.def("coarsen",
py::overload_cast<const iv_type&, int>(&coarsen<dim>));
m.def("refine",
py::overload_cast<const Dim3&, const iv_type&>(&refine<dim>));
}
}


void init_IntVect(py::module &m)
{
using namespace amrex;

py::class_< IntVect >(m, "IntVect")
.def("__repr__",
[](py::object& obj) {
py::str py_name = obj.attr("__class__").attr("__name__");
const std::string name = py_name;
const auto iv = obj.cast<IntVect>();
std::stringstream s;
s << iv;
return "<amrex." + name + " " + s.str() + ">";
}
)
.def("__str",
[](const IntVect& iv) {
std::stringstream s;
s << iv;
return s.str();
})

.def(py::init<>())
#if (AMREX_SPACEDIM > 1)
.def(py::init<AMREX_D_DECL(int, int, int)>())
#endif
.def(py::init<int>())
.def(py::init<const std::array<int, AMREX_SPACEDIM>&>())

.def_property_readonly("sum", &IntVect::sum)
.def_property_readonly("max",
py::overload_cast<>(&IntVect::max, py::const_))
.def_property_readonly("min",
py::overload_cast<>(&IntVect::min, py::const_))
.def_static("zero_vector", &IntVect::TheZeroVector)
.def_static("unit_vector", &IntVect::TheUnitVector)
.def_static("node_vector", &IntVect::TheNodeVector)
.def_static("cell_vector", &IntVect::TheCellVector)
.def_static("max_vector", &IntVect::TheMaxVector)
.def_static("min_vector", &IntVect::TheMinVector)

.def("dim3", &IntVect::dim3)
.def("__getitem__",
[](const IntVect& v, const int i) {
const int ii = (i >= 0) ? i : AMREX_SPACEDIM + i;
if ((ii < 0) || (ii >= AMREX_SPACEDIM))
throw py::index_error(
"Index must be between 0 and " +
std::to_string(AMREX_SPACEDIM));
return v[ii];
})
.def("__setitem__",
[](IntVect& v, const int i, const int& val) {
const int ii = (i >= 0) ? i : AMREX_SPACEDIM + i;
if ((ii < 0) || (ii >= AMREX_SPACEDIM))
throw py::index_error(
"Index must be between 0 and " +
std::to_string(AMREX_SPACEDIM));
return v[ii] = val;
})

.def("__len__", [](IntVect const &) { return AMREX_SPACEDIM; })
.def("__iter__", [](IntVect const & v) {
return py::make_iterator(v.begin(), v.end());
}, py::keep_alive<0, 1>()) /* Keep vector alive while iterator is used */

.def("__eq__",
py::overload_cast<int>(&IntVect::operator==, py::const_))
.def("__eq__",
py::overload_cast<const IntVect&>(&IntVect::operator==, py::const_))
.def("__ne__",
py::overload_cast<int>(&IntVect::operator!=, py::const_))
.def("__ne__",
py::overload_cast<const IntVect&>(&IntVect::operator!=, py::const_))
.def("__lt__", &IntVect::operator<)
.def("__le__", &IntVect::operator<=)
.def("__gt__", &IntVect::operator>)
.def("__ge__", &IntVect::operator>=)

.def("__add__",
py::overload_cast<int>(&IntVect::operator+, py::const_))
.def("__add__",
py::overload_cast<const IntVect&>(&IntVect::operator+, py::const_))
.def("__sub__",
py::overload_cast<int>(&IntVect::operator-, py::const_))
.def("__sub__",
py::overload_cast<const IntVect&>(&IntVect::operator-, py::const_))
.def("__mul__",
py::overload_cast<int>(&IntVect::operator*, py::const_))
.def("__mul__",
py::overload_cast<const IntVect&>(&IntVect::operator*, py::const_))
.def("__truediv__",
py::overload_cast<int>(&IntVect::operator/, py::const_))
.def("__truediv__",
py::overload_cast<const IntVect&>(&IntVect::operator/, py::const_))
.def("__iadd__",
py::overload_cast<int>(&IntVect::operator+=))
.def("__iadd__",
py::overload_cast<const IntVect&>(&IntVect::operator+=))
.def("__isub__",
py::overload_cast<int>(&IntVect::operator-=))
.def("__isub__",
py::overload_cast<const IntVect&>(&IntVect::operator-=))
.def("__imul__",
py::overload_cast<int>(&IntVect::operator*=))
.def("__imul__",
py::overload_cast<const IntVect&>(&IntVect::operator*=))
.def("__itruediv__",
py::overload_cast<int>(&IntVect::operator/=))
.def("__itruediv__",
py::overload_cast<const IntVect&>(&IntVect::operator/=))

.def("numpy",
[](const IntVect& iv) {
auto result = py::array(
py::buffer_info(
nullptr,
sizeof(int),
py::format_descriptor<int>::value,
1,
{ AMREX_SPACEDIM },
{ sizeof(int) }
));
auto buf = result.request();
int* ptr = static_cast<int*>(buf.ptr);
for (int i=0; i < AMREX_SPACEDIM; ++i)
ptr[i] = iv[0];

return result;
})
;

m.def("coarsen",
py::overload_cast<const IntVect&, const IntVect&>(&coarsen));
m.def("coarsen",
py::overload_cast<const Dim3&, const IntVect&>(&coarsen));
m.def("coarsen",
py::overload_cast<const IntVect&, int>(&coarsen));
m.def("refine",
py::overload_cast<const Dim3&, const IntVect&>(&refine));
init_IntVectND<1>(m);
init_IntVectND<2>(m);
init_IntVectND<3>(m);

make_Vector<IntVect> (m, "IntVect");
}

0 comments on commit 5070edc

Please sign in to comment.