[Staging PR] Tweak sinks behavior in SPH (better conservation)

src/pylib/shamrock/utils/analysis/AngularMomentumPlots.py

@@ -0,0 +1,176 @@
+import numpy as np
+
+import shamrock.sys
+
+from .StandardPlotHelper import StandardPlotHelper
+
+try:
+    from numba import njit
+
+    _HAS_NUMBA = True
+except ImportError:
+    _HAS_NUMBA = False
+
+
+def SliceAngularMomentum(
+    model,
+    ext_r,
+    nx,
+    ny,
+    ex,
+    ey,
+    center,
+    analysis_folder,
+    analysis_prefix,
+    do_normalization=True,
+    min_normalization=1e-9,
+    Lprojection=[0.0, 0.0, 1.0],
+):
+    def compute_angular_mom(helper):
+        if _HAS_NUMBA:
+            if shamrock.sys.world_rank() == 0:
+                print("Using numba for angular momentum in SliceAngularMomentum")
+
+        pmass = model.get_particle_mass()
+        hfact = model.get_hfact()
+
+        def internal(
+            size: int,
+            hpart: np.array,
+            x: np.array,
+            y: np.array,
+            z: np.array,
+            vx: np.array,
+            vy: np.array,
+            vz: np.array,
+        ) -> np.array:
+
+            rho = pmass * (hfact / hpart) ** 3
+
+            r = np.stack([x, y, z], axis=-1)
+            v = np.stack([vx, vy, vz], axis=-1)
+            L = rho[:, None] * np.cross(r, v)
+
+            # project L onto the Lprojection vector
+            L = np.sum(L * Lprojection, axis=-1)
+
+            return L
+
+        if _HAS_NUMBA:
+            internal = njit(internal)
+
+        def custom_getter(size: int, dic_out: dict) -> np.array:
+            return internal(
+                size,
+                dic_out["hpart"],
+                dic_out["xyz"][:, 0],
+                dic_out["xyz"][:, 1],
+                dic_out["xyz"][:, 2],
+                dic_out["vxyz"][:, 0],
+                dic_out["vxyz"][:, 1],
+                dic_out["vxyz"][:, 2],
+            )
+
+        arr_L = helper.slice_render(
+            "custom",
+            "f64",
+            do_normalization,
+            min_normalization,
+            custom_getter=custom_getter,
+        )
+
+        return arr_L
+
+    return StandardPlotHelper(
+        model,
+        ext_r,
+        nx,
+        ny,
+        ex,
+        ey,
+        center,
+        analysis_folder,
+        analysis_prefix,
+        compute_function=compute_angular_mom,
+    )
+
+
+def ColumnAverageAngularMomentum(
+    model,
+    ext_r,
+    nx,
+    ny,
+    ex,
+    ey,
+    center,
+    analysis_folder,
+    analysis_prefix,
+    min_normalization=1e-9,
+    Lprojection=[0.0, 0.0, 1.0],
+):
+    def compute_angular_mom(helper):
+        if _HAS_NUMBA:
+            if shamrock.sys.world_rank() == 0:
+                print("Using numba for angular momentum in SliceAngularMomentum")
+
+        pmass = model.get_particle_mass()
+        hfact = model.get_hfact()
+
+        def internal(
+            size: int,
+            hpart: np.array,
+            x: np.array,
+            y: np.array,
+            z: np.array,
+            vx: np.array,
+            vy: np.array,
+            vz: np.array,
+        ) -> np.array:
+
+            rho = pmass * (hfact / hpart) ** 3
+
+            r = np.stack([x, y, z], axis=-1)
+            v = np.stack([vx, vy, vz], axis=-1)
+            L = rho[:, None] * np.cross(r, v)
+
+            # project L onto the Lprojection vector
+            L = np.sum(L * Lprojection, axis=-1)
+
+            return L
+
+        if _HAS_NUMBA:
+            internal = njit(internal)
+
+        def custom_getter(size: int, dic_out: dict) -> np.array:
+            return internal(
+                size,
+                dic_out["hpart"],
+                dic_out["xyz"][:, 0],
+                dic_out["xyz"][:, 1],
+                dic_out["xyz"][:, 2],
+                dic_out["vxyz"][:, 0],
+                dic_out["vxyz"][:, 1],
+                dic_out["vxyz"][:, 2],
+            )
+
+        arr_L = helper.column_average_render(
+            "custom",
+            "f64",
+            min_normalization,
+            custom_getter=custom_getter,
+        )
+
+        return arr_L
+
+    return StandardPlotHelper(
+        model,
+        ext_r,
+        nx,
+        ny,
+        ex,
+        ey,
+        center,
+        analysis_folder,
+        analysis_prefix,
+        compute_function=compute_angular_mom,
+    )

