normalisation by dv and by zero-th moment as option in make_arbitrary_moment_product (#1308)

Author: slayoo

PySDM/products/size_spectral/arbitrary_moment.py

@@ -6,36 +6,74 @@
 
 
 def make_arbitrary_moment_product(**kwargs):
+    """returns a product class to be instantiated and passed to a builder"""
     for arg in kwargs:
-        assert arg in ("rank", "attr", "attr_unit")
+        assert arg in (
+            "rank",
+            "attr",
+            "attr_unit",
+            "skip_division_by_m0",
+            "skip_division_by_dv",
+        )
 
     class ArbitraryMoment(MomentProduct):
         def __init__(
-            self, name=None, unit=f"({kwargs['attr_unit']})**{kwargs['rank']}"
+            self,
+            name=None,
+            unit=f"({kwargs['attr_unit']})**{kwargs['rank']}"
+            + ("" if kwargs["skip_division_by_dv"] else " / m**3"),
         ):
             super().__init__(name=name, unit=unit)
-            self.attr = kwargs["attr"]
-            self.rank = kwargs["rank"]
 
-        def _impl(self, **kwargs):
+        def _impl(self, **_):
             self._download_moment_to_buffer(
-                attr=self.attr, rank=self.rank, skip_division_by_m0=True
+                attr=kwargs["attr"],
+                rank=kwargs["rank"],
+                skip_division_by_m0=kwargs["skip_division_by_m0"],
             )
+            if not kwargs["skip_division_by_dv"]:
+                self.buffer /= self.particulator.mesh.dv
             return self.buffer
 
     return ArbitraryMoment
 
 
-ZerothMoment = make_arbitrary_moment_product(rank=0, attr="volume", attr_unit="m^3")
+ZerothMoment = make_arbitrary_moment_product(
+    rank=0,
+    attr="volume",
+    attr_unit="m^3",
+    skip_division_by_m0=True,
+    skip_division_by_dv=True,
+)
 
 VolumeFirstMoment = make_arbitrary_moment_product(
-    rank=1, attr="volume", attr_unit="m^3"
+    rank=1,
+    attr="volume",
+    attr_unit="m^3",
+    skip_division_by_m0=True,
+    skip_division_by_dv=True,
 )
 
 VolumeSecondMoment = make_arbitrary_moment_product(
-    rank=2, attr="volume", attr_unit="m^3"
+    rank=2,
+    attr="volume",
+    attr_unit="m^3",
+    skip_division_by_m0=True,
+    skip_division_by_dv=True,
 )
 
-RadiusSixthMoment = make_arbitrary_moment_product(rank=6, attr="radius", attr_unit="m")
+RadiusSixthMoment = make_arbitrary_moment_product(
+    rank=6,
+    attr="radius",
+    attr_unit="m",
+    skip_division_by_m0=True,
+    skip_division_by_dv=True,
+)
 
-RadiusFirstMoment = make_arbitrary_moment_product(rank=1, attr="radius", attr_unit="m")
+RadiusFirstMoment = make_arbitrary_moment_product(
+    rank=1,
+    attr="radius",
+    attr_unit="m",
+    skip_division_by_m0=True,
+    skip_division_by_dv=True,
+)

examples/PySDM_examples/Bieli_et_al_2022/simulation.py

@@ -26,8 +26,14 @@ def make_core(settings, coal_eff):
         adaptive=settings.adaptive,
     )
     builder.add_dynamic(collision)
-    M0 = am.make_arbitrary_moment_product(rank=0, attr="volume", attr_unit="m^3")
-    M1 = am.make_arbitrary_moment_product(rank=1, attr="volume", attr_unit="m^3")
-    M2 = am.make_arbitrary_moment_product(rank=2, attr="volume", attr_unit="m^3")
-    products = (M0(name="M0"), M1(name="M1"), M2(name="M2"))
+    common_args = {
+        "attr": "volume",
+        "attr_unit": "m^3",
+        "skip_division_by_m0": True,
+        "skip_division_by_dv": True,
+    }
+    products = tuple(
+        am.make_arbitrary_moment_product(rank=rank, **common_args)(name=f"M{rank}")
+        for rank in range(3)
+    )
     return builder.build(attributes, products)

