Skip to content

[Ramses] allow setting field values in bulk (with numba) #1676

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
tdavidcl wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

[Ramses] allow setting field values in bulk (with numba) #1676

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 94 additions & 2 deletions examples/ramses/run_kh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
from numba import njit

import shamrock

Expand All @@ -25,7 +26,7 @@
# Setup parameters

# plot
nx, ny = 512, 512
nx, ny = 1024, 1024

# Physical parameters
vslip = 1 # slip speed between the two layers
Expand All @@ -48,7 +49,7 @@
multz = 1

sz = 1 << 1
base = 32
base = 64

sim_folder = "_to_trash/ramses_kh/"

Expand Down Expand Up @@ -211,10 +212,101 @@ def rhoetot_map(rmin, rmax):

return rhoekin + P / (gamma - 1)

print("no @njit")
model.set_field_value_lambda_f64("rho", rho_map)
model.set_field_value_lambda_f64("rhoetot", rhoetot_map)
model.set_field_value_lambda_f64_3("rhovel", rhovel_map)

rho_map = njit(rho_map)
rhovel_map = njit(rhovel_map)
rhoetot_map = njit(rhoetot_map)

print("with @njit")
model.set_field_value_lambda_f64("rho", rho_map)
model.set_field_value_lambda_f64("rhoetot", rhoetot_map)
model.set_field_value_lambda_f64_3("rhovel", rhovel_map)

def rho_map_all(in_min_cell: np.ndarray, in_max_cell: np.ndarray) -> np.array:
ncells = in_min_cell.shape[0]
rho = np.zeros(ncells)
for i in range(ncells):
x, y, z = in_min_cell[i]
if y > y_interface:
rho[i] = rho_2
else:
rho[i] = rho_1
Comment on lines +231 to +237
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The for loop used to populate the rho array can be replaced by a more idiomatic and efficient vectorized NumPy operation. Using np.where will improve both performance and code readability.

        y = in_min_cell[:, 1]
        rho = np.where(y > y_interface, rho_2, rho_1)

return rho

def rhovel_map_all(in_min_cell: np.ndarray, in_max_cell: np.ndarray) -> np.array:
rho = rho_map_all(in_min_cell, in_max_cell)
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This function recalculates rho by calling rho_map_all, which is inefficient, especially since rhoetot_map_all will also call this function, leading to redundant computations. Consider calculating rho once and passing it as an argument if these functions are to be used together.

Additionally, the loop over ncells can be vectorized using NumPy for significant performance improvements, which aligns with the goal of using numba.


ncells = in_min_cell.shape[0]
rhovel = np.zeros((ncells, 3))
for i in range(ncells):
x, y, z = in_min_cell[i, 0], in_min_cell[i, 1], in_min_cell[i, 2]

ampl = 0.01
n = 2
pert = np.sin(2 * np.pi * n * x / (xs))

sigma = 0.05 / (2**0.5)
gauss1 = np.exp(-((y - y_interface) ** 2) / (2 * sigma * sigma))
gauss2 = np.exp(-((y + y_interface) ** 2) / (2 * sigma * sigma))
pert *= gauss1 + gauss2

# Alternative formula (See T. Tricco paper)
# interf_sz = ys/32
# vx = np.arctan(y/interf_sz)/np.pi

vx = 0
if np.abs(y) > y_interface:
vx = vslip / 2
else:
vx = -vslip / 2

rhovel[i, 0] = vx * rho[i]
rhovel[i, 1] = ampl * pert * rho[i]
rhovel[i, 2] = 0.0
return rhovel

def rhoetot_map_all(in_min_cell: np.ndarray, in_max_cell: np.ndarray) -> np.array:
rho = rho_map_all(in_min_cell, in_max_cell)
rhovel = rhovel_map_all(in_min_cell, in_max_cell)

ncells = in_min_cell.shape[0]
rhoekin = np.zeros(ncells)

for i in range(ncells):
x, y, z = in_min_cell[i]

rhovel2 = (
rhovel[i][0] * rhovel[i][0]
+ rhovel[i][1] * rhovel[i][1]
+ rhovel[i][2] * rhovel[i][2]
)
rhoekin[i] = 0.5 * rhovel2 / rho[i]

