[Examples] dustydisc setup

examples/sph/run_dustysettle_tvi.py

@@ -28,32 +28,8 @@
     shamrock.sys.init("0:0")
 
 # %%
-# mpl style
-mpl.rcParams.update(
-    {
-        "font.family": "serif",
-        "mathtext.fontset": "cm",
-        "font.size": 14,
-        "axes.labelsize": 16,
-        "axes.titlesize": 16,
-        "xtick.labelsize": 13,
-        "ytick.labelsize": 13,
-        "legend.fontsize": 13,
-        "axes.facecolor": "#f2f2f2",
-        "axes.linewidth": 1.0,
-        "xtick.direction": "in",
-        "ytick.direction": "in",
-        "xtick.top": True,
-        "ytick.right": True,
-        "xtick.major.size": 8,
-        "ytick.major.size": 8,
-        "xtick.minor.visible": True,
-        "ytick.minor.visible": True,
-        "legend.frameon": True,
-        "legend.fancybox": False,
-        "legend.edgecolor": "black",
-    }
-)
+# Use shamrock documentation style for matplotlib
+shamrock.matplotlib.set_shamrock_mpl_style()
 
 # %%
 # Sim parameters

examples/sph/run_dustywave_tvi.py

@@ -37,32 +37,8 @@
 N_target = 1e3
 
 # %%
-# mpl style
-mpl.rcParams.update(
-    {
-        "font.family": "serif",
-        "mathtext.fontset": "cm",
-        "font.size": 14,
-        "axes.labelsize": 16,
-        "axes.titlesize": 16,
-        "xtick.labelsize": 13,
-        "ytick.labelsize": 13,
-        "legend.fontsize": 13,
-        "axes.facecolor": "#f2f2f2",
-        "axes.linewidth": 1.0,
-        "xtick.direction": "in",
-        "ytick.direction": "in",
-        "xtick.top": True,
-        "ytick.right": True,
-        "xtick.major.size": 8,
-        "ytick.major.size": 8,
-        "xtick.minor.visible": True,
-        "ytick.minor.visible": True,
-        "legend.frameon": True,
-        "legend.fancybox": False,
-        "legend.edgecolor": "black",
-    }
-)
+# Use shamrock documentation style for matplotlib
+shamrock.matplotlib.set_shamrock_mpl_style()
 
 # %%
 # Do setup

src/pylib/shamrock/__init__.py

@@ -24,7 +24,7 @@
         # Some C-extension objects or builtins don't allow rebinding __module__
         pass
 
-from . import utils
+from . import matplotlib, utils
 
 # print(f"shamrock.__all__: {__all__}")
 # print(f"shamrock imported from {__file__}")

src/pylib/shamrock/matplotlib/__init__.py

@@ -0,0 +1 @@
+from .style import set_shamrock_mpl_style

src/pylib/shamrock/matplotlib/style.py

@@ -0,0 +1,37 @@
+"""
+Set the matplotlib style for shamrock (doc and standard plots).
+"""
+
+import matplotlib as mpl
+
+
+def set_shamrock_mpl_style():
+    """
+    Set the matplotlib style for shamrock (doc and standard plots).
+    """
+
+    mpl.rcParams.update(
+        {
+            "font.family": "serif",
+            "mathtext.fontset": "cm",
+            "font.size": 14,
+            "axes.labelsize": 16,
+            "axes.titlesize": 16,
+            "xtick.labelsize": 13,
+            "ytick.labelsize": 13,
+            "legend.fontsize": 13,
+            "axes.facecolor": "#f2f2f2",
+            "axes.linewidth": 1.0,
+            "xtick.direction": "in",
+            "ytick.direction": "in",
+            "xtick.top": True,
+            "ytick.right": True,
+            "xtick.major.size": 8,
+            "ytick.major.size": 8,
+            "xtick.minor.visible": True,
+            "ytick.minor.visible": True,
+            "legend.frameon": True,
+            "legend.fancybox": False,
+            "legend.edgecolor": "black",
+        }
+    )

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

@@ -1,6 +1,68 @@
+import numpy as np
+
 from .StandardPlotHelper import StandardPlotHelper
 
 