tests/unit_tests/products/test_arbitrary_moment.py

@@ -0,0 +1,119 @@
+""" tests of the product-class-generating `make_arbitrary_moment_product` function """
+
+import pytest
+import numpy as np
+
+from PySDM.products.size_spectral.arbitrary_moment import make_arbitrary_moment_product
+from PySDM import Builder
+from PySDM.backends import CPU
+from PySDM.environments import Box
+from PySDM.physics import si
+
+
+class TestArbitraryMoment:
+    """groups tests to make [de]selecting easier"""
+
+    @staticmethod
+    @pytest.mark.parametrize(
+        "kwargs, expected_unit",
+        (
+            (
+                {
+                    "rank": 0,
+                    "attr": "radius",
+                    "attr_unit": "m",
+                    "skip_division_by_m0": True,
+                    "skip_division_by_dv": True,
+                },
+                "1 dimensionless",
+            ),
+            (
+                {
+                    "rank": 1,
+                    "attr": "radius",
+                    "attr_unit": "m",
+                    "skip_division_by_m0": False,
+                    "skip_division_by_dv": True,
+                },
+                "1 meter",
+            ),
+            (
+                {
+                    "rank": 6,
+                    "attr": "radius",
+                    "attr_unit": "m",
+                    "skip_division_by_m0": True,
+                    "skip_division_by_dv": True,
+                },
+                "1 meter ** 6",
+            ),
+            (
+                {
+                    "rank": 1,
+                    "attr": "water mass",
+                    "attr_unit": "kg",
+                    "skip_division_by_m0": False,
+                    "skip_division_by_dv": True,
+                },
+                "1 kilogram",
+            ),
+            (
+                {
+                    "rank": 1,
+                    "attr": "water mass",
+                    "attr_unit": "kg",
+                    "skip_division_by_m0": False,
+                    "skip_division_by_dv": True,
+                },
+                "1 kilogram",
+            ),
+            (
+                {
+                    "rank": 1,
+                    "attr": "water mass",
+                    "attr_unit": "kg",
+                    "skip_division_by_m0": False,
+                    "skip_division_by_dv": False,
+                },
+                "1.0 kilogram / meter ** 3",
+            ),
+        ),
+    )
+    def test_unit(kwargs, expected_unit):
+        """tests if the product unit respects arguments (incl. division by dv)"""
+        # arrange
+        product_class = make_arbitrary_moment_product(**kwargs)
+
+        # act
+        sut = product_class()
+
+        # assert
+        assert sut.unit == expected_unit
+
+    @staticmethod
+    @pytest.mark.parametrize("skip_division_by_dv", (True, False))
+    def test_division_by_dv(skip_division_by_dv):
+        """tests, using a single-superdroplet setup, if the volume normalisation logic works"""
+        # arrange
+        dv = 666 * si.m**3
+        product_class = make_arbitrary_moment_product(
+            rank=1,
+            attr="water mass",
+            attr_unit="kg",
+            skip_division_by_m0=False,
+            skip_division_by_dv=skip_division_by_dv,
+        )
+        builder = Builder(n_sd=1, backend=CPU(), environment=Box(dv=dv, dt=np.nan))
+        particulator = builder.build(
+            attributes={
+                k: np.ones(builder.particulator.n_sd)
+                for k in ("multiplicity", "water mass")
+            },
+            products=(product_class(name="sut"),),
+        )
+
+        # act
+        values = particulator.products["sut"].get()
+
+        # assert
+        assert values == 1 / (1 if skip_division_by_dv else dv)