[Ramses] draft new setup routine for grid solver

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 = []
 

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(

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 <pybind11/functional.h>
+#include <functional>
 
 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<void(shamrock::PatchDataSetup &)> &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<TgridVec> &buf_cell_min = pdat.get_field_buf_ref<TgridVec>(0);
+                sham::DeviceBuffer<TgridVec> &buf_cell_max = pdat.get_field_buf_ref<TgridVec>(1);
+
+                Tscal scale_factor = solver.solver_config.grid_coord_to_pos_fact;
+
+                setup.register_getter("cell_center", [&]() -> py::array_t<f64> {
+                    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<Tvec> centers(ncell);
+
+                    for (u32 i = 0; i < pdat.get_obj_cnt(); i++) {
+                        Tvec block_min  = cell_min[i].template convert<Tscal>() * scale_factor;
+                        Tvec block_max  = cell_max[i].template convert<Tscal>() * 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<Tvec>::convert(centers);
+                });
+
+                updater(setup);
+            });
+        }
+
         inline std::pair<Tvec, Tvec> get_cell_coords(
             std::pair<TgridVec, TgridVec> block_coords, u32 lid) {
             using Block = typename Solver::Config::AMRBlock;

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 <pybind11/functional.h>
 #include <pybind11/numpy.h>
 #include <memory>
@@ -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<void(shamrock::PatchDataSetup &)> &cb) {
+                    self.update_fields(cb);
+                })
             .def(
                 "gen_default_config",
                 [](T &self) -> TConfig {

src/shampylib/include/shampylib/PatchDataSetup.hpp

@@ -0,0 +1,69 @@
+// -------------------------------------------------------//
+//
+// SHAMROCK code for hydrodynamics
+// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
+// 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 <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <unordered_map>
+#include <functional>
+#include <string>
+
+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<std::string, std::function<py::array_t<f64>()>> getters;
+        std::unordered_map<std::string, std::function<void(py::array_t<f64>)>> setters;
+
+        public:
+        void register_getter(std::string name, std::function<py::array_t<f64>()> fn) {
+            getters[std::move(name)] = std::move(fn);
+        }
+
+        void register_setter(std::string name, std::function<void(py::array_t<f64>)> fn) {
+            setters[std::move(name)] = std::move(fn);
+        }
+
+        py::array_t<f64> get(const std::string &name) const {
+            auto it = getters.find(name);
+            if (it == getters.end()) {
+                throw shambase::make_except_with_loc<std::invalid_argument>(shambase::format(
+                    "PatchDataSetup: no getter registered for field \"{}\"", name));
+            }
+            return it->second();
+        }
+
+        void set(const std::string &name, py::array_t<f64> value) const {
+            auto it = setters.find(name);
+            if (it == setters.end()) {
+                throw shambase::make_except_with_loc<std::invalid_argument>(shambase::format(
+                    "PatchDataSetup: no setter registered for field \"{}\"", name));
+            }
+            it->second(std::move(value));
+        }
+    };
+
+} // namespace shamrock

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 <pybind11/numpy.h>
 #include <pybind11/pybind11.h>
@@ -248,4 +251,89 @@ namespace shamrock {
 
         return dic_out;
     }
+
+    template<class T>
+    class NumpyToVec;
+
+    template<>
+    class NumpyToVec<f64> {
+        public:
+        static std::vector<f64> convert(py::array_t<f64> arr) {
+            if (arr.ndim() != 1) {
+                throw shambase::make_except_with_loc<std::invalid_argument>(
+                    shambase::format("expected 1D array for f64 field, got ndim={}", arr.ndim()));
+            }
+
+            auto r = arr.unchecked<1>();
+            std::vector<f64> vec(static_cast<size_t>(r.shape(0)));
+            for (py::ssize_t i = 0; i < r.shape(0); i++) {
+                vec[static_cast<size_t>(i)] = r(i);
+            }
+            return vec;
+        }
+    };
+
+    template<>
+    class NumpyToVec<f64_3> {
+        public:
+        static std::vector<f64_3> convert(py::array_t<f64> arr) {
+            if (arr.ndim() != 2 || arr.shape(1) != 3) {
+                throw shambase::make_except_with_loc<std::invalid_argument>(
+                    "expected (N, 3) array for f64_3 field");
+            }
+
+            auto r = arr.unchecked<2>();
+            std::vector<f64_3> vec(static_cast<size_t>(r.shape(0)));
+            for (py::ssize_t i = 0; i < r.shape(0); i++) {
+                vec[static_cast<size_t>(i)] = f64_3{r(i, 0), r(i, 1), r(i, 2)};
+            }
+            return vec;
+        }
+    };
+
+    inline void register_field_io_f64(PatchDataSetup &setup, PatchDataField<f64> &field) {
+        std::string name = field.get_name();
+        setup.register_getter(name, [&field]() -> py::array_t<f64> {
+            return VecToNumpy<f64>::convert(field.get_buf().copy_to_stdvec());
+        });
+        setup.register_setter(name, [&field](py::array_t<f64> arr) {
+            auto vec = NumpyToVec<f64>::convert(arr);
+            if (vec.size() != field.get_val_cnt()) {
+                throw shambase::make_except_with_loc<std::invalid_argument>(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<f64_3> &field) {
+        std::string name = field.get_name();
+        setup.register_getter(name, [&field]() -> py::array_t<f64> {
+            return VecToNumpy<f64_3>::convert(field.get_buf().copy_to_stdvec());
+        });
+        setup.register_setter(name, [&field](py::array_t<f64> arr) {
+            auto vec = NumpyToVec<f64_3>::convert(arr);
+            if (vec.size() != field.get_val_cnt()) {
+                throw shambase::make_except_with_loc<std::invalid_argument>(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<f64>([&](auto &field) {
+            register_field_io_f64(setup, field);
+        });
+        pdat.for_each_field<f64_3>([&](auto &field) {
+            register_field_io_f64_3(setup, field);
+        });
+    }
 } // namespace shamrock

src/shampylib/src/pyPatchDataSetup.cpp

@@ -0,0 +1,28 @@
+// -------------------------------------------------------//
+//
+// SHAMROCK code for hydrodynamics
+// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
+// 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_<shamrock::PatchDataSetup>(m, "PatchDataSetup")
+        .def("get", &shamrock::PatchDataSetup::get, py::arg("name"))
+        .def("set", &shamrock::PatchDataSetup::set, py::arg("name"), py::arg("value"));
+}