diff --git a/examples/ramses/run_advect.py b/examples/ramses/run_advect.py index f06d2f163b..b3b2f222d9 100644 --- a/examples/ramses/run_advect.py +++ b/examples/ramses/run_advect.py @@ -77,26 +77,40 @@ def run_advect(slope_limiter: str, riemann_solver: str, only_last_step: bool = T model.init_scheduler(int(1e7), 1) model.make_base_grid((0, 0, 0), (sz, sz, sz), (base * multx, base * multy, base * multz)) - def rho_map(rmin, rmax): - x, y, z = rmin - + def rho_x(x): if x < 0.6 and x > 0.4: return 2 + return 1.0 + + def v_x(x): + return 1.0 + def e_x(): return 1.0 - def rhoe_map(rmin, rmax): - rho = rho_map(rmin, rmax) - return 1.0 * rho + def setter(patchdata): + # will return (ncell, 3) as shape + # (cell_center does not exist so it will be computed on the fly) + cell_center = patchdata.get("cell_center") + + x = cell_center[:, 0] + rho = np.where((x < 0.6) & (x > 0.4), 2.0, 1.0) + + ncell = cell_center.shape[0] + v = np.zeros((ncell, 3)) + v[:, 0] = 1.0 + e = np.ones(ncell) + + rho_v = rho[:, None] * v + rho_e = rho * e - def rhovel_map(rmin, rmax): - x, y, z = rmin - rho = rho_map(rmin, rmax) - return (1 * rho, 0 * rho, 0 * rho) + patchdata.set("rho", rho) # will set rho (f64) + patchdata.set("rhovel", rho_v) # will set rhovel (f64_3) + patchdata.set("rhoetot", rho_e) # will set rhoetot (f64) - model.set_field_value_lambda_f64("rho", rho_map) - model.set_field_value_lambda_f64("rhoetot", rhoe_map) - model.set_field_value_lambda_f64_3("rhovel", rhovel_map) + # Will call the function passed as argument for each patchdata + # patchdata are instances of a custom PatchDataSetup class used to have getters and setters + model.update_fields(setter) results = [] diff --git a/src/shammodels/ramses/CMakeLists.txt b/src/shammodels/ramses/CMakeLists.txt index eca6794e83..5fc3b48db8 100644 --- a/src/shammodels/ramses/CMakeLists.txt +++ b/src/shammodels/ramses/CMakeLists.txt @@ -62,6 +62,7 @@ target_link_libraries(shammodels_ramses PUBLIC shammath) target_link_libraries(shammodels_ramses PUBLIC shamphys) target_link_libraries(shammodels_ramses PUBLIC shamsys) target_link_libraries(shammodels_ramses PUBLIC shammodels_common) +target_link_libraries(shammodels_ramses PUBLIC shampylib) target_link_libraries(shammodels_ramses PUBLIC nlohmann_json::nlohmann_json) target_include_directories( diff --git a/src/shammodels/ramses/include/shammodels/ramses/Model.hpp b/src/shammodels/ramses/include/shammodels/ramses/Model.hpp index 4d5e2e666c..29fb22f05f 100644 --- a/src/shammodels/ramses/include/shammodels/ramses/Model.hpp +++ b/src/shammodels/ramses/include/shammodels/ramses/Model.hpp @@ -21,12 +21,16 @@ #include "shambase/memory.hpp" #include "shambackends/vec.hpp" #include "shammodels/ramses/Solver.hpp" +#include "shampylib/PatchDataSetup.hpp" +#include "shampylib/PatchDataToPy.hpp" #include "shamrock/amr/AMRGrid.hpp" #include "shamrock/io/ShamrockDump.hpp" #include "shamrock/patch/PatchDataLayer.hpp" #include "shamrock/scheduler/ReattributeDataUtility.hpp" #include "shamrock/scheduler/ShamrockCtx.hpp" #include "shamtree/kernels/geometry_utils.hpp" +#include +#include namespace shammodels::basegodunov { @@ -104,6 +108,53 @@ namespace shammodels::basegodunov { }); } + /** + * @brief Call a Python callback once per owned patch to get/set float64 fields. + * + * Registers all f64 / f64_3 layout fields plus a virtual "cell_center" getter. + */ + inline void update_fields(const std::function &updater) { + StackEntry stack_loc{}; + + using Block = typename Solver::Config::AMRBlock; + + PatchScheduler &sched = shambase::get_check_ref(ctx.sched); + sched.patch_data.for_each_patchdata([&](u64 /*patch_id*/, + shamrock::patch::PatchDataLayer &pdat) { + shamrock::PatchDataSetup setup; + shamrock::register_f64_layout_fields(setup, pdat); + + sham::DeviceBuffer &buf_cell_min = pdat.get_field_buf_ref(0); + sham::DeviceBuffer &buf_cell_max = pdat.get_field_buf_ref(1); + + Tscal scale_factor = solver.solver_config.grid_coord_to_pos_fact; + + setup.register_getter("cell_center", [&]() -> py::array_t { + auto cell_min = buf_cell_min.copy_to_stdvec(); + auto cell_max = buf_cell_max.copy_to_stdvec(); + + u32 ncell = pdat.get_obj_cnt() * Block::block_size; + std::vector centers(ncell); + + for (u32 i = 0; i < pdat.get_obj_cnt(); i++) { + Tvec block_min = cell_min[i].template convert() * scale_factor; + Tvec block_max = cell_max[i].template convert() * scale_factor; + Tvec delta_cell = (block_max - block_min) / Block::side_size; + + Block::for_each_cell_in_block(delta_cell, [&](u32 lid, Tvec delta) { + Tvec bmin = block_min + delta; + Tvec bmax = bmin + delta_cell; + centers[i * Block::block_size + lid] = (bmin + bmax) * Tscal{0.5}; + }); + } + + return shamrock::VecToNumpy::convert(centers); + }); + + updater(setup); + }); + } + inline std::pair get_cell_coords( std::pair block_coords, u32 lid) { using Block = typename Solver::Config::AMRBlock; diff --git a/src/shammodels/ramses/src/pyRamsesModel.cpp b/src/shammodels/ramses/src/pyRamsesModel.cpp index d52a194d6d..f36bdced64 100644 --- a/src/shammodels/ramses/src/pyRamsesModel.cpp +++ b/src/shammodels/ramses/src/pyRamsesModel.cpp @@ -24,6 +24,7 @@ #include "shammodels/ramses/Solver.hpp" #include "shammodels/ramses/modules/AnalysisSodTube.hpp" #include "shammodels/ramses/modules/render/GridRender.hpp" +#include "shampylib/PatchDataSetup.hpp" #include #include #include @@ -310,6 +311,11 @@ namespace shammodels::basegodunov { py::arg("field_name"), py::arg("pos_to_val"), py::arg("offset") = 0) + .def( + "update_fields", + [](T &self, const std::function &cb) { + self.update_fields(cb); + }) .def( "gen_default_config", [](T &self) -> TConfig { diff --git a/src/shampylib/include/shampylib/PatchDataSetup.hpp b/src/shampylib/include/shampylib/PatchDataSetup.hpp new file mode 100644 index 0000000000..25ee7a781e --- /dev/null +++ b/src/shampylib/include/shampylib/PatchDataSetup.hpp @@ -0,0 +1,69 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +#pragma once + +/** + * @file PatchDataSetup.hpp + * @author Timothée David--Cléris (tim.shamrock@proton.me) + * @brief Proxy for patch field get/set during Python IC / field setup. + */ + +#include "shambase/exception.hpp" +#include "shambase/string.hpp" +#include "shambackends/typeAliasVec.hpp" +#include +#include +#include +#include +#include + +namespace py = pybind11; + +namespace shamrock { + + /** + * @brief Thin proxy over named field getters/setters as float64 numpy arrays. + * + * Does not own PatchData. Real fields and virtual (computed) fields share the + * same get/set path via registered lambdas. + */ + class PatchDataSetup { + std::unordered_map()>> getters; + std::unordered_map)>> setters; + + public: + void register_getter(std::string name, std::function()> fn) { + getters[std::move(name)] = std::move(fn); + } + + void register_setter(std::string name, std::function)> fn) { + setters[std::move(name)] = std::move(fn); + } + + py::array_t get(const std::string &name) const { + auto it = getters.find(name); + if (it == getters.end()) { + throw shambase::make_except_with_loc(shambase::format( + "PatchDataSetup: no getter registered for field \"{}\"", name)); + } + return it->second(); + } + + void set(const std::string &name, py::array_t value) const { + auto it = setters.find(name); + if (it == setters.end()) { + throw shambase::make_except_with_loc(shambase::format( + "PatchDataSetup: no setter registered for field \"{}\"", name)); + } + it->second(std::move(value)); + } + }; + +} // namespace shamrock diff --git a/src/shampylib/include/shampylib/PatchDataToPy.hpp b/src/shampylib/include/shampylib/PatchDataToPy.hpp index 5c985a75ad..08e4cac96f 100644 --- a/src/shampylib/include/shampylib/PatchDataToPy.hpp +++ b/src/shampylib/include/shampylib/PatchDataToPy.hpp @@ -15,8 +15,11 @@ * @brief */ +#include "shambase/exception.hpp" +#include "shambase/string.hpp" #include "shambindings/pybind11_stl.hpp" #include "shambindings/pybindaliases.hpp" +#include "shampylib/PatchDataSetup.hpp" #include "shamrock/patch/PatchDataLayer.hpp" #include #include @@ -248,4 +251,89 @@ namespace shamrock { return dic_out; } + + template + class NumpyToVec; + + template<> + class NumpyToVec { + public: + static std::vector convert(py::array_t arr) { + if (arr.ndim() != 1) { + throw shambase::make_except_with_loc( + shambase::format("expected 1D array for f64 field, got ndim={}", arr.ndim())); + } + + auto r = arr.unchecked<1>(); + std::vector vec(static_cast(r.shape(0))); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + vec[static_cast(i)] = r(i); + } + return vec; + } + }; + + template<> + class NumpyToVec { + public: + static std::vector convert(py::array_t arr) { + if (arr.ndim() != 2 || arr.shape(1) != 3) { + throw shambase::make_except_with_loc( + "expected (N, 3) array for f64_3 field"); + } + + auto r = arr.unchecked<2>(); + std::vector vec(static_cast(r.shape(0))); + for (py::ssize_t i = 0; i < r.shape(0); i++) { + vec[static_cast(i)] = f64_3{r(i, 0), r(i, 1), r(i, 2)}; + } + return vec; + } + }; + + inline void register_field_io_f64(PatchDataSetup &setup, PatchDataField &field) { + std::string name = field.get_name(); + setup.register_getter(name, [&field]() -> py::array_t { + return VecToNumpy::convert(field.get_buf().copy_to_stdvec()); + }); + setup.register_setter(name, [&field](py::array_t arr) { + auto vec = NumpyToVec::convert(arr); + if (vec.size() != field.get_val_cnt()) { + throw shambase::make_except_with_loc(shambase::format( + "field \"{}\": array size {} does not match field val_cnt {}", + field.get_name(), + vec.size(), + field.get_val_cnt())); + } + field.get_buf().copy_from_stdvec(vec); + }); + } + + inline void register_field_io_f64_3(PatchDataSetup &setup, PatchDataField &field) { + std::string name = field.get_name(); + setup.register_getter(name, [&field]() -> py::array_t { + return VecToNumpy::convert(field.get_buf().copy_to_stdvec()); + }); + setup.register_setter(name, [&field](py::array_t arr) { + auto vec = NumpyToVec::convert(arr); + if (vec.size() != field.get_val_cnt()) { + throw shambase::make_except_with_loc(shambase::format( + "field \"{}\": array size {} does not match field val_cnt {}", + field.get_name(), + vec.size(), + field.get_val_cnt())); + } + field.get_buf().copy_from_stdvec(vec); + }); + } + + inline void register_f64_layout_fields( + PatchDataSetup &setup, shamrock::patch::PatchDataLayer &pdat) { + pdat.for_each_field([&](auto &field) { + register_field_io_f64(setup, field); + }); + pdat.for_each_field([&](auto &field) { + register_field_io_f64_3(setup, field); + }); + } } // namespace shamrock diff --git a/src/shampylib/src/pyPatchDataSetup.cpp b/src/shampylib/src/pyPatchDataSetup.cpp new file mode 100644 index 0000000000..971bf72677 --- /dev/null +++ b/src/shampylib/src/pyPatchDataSetup.cpp @@ -0,0 +1,28 @@ +// -------------------------------------------------------// +// +// SHAMROCK code for hydrodynamics +// Copyright (c) 2021-2026 Timothée David--Cléris +// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1 +// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information +// +// -------------------------------------------------------// + +/** + * @file pyPatchDataSetup.cpp + * @author Timothée David--Cléris (tim.shamrock@proton.me) + * @brief Python bindings for PatchDataSetup + */ + +#include "shambindings/pybindaliases.hpp" +#include "shamcomm/logs.hpp" +#include "shampylib/PatchDataSetup.hpp" + +ON_PYTHON_INIT { + auto &m = root_module; + + shamlog_debug_ln("[Py]", "registering shamrock.PatchDataSetup"); + + py::class_(m, "PatchDataSetup") + .def("get", &shamrock::PatchDataSetup::get, py::arg("name")) + .def("set", &shamrock::PatchDataSetup::set, py::arg("name"), py::arg("value")); +}