add STP option to ParticleSizeSpectrum* products (#1346)

PySDM/products/impl/concentration_product.py

@@ -7,19 +7,23 @@
 
 
 class ConcentrationProduct(MomentProduct):
+    @staticmethod
+    def check_ctor_arguments(specific, stp):
+        if stp and specific:
+            raise ValueError(
+                "std-temperature-and-pressure precludes specific conc. option"
+            )
+
     def __init__(self, *, unit: str, name: str, specific: bool, stp: bool):
         """
         `stp` toggles expressing the concentration in terms of standard temperature
         and pressure conditions (ground level of the ICAO standard atmosphere, zero humidity)
         """
+        self.check_ctor_arguments(specific, stp)
         super().__init__(unit=unit, name=name)
         self.specific = specific
         self.stp = stp
         self.rho_stp = None
-        if self.stp and self.specific:
-            raise ValueError(
-                "std-temperature-and-pressure precludes specific conc. option"
-            )
 
     def register(self, builder):
         super().register(builder)

PySDM/products/size_spectral/particle_size_spectrum.py

@@ -7,13 +7,26 @@
 import numpy as np
 
 from PySDM.products.impl.spectrum_moment_product import SpectrumMomentProduct
+from PySDM.products.impl.concentration_product import ConcentrationProduct
 
 
 class ParticleSizeSpectrum(SpectrumMomentProduct, ABC):
-    def __init__(self, *, radius_bins_edges, name, unit, dry=False, specific=False):
+    def __init__(
+        self,
+        *,
+        stp: bool,
+        radius_bins_edges,
+        name: str,
+        unit: str,
+        dry: bool = False,
+        specific: bool = False,
+    ):
+        ConcentrationProduct.check_ctor_arguments(specific, stp)
         self.volume_attr = "dry volume" if dry else "volume"
         self.radius_bins_edges = radius_bins_edges
         self.specific = specific
+        self.stp = stp
+        self.rho_stp = None
         super().__init__(name=name, unit=unit, attr_unit="m^3")
 
     def register(self, builder):
@@ -29,6 +42,7 @@ def register(self, builder):
         super().register(builder)
 
         self.shape = (*builder.particulator.mesh.grid, len(self.attr_bins_edges) - 1)
+        self.rho_stp = builder.formulae.constants.rho_STP
 
     def _impl(self, **kwargs):
         vals = np.empty([self.particulator.mesh.n_cell, len(self.attr_bins_edges) - 1])
@@ -42,37 +56,50 @@ def _impl(self, **kwargs):
 
         vals[:] /= self.particulator.mesh.dv
 
-        if self.specific:
+        if self.specific or self.stp:
             self._download_to_buffer(self.particulator.environment["rhod"])
 
         for i in range(len(self.attr_bins_edges) - 1):
             dr = self.formulae.trivia.radius(
                 volume=self.attr_bins_edges[i + 1]
             ) - self.formulae.trivia.radius(volume=self.attr_bins_edges[i])
             vals[:, i] /= dr
-            if self.specific:
+            if self.specific or self.stp:
                 vals[:, i] /= self.buffer.ravel()
+                if self.stp:
+                    vals[:, i] *= self.rho_stp
 
         return np.squeeze(vals.reshape(self.shape))
 
 
 class ParticleSizeSpectrumPerMassOfDryAir(ParticleSizeSpectrum):
-    def __init__(self, radius_bins_edges, dry=False, name=None, unit="kg^-1 m^-1"):
+    def __init__(
+        self,
+        *,
+        radius_bins_edges,
+        dry=False,
+        name=None,
+        unit="kg^-1 m^-1",
+    ):
         super().__init__(
             radius_bins_edges=radius_bins_edges,
             dry=dry,
             specific=True,
             name=name,
             unit=unit,
+            stp=False,
         )
 
 
 class ParticleSizeSpectrumPerVolume(ParticleSizeSpectrum):
-    def __init__(self, radius_bins_edges, dry=False, name=None, unit="m^-3 m^-1"):
+    def __init__(
+        self, *, radius_bins_edges, dry=False, name=None, unit="m^-3 m^-1", stp=False
+    ):
         super().__init__(
             radius_bins_edges=radius_bins_edges,
             dry=dry,
             specific=False,
             name=name,
             unit=unit,
+            stp=stp,
         )

tests/unit_tests/products/test_particle_size_spectrum.py

@@ -11,18 +11,25 @@
     ParticleSizeSpectrumPerMassOfDryAir,
     ParticleSizeSpectrumPerVolume,
 )
+from PySDM.products.size_spectral.particle_size_spectrum import ParticleSizeSpectrum
 
 
 class TestParticleSizeSpectrum:
     @staticmethod
     @pytest.mark.parametrize(
-        "product_class, specific, unit",
+        "product_class, specific, unit, stp",
         (
-            (ParticleSizeSpectrumPerVolume, False, "1.0 / meter ** 4"),
-            (ParticleSizeSpectrumPerMassOfDryAir, True, "1.0 / kilogram / meter"),
+            (ParticleSizeSpectrumPerVolume, False, "1.0 / meter ** 4", True),
+            (ParticleSizeSpectrumPerVolume, False, "1.0 / meter ** 4", False),
+            (
+                ParticleSizeSpectrumPerMassOfDryAir,
+                True,
+                "1.0 / kilogram / meter",
+                False,
+            ),
         ),
     )
-    def test_specific_flag(product_class, specific, unit, backend_instance):
+    def test_specific_flag(product_class, specific, unit, backend_instance, stp):
         """checks if the reported concentration is correctly normalised per
         volume or mass of air, and per bin size"""
         # arrange
@@ -35,7 +42,9 @@ def test_specific_flag(product_class, specific, unit, backend_instance):
         n_sd = 1
         box = Box(dt=np.nan, dv=dv)
         builder = Builder(n_sd=n_sd, backend=backend_instance, environment=box)
-        sut = product_class(radius_bins_edges=(min_size, max_size))
+        sut = product_class(
+            radius_bins_edges=(min_size, max_size), **({"stp": stp} if stp else {})
+        )
         builder.build(
             products=(sut,),
             attributes={
@@ -56,7 +65,8 @@ def test_specific_flag(product_class, specific, unit, backend_instance):
             desired=multiplicity
             / dv
             / (max_size - min_size)
-            / (rhod if specific else 1),
+            / (rhod if specific or stp else 1)
+            * (backend_instance.formulae.constants.rho_STP if stp else 1),
             significant=10,
         )
 
@@ -103,3 +113,10 @@ def test_dry_flag(product_class, dry, backend_instance):
             / (rhod if sut.specific else 1),
             significant=10,
         )
+
+    @staticmethod
+    def test_stp_flag_incompatible_with_specific():
+        with pytest.raises(Exception, match="precludes specific"):
+            ParticleSizeSpectrum(
+                stp=True, specific=True, name="", unit="", radius_bins_edges=()
+            )