Shamrock 2025.10.0
Astrophysical Code
Loading...
Searching...
No Matches
pyRamsesModel.cpp
Go to the documentation of this file.
1// -------------------------------------------------------//
2//
3// SHAMROCK code for hydrodynamics
4// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
5// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
6// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
7//
8// -------------------------------------------------------//
9
19
27#include <pybind11/functional.h>
28#include <pybind11/numpy.h>
29#include <memory>
30
32 template<class Tvec, class TgridVec>
33 void add_instance(py::module &m, std::string name_config, std::string name_model) {
34
35 using Tscal = shambase::VecComponent<Tvec>;
36 using Tgridscal = shambase::VecComponent<TgridVec>;
37
38 using T = Model<Tvec, TgridVec>;
39 using TConfig = typename T::Solver::Config;
40 using TAnalysisSodTube = shammodels::basegodunov::modules::AnalysisSodTube<Tvec, TgridVec>;
41
42 shamlog_debug_ln("[Py]", "registering class :", name_config, typeid(T).name());
43 shamlog_debug_ln("[Py]", "registering class :", name_model, typeid(T).name());
44
45 py::class_<TConfig> config_cls(m, name_config.c_str());
46
47 shammodels::common::add_json_defs<TConfig>(config_cls);
48
49 config_cls
50 .def(
51 "set_scale_factor",
52 [](TConfig &self, Tscal scale_factor) {
53 self.grid_coord_to_pos_fact = scale_factor;
54 })
55 .def(
56 "set_Csafe",
57 [](TConfig &self, Tscal Csafe) {
58 self.Csafe = Csafe;
59 })
60 .def(
61 "set_eos_gamma",
62 [](TConfig &self, Tscal eos_gamma) {
63 self.set_eos_gamma(eos_gamma);
64 })
65 .def(
66 "set_riemann_solver_hll",
67 [](TConfig &self) {
68 self.riemann_config = HLL;
69 })
70 .def(
71 "set_riemann_solver_hllc",
72 [](TConfig &self) {
73 self.riemann_config = HLLC;
74 })
75 .def(
76 "set_riemann_solver_rusanov",
77 [](TConfig &self) {
78 self.riemann_config = Rusanov;
79 })
80 .def(
81 "set_slope_lim_none",
82 [](TConfig &self) {
83 self.slope_config = None;
84 })
85 .def(
86 "set_slope_lim_vanleer_f",
87 [](TConfig &self) {
88 self.slope_config = VanLeer_f;
89 })
90 .def(
91 "set_slope_lim_vanleer_std",
92 [](TConfig &self) {
93 self.slope_config = VanLeer_std;
94 })
95 .def(
96 "set_slope_lim_vanleer_sym",
97 [](TConfig &self) {
98 self.slope_config = VanLeer_sym;
99 })
100 .def(
101 "set_slope_lim_minmod",
102 [](TConfig &self) {
103 self.slope_config = Minmod;
104 })
105 .def(
106 "set_scheduler_config",
107 [](TConfig &self, u64 split_crit, u64 merge_crit) {
108 self.scheduler_conf.split_load_value = split_crit;
109 self.scheduler_conf.merge_load_value = merge_crit;
110 },
111 py::kw_only(),
112 py::arg("split_load_value"),
113 py::arg("merge_load_value"))
114 .def(
115 "set_face_time_interpolation",
116 [](TConfig &self, bool face_time_interpolate) {
117 self.face_half_time_interpolation = face_time_interpolate;
118 })
119 .def(
120 "set_boundary_condition",
121 [](TConfig &self, const std::string &axis, const std::string &bc_type) {
122 BCConfig::GhostType ghost_type;
123 if (bc_type == "periodic") {
124 ghost_type = BCConfig::GhostType::Periodic;
125 } else if (bc_type == "reflective") {
126 ghost_type = BCConfig::GhostType::Reflective;
127 } else if (bc_type == "outflow") {
128 ghost_type = BCConfig::GhostType::Outflow;
129 } else {
130 throw std::invalid_argument(
131 "Unsupported boundary condition type: " + bc_type);
132 }
133
134 if (axis == "x") {
135 self.bc_config.set_x(ghost_type);
136 } else if (axis == "y") {
137 self.bc_config.set_y(ghost_type);
138 } else if (axis == "z") {
139 self.bc_config.set_z(ghost_type);
140 } else {
141 throw std::invalid_argument("Unsupported axis: " + axis);
142 }
143 },
144 py::arg("axis"),
145 py::arg("bc_type"))
146 .def(
147 "set_dust_mode_dhll",
148 [](TConfig &self, u32 ndust) {
149 self.dust_config = {DHLL, ndust};
150 })
151 .def(
152 "set_dust_mode_hb",
153 [](TConfig &self, u32 ndust) {
154 self.dust_config = {HB, ndust};
155 })
156 .def(
157 "set_dust_mode_none",
158 [](TConfig &self) {
159 self.dust_config = {NoDust, 0};
160 })
161 .def(
162 "set_alpha_values",
163 [](TConfig &self, f32 alpha_values) {
164 return self.set_alphas_static(alpha_values);
165 })
166 .def(
167 "set_drag_mode_no_drag",
168 [](TConfig &self) {
169 self.drag_config.drag_solver_config = NoDrag;
170 self.drag_config.enable_frictional_heating = false;
171 })
172 .def(
173 "set_drag_mode_irk1",
174 [](TConfig &self, bool frictional_status) {
175 self.drag_config.drag_solver_config = IRK1;
176 self.drag_config.enable_frictional_heating = frictional_status;
177 })
178 .def(
179 "set_drag_mode_irk2",
180 [](TConfig &self, bool frictional_status) {
181 self.drag_config.drag_solver_config = IRK2;
182 self.drag_config.enable_frictional_heating = frictional_status;
183 })
184 .def(
185 "set_drag_mode_expo",
186 [](TConfig &self, bool frictional_status) {
187 self.drag_config.drag_solver_config = EXPO;
188 self.drag_config.enable_frictional_heating = frictional_status;
189 })
190 .def(
191 "set_amr_mode_none",
192 [](TConfig &self) {
193 self.amr_mode.set_refine_none();
194 })
195 .def(
196 "set_amr_mode_density_based",
197 [](TConfig &self, Tscal crit_mass) {
198 self.amr_mode.set_refine_density_based(crit_mass);
199 },
200 py::kw_only(),
201 py::arg("crit_mass"))
202 .def(
203 "set_amr_mode_pseudo_gradient_based",
204 [](TConfig &self, Tscal error_min, Tscal error_max) {
205 self.amr_mode.set_refine_pseudo_gradient_based(error_min, error_max);
206 },
207 py::kw_only(),
208 py::arg("error_min"),
209 py::arg("error_max"))
210 .def(
211 "set_amr_mode_jeans_length_based",
212 [](TConfig &self, u32 N_jeans, Tscal T_init) {
213 self.amr_mode.set_refine_jeans_length_based(N_jeans, T_init);
214 },
215 py::kw_only(),
216 py::arg("N_jeans"),
217 py::arg("T_init"))
218 .def(
219 "set_amr_mode_shear_based",
220 [](TConfig &self, Tscal threshold) {
221 self.amr_mode.set_refine_shear_based(threshold);
222 },
223 py::kw_only(),
224 py::arg("Threshold"))
225 .def(
226 "set_amr_mode_old",
227 [](TConfig &self, bool use_old_amr) {
228 self.amr_mode.old_amr = use_old_amr;
229 })
230 .def(
231 "set_first_order_interpolation_mode",
232 [](TConfig &self) {
233 self.amr_interp_mode = FIRST_ORDER;
234 })
235 .def(
236 "set_second_order_interpolation_mode",
237 [](TConfig &self) {
238 self.amr_interp_mode = SECOND_ORDER;
239 })
240 .def(
241 "set_gravity_mode_no_gravity",
242 [](TConfig &self) {
243 self.gravity_config.gravity_mode = NoGravity;
244 })
245 .def(
246 "set_gravity_mode_cg",
247 [](TConfig &self) {
248 self.gravity_config.gravity_mode = CG;
249 })
250 .def(
251 "set_gravity_mode_pcg",
252 [](TConfig &self) {
253 self.gravity_config.gravity_mode = PCG;
254 })
255 .def(
256 "set_gravity_mode_bicgstab",
257 [](TConfig &self) {
258 self.gravity_config.gravity_mode = BICGSTAB;
259 })
260 .def("set_npscal_gas", [](TConfig &self, u32 npscal_gas) {
261 self.npscal_gas_config.npscal_gas = npscal_gas;
262 });
263
264 std::string sod_tube_analysis_name = name_model + "_AnalysisSodTube";
265 py::class_<TAnalysisSodTube>(m, sod_tube_analysis_name.c_str())
266 .def("compute_L2_dist", [](TAnalysisSodTube &self) -> std::tuple<Tscal, Tvec, Tscal> {
267 auto ret = self.compute_L2_dist();
268 return {ret.rho, ret.v, ret.P};
269 });
270
271 py::class_<T>(m, name_model.c_str())
272 .def("init", &T::init)
273 .def("init_scheduler", &T::init_scheduler)
274 .def("make_base_grid", &T::make_base_grid)
275 .def("dump_vtk", &T::dump_vtk)
276 .def("dump", &T::dump)
277 .def("load_from_dump", &T::load_from_dump)
278 .def("evolve_once_override_time", &T::evolve_once_time_expl)
279 .def("evolve_once", &T::evolve_once)
280 .def(
281 "evolve_until",
282 [](T &self, f64 target_time, i32 niter_max) {
283 return self.evolve_until(target_time, niter_max);
284 },
285 py::arg("target_time"),
286 py::kw_only(),
287 py::arg("niter_max") = -1)
288 .def("timestep", &T::timestep)
289 .def("solver_logs_last_rate", &T::solver_logs_last_rate)
290 .def("solver_logs_last_obj_count", &T::solver_logs_last_obj_count)
291 .def(
292 "solver_logs_last_system_metrics",
293 [](T &self) {
294 auto system_metrics = self.solver.solve_logs.get_last_system_metrics();
295 py::dict ret;
296 ret["duration"] = system_metrics.wall_time;
297 if (system_metrics.rank_energy_consummed.has_value()) {
298 ret["rank_energy_consummed"] = system_metrics.rank_energy_consummed.value();
299 }
300 if (system_metrics.gpu_energy_consummed.has_value()) {
301 ret["gpu_energy_consummed"] = system_metrics.gpu_energy_consummed.value();
302 }
303 if (system_metrics.cpu_energy_consummed.has_value()) {
304 ret["cpu_energy_consummed"] = system_metrics.cpu_energy_consummed.value();
305 }
306 if (system_metrics.dram_energy_consummed.has_value()) {
307 ret["dram_energy_consummed"] = system_metrics.dram_energy_consummed.value();
308 }
309 return ret;
310 })
311 .def(
312 "set_field_value_lambda_f64",
313 [](T &self,
314 std::string field_name,
315 const std::function<f64(Tvec, Tvec)> pos_to_val,
316 const i32 offset) {
317 return self.template set_field_value_lambda<f64>(
318 field_name, pos_to_val, offset);
319 },
320 py::arg("field_name"),
321 py::arg("pos_to_val"),
322 py::arg("offset") = 0)
323 .def(
324 "set_field_value_lambda_f64_3",
325 [](T &self,
326 std::string field_name,
327 const std::function<f64_3(Tvec, Tvec)> pos_to_val,
328 const i32 offset) {
329 return self.template set_field_value_lambda<f64_3>(
330 field_name, pos_to_val, offset);
331 },
332 py::arg("field_name"),
333 py::arg("pos_to_val"),
334 py::arg("offset") = 0)
335 .def(
336 "gen_default_config",
337 [](T &self) -> TConfig {
338 return TConfig();
339 })
340 .def(
341 "set_solver_config",
342 [](T &self, TConfig cfg) {
343 if (self.ctx.is_scheduler_initialized()) {
345 "Cannot change solver config after scheduler is initialized");
346 }
347 cfg.check_config();
348 self.solver.solver_config = cfg;
349 })
350 .def(
351 "get_cell_coords",
352 [](T &self, std::pair<TgridVec, TgridVec> block_coord, u32 cell_local_id) {
353 return self.get_cell_coords(block_coord, cell_local_id);
354 })
355 .def(
356 "make_analysis_sodtube",
357 [](T &self,
358 shamphys::SodTube sod,
359 Tvec direction,
360 Tscal time_val,
361 Tscal x_ref,
362 Tscal x_min,
363 Tscal x_max) {
364 return std::make_unique<TAnalysisSodTube>(
365 self.ctx,
366 self.solver.solver_config,
367 self.solver.storage,
368 sod,
369 direction,
370 time_val,
371 x_ref,
372 x_min,
373 x_max);
374 })
375 .def(
376 "add_timestep_callback",
377 [](T &self,
378 std::optional<std::function<void(void)>> step_begin_callback,
379 std::optional<std::function<void(void)>> step_end_callback) {
380 self.solver.timestep_callbacks.push_back(
381 {std::move(step_begin_callback), std::move(step_end_callback)});
382 },
383 py::kw_only(),
384 py::arg("step_begin") = std::nullopt,
385 py::arg("step_end") = std::nullopt)
386 .def(
387 "get_solver_tex",
388 [](T &self) {
389 return shambase::get_check_ref(self.solver.storage.solver_sequence).get_tex();
390 })
391 .def(
392 "get_solver_dot_graph",
393 [](T &self) {
394 return shambase::get_check_ref(self.solver.storage.solver_sequence)
395 .get_dot_graph();
396 })
397 .def(
398 "render_slice",
399 [](T &self, std::string name, std::string field_type, std::vector<Tvec> positions)
400 -> std::variant<std::vector<f64>, std::vector<f64_3>> {
401 if (field_type == "f64") {
402 ramses::modules::GridRender<Tvec, TgridVec, f64> render(
403 self.ctx, self.solver.solver_config, self.solver.storage);
404 return render.compute_slice(name, positions).copy_to_stdvec();
405 }
406
407 if (field_type == "f64_3") {
408 ramses::modules::GridRender<Tvec, TgridVec, f64_3> render(
409 self.ctx, self.solver.solver_config, self.solver.storage);
410 return render.compute_slice(name, positions).copy_to_stdvec();
411 }
412
413 throw shambase::make_except_with_loc<std::runtime_error>("unknown field type");
414 })
415 .def(
416 "get_time",
417 [](T &self) {
418 return self.solver.get_time();
419 })
420 .def(
421 "get_dt",
422 [](T &self) {
423 return self.solver.get_dt();
424 })
425 .def(
426 "set_time",
427 [](T &self, Tscal t) {
428 return self.solver.set_time(t);
429 })
430 .def("set_next_dt", [](T &self, Tscal dt) {
431 return self.solver.set_next_dt(dt);
432 });
433 }
434} // namespace shammodels::basegodunov
435
437 auto &m = root_module;
438
439 py::module mramses = m.def_submodule("model_ramses", "Shamrock Ramses solver");
440
441 std::string base_name = "RamsesModel";
442 using namespace shammodels::basegodunov;
443
444 add_instance<f64_3, i64_3>(
445 mramses, base_name + "_f64_3_i64_3_SolverConfig", base_name + "_f64_3_i64_3_Model");
446
447 using VariantAMRGodunovBind = std::variant<std::unique_ptr<Model<f64_3, i64_3>>>;
448
449 m.def(
450 "get_Model_Ramses",
451 [](ShamrockCtx &ctx,
452 std::string vector_type,
453 std::string grid_repr) -> VariantAMRGodunovBind {
454 VariantAMRGodunovBind ret;
455
456 if (vector_type == "f64_3" && grid_repr == "i64_3") {
457 ret = std::make_unique<Model<f64_3, i64_3>>(ctx);
458 } else {
460 "unknown combination of representation and grid_repr");
461 }
462
463 return ret;
464 },
465 py::kw_only(),
466 py::arg("context"),
467 py::arg("vector_type"),
468 py::arg("grid_repr"));
469}
double f64
Alias for double.
float f32
Alias for float.
std::uint32_t u32
32 bit unsigned integer
std::uint64_t u64
64 bit unsigned integer
std::int32_t i32
32 bit integer
void throw_with_loc(std::string message, SourceLocation loc=SourceLocation{})
Throw an exception and append the source location to it.
T & get_check_ref(const std::unique_ptr< T > &ptr, SourceLocation loc=SourceLocation())
Takes a std::unique_ptr and returns a reference to the object it holds. It throws a std::runtime_erro...
Definition memory.hpp:112
ExcptTypes make_except_with_loc(std::string message, SourceLocation loc=SourceLocation{})
Create an exception with a message and a location.
namespace for the basegodunov model
@ HB
Huang and Bai. Pressureless Riemann solver by Huang and Bai (2022) in Athena++.
@ DHLL
Dust HLL. This is merely the HLL solver for dust. It's then a Rusanov like.
@ NoDust
No dust, so no Riemann solver is used.
Pybind11 include and definitions.
#define ON_PYTHON_INIT
Register a Python module init function using static initialization.
Utilities to convert JSON objects to Python objects and vice versa. TODO: try to convert directly wit...