if y > y_interface:
P = P_2
else:
P = P_1

return rhoekin + P / (gamma - 1)
Comment on lines +276 to +294
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Feb 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

This function contains a critical bug: the return statement is inside the for loop, which will cause it to exit after the first iteration with partially computed data. This leads to incorrect results.

Additionally, the function can be fully vectorized using NumPy to improve performance and readability, which also resolves the bug. The current implementation also performs redundant calculations by calling rho_map_all and rhovel_map_all.

        y = in_min_cell[:, 1]
        rhovel2 = np.sum(rhovel**2, axis=1)
        rhoekin = 0.5 * rhovel2 / rho

        P = np.where(y > y_interface, P_2, P_1)

        return rhoekin + P / (gamma - 1)


print("bulk")
model.set_field_value_lambda_f64_all("rho", rho_map_all)
model.set_field_value_lambda_f64_all("rhoetot", rhoetot_map_all)
model.set_field_value_lambda_f64_3_all("rhovel", rhovel_map_all)

rho_map_all = njit(rho_map_all)
rhovel_map_all = njit(rhovel_map_all)
rhoetot_map_all = njit(rhoetot_map_all)

print("bulk @njit")
model.set_field_value_lambda_f64_all("rho", rho_map_all)
model.set_field_value_lambda_f64_all("rhoetot", rhoetot_map_all)
model.set_field_value_lambda_f64_3_all("rhovel", rhovel_map_all)

# model.evolve_once(0,0.1)
fact = 25
tmax = 0.127 * fact
Expand Down
54 changes: 51 additions & 3 deletions src/shammodels/ramses/include/shammodels/ramses/Model.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,10 @@ namespace shammodels::basegodunov {
void dump_vtk(std::string filename);

template<class T>
inline void set_field_value_lambda(
inline void set_field_value_lambda_all(
std::string field_name,
const std::function<T(Tvec, Tvec)> pos_to_val,
const std::function<
std::vector<T>(const std::vector<Tvec> &, const std::vector<Tvec> &)> pos_to_val,
const i32 offset) {

StackEntry stack_loc{};
Expand All @@ -78,7 +79,34 @@ namespace shammodels::basegodunov {
auto cell_min = buf_cell_min.copy_to_stdvec();
auto cell_max = buf_cell_max.copy_to_stdvec();

std::vector<Tvec> in_min_cell(pdat.get_obj_cnt() * Block::block_size);
std::vector<Tvec> in_max_cell(pdat.get_obj_cnt() * Block::block_size);

Tscal scale_factor = solver.solver_config.grid_coord_to_pos_fact;
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;
in_min_cell[i * Block::block_size + lid] = bmin;
in_max_cell[i * Block::block_size + lid] = bmin + delta_cell;
});
}

shambase::Timer timer;
timer.start();
std::vector<T> in_values = pos_to_val(in_min_cell, in_max_cell);
timer.end();
logger::info_ln(
"Godunov",
shambase::format(
"Time to compute in_values: {} s for {} objects for field {}",
timer.elasped_sec(),
pdat.get_obj_cnt(),
field_name));

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;
Expand All @@ -87,14 +115,34 @@ namespace shammodels::basegodunov {
Block::for_each_cell_in_block(delta_cell, [&](u32 lid, Tvec delta) {
Tvec bmin = block_min + delta;
acc[(i * Block::block_size + lid) * f_nvar + offset]
= pos_to_val(bmin, bmin + delta_cell);
= in_values[i * Block::block_size + lid];
});
}

f.get_buf().copy_from_stdvec(acc);
});
}

template<class T>
inline void set_field_value_lambda(
std::string field_name,
const std::function<T(Tvec, Tvec)> pos_to_val,
const i32 offset) {

StackEntry stack_loc{};

auto lambda = [&](const std::vector<Tvec> &in_min_cell,
const std::vector<Tvec> &in_max_cell) -> std::vector<T> {
std::vector<T> in_values(in_min_cell.size());
for (u32 i = 0; i < in_min_cell.size(); i++) {
in_values[i] = pos_to_val(in_min_cell[i], in_max_cell[i]);
}
return in_values;
};

set_field_value_lambda_all<T>(field_name, lambda, offset);
}