+def get_rhod_j_getter(model, jdust, ndust):
+
+    def int_getter(size: int, dic_out: dict) -> np.array:
+        s_j = dic_out["s_j"].reshape(-1, ndust)
+        return s_j[:, jdust] ** 2
+
+    return int_getter
+
+
+def get_rhod_getter(model, ndust):
+
+    def int_getter(size: int, dic_out: dict) -> np.array:
+        s_j = dic_out["s_j"].reshape(-1, ndust)
+        rho_d_sum = np.sum(s_j**2, axis=1)
+        return rho_d_sum
+
+    return int_getter
+
+
+def get_rhog_getter(model, ndust):
+
+    pmass = model.get_particle_mass()
+    hfact = model.get_hfact()
+
+    rhod_getter = get_rhod_getter(model, ndust)
+
+    def int_getter(size: int, dic_out: dict) -> np.array:
+        rho_d_sum = rhod_getter(size, dic_out)
+
+        hpart = dic_out["hpart"]
+        rho = pmass * (hfact / np.array(hpart)) ** 3
+
+        return rho - rho_d_sum
+
+    return int_getter
+
+
+def get_epsilon_j_getter(model, jdust, ndust):
+
+    pmass = model.get_particle_mass()
+    hfact = model.get_hfact()
+
+    rhod_j_getter = get_rhod_j_getter(model, jdust, ndust)
+
+    def int_getter(size: int, dic_out: dict) -> np.array:
+
+        rhod = rhod_j_getter(size, dic_out)
+
+        hpart = dic_out["hpart"]
+        rho = pmass * (hfact / np.array(hpart)) ** 3
+
+        tmp = rhod / rho
+
+        print(tmp.min(), tmp.max())
+        print(tmp)
+        return tmp
+
+    return int_getter
+
+
 def ColumnDensityPlot(model, ext_r, nx, ny, ex, ey, center, analysis_folder, analysis_prefix):
     def compute_rho_integ(helper):
         return helper.column_integ_render("rho", "f64")
@@ -19,6 +81,28 @@ def compute_rho_integ(helper):
     )
 
 
+def ColumnDensityPlotDust(
+    model, ext_r, nx, ny, ex, ey, center, analysis_folder, analysis_prefix, jdust, ndust
+):
+    def compute_rhod_integ(helper):
+        return helper.column_integ_render(
+            "custom", "f64", custom_getter=get_rhod_j_getter(model, jdust, ndust)
+        )
+
+    return StandardPlotHelper(
+        model,
+        ext_r,
+        nx,
+        ny,
+        ex,
+        ey,
+        center,
+        analysis_folder,
+        analysis_prefix,
+        compute_function=compute_rhod_integ,
+    )
+
+
 def SliceDensityPlot(
     model,
     ext_r,
@@ -47,3 +131,41 @@ def compute_rho_slice(helper):
         analysis_prefix,
         compute_function=compute_rho_slice,
     )
+
+
+def SliceDensityPlotDust(
+    model,
+    ext_r,
+    nx,
+    ny,
+    ex,
+    ey,
+    center,
+    analysis_folder,
+    analysis_prefix,
+    jdust,
+    ndust,
+    do_normalization=True,
+    min_normalization=1e-9,
+):
+    def compute_rho_slice(helper):
+        return helper.slice_render(
+            "custom",
+            "f64",
+            do_normalization,
+            min_normalization,
+            custom_getter=get_rhod_j_getter(model, jdust, ndust),
+        )
+
+    return StandardPlotHelper(
+        model,
+        ext_r,
+        nx,
+        ny,
+        ex,
+        ey,
+        center,
+        analysis_folder,
+        analysis_prefix,
+        compute_function=compute_rho_slice,
+    )

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

@@ -348,6 +348,7 @@ def make_plot(
         sink_linewidth=1,
         sink_fill=False,
         save_plot=True,
+        extra_title=None,
         **kwargs,
     ):
         if shamrock.sys.world_rank() == 0:
@@ -403,6 +404,8 @@ def make_plot(
             plt.ylabel(f"{y_label} {dist_label_y}")
 
             text = f"t = {metadata['time']:0.3f} {time_label}"
+            if extra_title is not None:
+                text += f" {extra_title}"
             self.figure_add_time_info(text, holywood_mode)
 
             cmap_label = f"{field_label} {field_unit_label}"

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

@@ -4,7 +4,16 @@
 from .UnitHelper import plot_codeu_to_unit  # noqa: I001
 
 # Render based analysis
-from .DensityPlots import ColumnDensityPlot, SliceDensityPlot
+from .DensityPlots import (
+    ColumnDensityPlot,
+    SliceDensityPlot,
+    ColumnDensityPlotDust,
+    SliceDensityPlotDust,
+    get_epsilon_j_getter,
+    get_rhod_j_getter,
+    get_rhod_getter,
+    get_rhog_getter,
+)
 from .ColumnParticleCount import ColumnParticleCount
 from .ParticlesDt import SliceDtPart
 from .VelocityPlots import (