src/pylib/shamrock/utils/analysis/__init__.py

@@ -22,5 +22,10 @@
     SliceBVerticalShearGradient,
 )
 
+from .AngularMomentumPlots import (
+    SliceAngularMomentum,
+    ColumnAverageAngularMomentum,
+)
+
 # Performance analysis
 from .PerfHistory import PerfHistory

src/shammath/include/shammath/matrix_op.hpp

@@ -334,6 +334,35 @@ namespace shammath {
         }
     }
 
+    /**
+     * @brief Compute the determinant of a 3x3 matrix.
+     *
+     * @param input The input matrix to invert.
+     * @return the determinant
+     */
+    template<class T, class SizeType, class Layout, class Accessor>
+    inline T mat_det_33(
+        const std::mdspan<T, std::extents<SizeType, 3, 3>, Layout, Accessor> &input) {
+
+        T &a00 = input(0, 0);
+        T &a10 = input(1, 0);
+        T &a20 = input(2, 0);
+
+        T &a01 = input(0, 1);
+        T &a11 = input(1, 1);
+        T &a21 = input(2, 1);
+
+        T &a02 = input(0, 2);
+        T &a12 = input(1, 2);
+        T &a22 = input(2, 2);
+
+        T det
+            = (-a02 * a11 * a20 + a01 * a12 * a20 + a02 * a10 * a21 - a00 * a12 * a21
+               - a01 * a10 * a22 + a00 * a11 * a22);
+
+        return det;
+    }
+
     /**
      * @brief Compute the inverse of a 3x3 matrix.
      *

src/shammodels/sph/include/shammodels/sph/Model.hpp

@@ -173,15 +173,17 @@ namespace shammodels::sph {
             Tscal q,
             std::mt19937 eng);
 
-        inline void add_sink(Tscal mass, Tvec pos, Tvec velocity, Tscal accretion_radius) {
+        inline void add_sink(
+            Tscal mass, Tvec pos, Tvec velocity, Tscal accretion_radius, bool is_torque_free) {
             if (solver.storage.sinks.is_empty()) {
                 solver.storage.sinks.set({});
             }
 
-            shamlog_debug_ln("SPH", "add sink :", mass, pos, velocity, accretion_radius);
+            shamlog_debug_ln(
+                "SPH", "add sink :", mass, pos, velocity, accretion_radius, is_torque_free);
 
             solver.storage.sinks.get().push_back(
-                {pos, velocity, {}, {}, mass, {}, accretion_radius});
+                {pos, velocity, {}, {}, mass, {}, accretion_radius, is_torque_free});
         }
 
         template<class T>

src/shammodels/sph/include/shammodels/sph/SinkPartStruct.hpp

@@ -34,6 +34,8 @@ namespace shammodels::sph {
         Tscal mass;
         Tvec angular_momentum;
         Tscal accretion_radius;
+
+        bool is_torque_free;
     };
 
     template<class Tvec>
@@ -50,6 +52,7 @@ namespace shammodels::sph {
             {"mass", p.mass},
             {"angular_momentum", p.angular_momentum},
             {"accretion_radius", p.accretion_radius},
+            {"is_torque_free", p.is_torque_free},
         };
     }
 
@@ -65,6 +68,7 @@ namespace shammodels::sph {
         j.at("mass").get_to(p.mass);
         j.at("angular_momentum").get_to(p.angular_momentum);
         j.at("accretion_radius").get_to(p.accretion_radius);
+        j.at("is_torque_free").get_to(p.is_torque_free);
     }
 
 } // namespace shammodels::sph

src/shammodels/sph/src/Model.cpp

@@ -1437,7 +1437,8 @@ void shammodels::sph::Model<Tvec, SPHKernel>::init_from_phantom_dump(
                     mass[i],
                     {xsink[i], ysink[i], zsink[i]},
                     {vxsink[i], vysink[i], vzsink[i]},
-                    Racc[i]);
+                    Racc[i],
+                    false);
             }
         }
     }