inline std::pair<Tvec, Tvec> get_cell_coords(
std::pair<TgridVec, TgridVec> block_coords, u32 lid) {
using Block = typename Solver::Config::AMRBlock;
Expand Down
99 changes: 99 additions & 0 deletions src/shammodels/ramses/src/pyRamsesModel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,44 @@
#include <pybind11/numpy.h>
#include <memory>

// Should be moved to a global utility
template<class T>
class VecToNumpy {
public:
static py::array_t<T> convert(const std::vector<T> &vec) {

u32 len = vec.size();

py::array_t<T> ret({len});
auto r = ret.mutable_unchecked();

for (u32 i = 0; i < len; i++) {
r(i) = vec[i];
}

return std::move(ret);
}
};

template<class T>
class VecToNumpy<sycl::vec<T, 3>> {
public:
static py::array_t<T> convert(const std::vector<sycl::vec<T, 3>> &vec) {

u32 len = vec.size();

py::array_t<T> ret({len, 3U});
auto r = ret.mutable_unchecked();

for (u32 i = 0; i < len; i++) {
r(i, 0) = vec[i].x();
r(i, 1) = vec[i].y();
r(i, 2) = vec[i].z();
}
return std::move(ret);
}
};

namespace shammodels::basegodunov {
template<class Tvec, class TgridVec>
void add_instance(py::module &m, std::string name_config, std::string name_model) {
Expand Down Expand Up @@ -264,6 +302,67 @@ namespace shammodels::basegodunov {
py::arg("field_name"),
py::arg("pos_to_val"),
py::arg("offset") = 0)
.def(
"set_field_value_lambda_f64_all",
[](T &self,
std::string field_name,
const std::function<pybind11::array_t<f64>(
const pybind11::array_t<f64> &, const pybind11::array_t<f64> &)> pos_to_val,
const i32 offset) {
return self.template set_field_value_lambda_all<f64>(
field_name,
[&](const std::vector<f64_3> &in_min_cell,
const std::vector<f64_3> &in_max_cell) {
auto flattened_in_min_cell_arr
= VecToNumpy<f64_3>::convert(in_min_cell);
auto flattened_in_max_cell_arr
= VecToNumpy<f64_3>::convert(in_max_cell);

pybind11::array_t<f64> ret
= pos_to_val(flattened_in_min_cell_arr, flattened_in_max_cell_arr);

return ret.template cast<std::vector<f64>>();
},
offset);
},
py::arg("field_name"),
py::arg("pos_to_val"),
py::arg("offset") = 0)
.def(
"set_field_value_lambda_f64_3_all",
[](T &self,
std::string field_name,
const std::function<pybind11::array_t<f64>(
const pybind11::array_t<f64> &, const pybind11::array_t<f64> &)> pos_to_val,
const i32 offset) {
return self.template set_field_value_lambda_all<f64_3>(
field_name,
[&](const std::vector<f64_3> &in_min_cell,
const std::vector<f64_3> &in_max_cell) {
auto flattened_in_min_cell_arr
= VecToNumpy<f64_3>::convert(in_min_cell);
auto flattened_in_max_cell_arr
= VecToNumpy<f64_3>::convert(in_max_cell);

pybind11::array_t<f64> ret
= pos_to_val(flattened_in_min_cell_arr, flattened_in_max_cell_arr);

std::vector<f64> ret_vec = ret.reshape({in_min_cell.size() * 3})
.template cast<std::vector<f64>>();

std::vector<f64_3> ret2{};
ret2.reserve(ret_vec.size() / 3);
for (u32 i = 0; i < ret_vec.size(); i += 3) {
ret2.push_back({ret_vec[i], ret_vec[i + 1], ret_vec[i + 2]});
}

return ret2;
},
offset);
},
py::arg("field_name"),
py::arg("pos_to_val"),
py::arg("offset") = 0)
.def(
"gen_default_config",
[](T &self) -> TConfig {
Expand Down