sedimentation removal dynamic (for 0-d environments)

PySDM/backends/impl_numba/methods/__init__.py

@@ -14,3 +14,4 @@
 from .terminal_velocity_methods import TerminalVelocityMethods
 from .seeding_methods import SeedingMethods
 from .deposition_methods import DepositionMethods
+from .sedimentation_removal_0d_methods import SedimentationRemoval0DMethods

PySDM/backends/impl_numba/methods/sedimentation_removal_0d_methods.py

@@ -0,0 +1,58 @@
+"""
+CPU implementation of backend methods for removal due to sedimentation
+in a 0D environment
+"""
+
+from functools import cached_property
+
+import numba
+
+from PySDM.backends.impl_common.backend_methods import BackendMethods
+
+
+class SedimentationRemoval0DMethods(BackendMethods):
+
+    @cached_property
+    def _sedimentation_removal_deterministic_body(self):
+
+        @numba.njit(**self.default_jit_flags)
+        def body(relative_fall_velocity, multiplicity, length_scale, timestep):
+            n_sd = len(relative_fall_velocity)
+            for i in numba.prange(n_sd):  # pylint: disable=not-an-iterable
+                multiplicity[i] -= (
+                    multiplicity[i]
+                    * relative_fall_velocity[i]
+                    * timestep
+                    / length_scale
+                )
+
+        return body
+
+    @cached_property
+    def _sedimentation_removal_stochastic_body(self):
+
+        prob_zero_events = self.formulae.trivia.poissonian_avoidance_function
+
+        @numba.njit(**self.default_jit_flags)
+        def body(relative_fall_velocity, multiplicity, length_scale, timestep):
+            n_sd = len(relative_fall_velocity)
+            for i in numba.prange(n_sd):  # pylint: disable=not-an-iterable
+                removal_rate = relative_fall_velocity[i] / length_scale
+                survive_prob = prob_zero_events(r=removal_rate, dt=timestep)
+                multiplicity[i] *= survive_prob
+
+        return body
+
+    def sedimentation_removal_deterministic(
+        self, *, relative_fall_velocity, multiplicity, length_scale, timestep
+    ):
+        self._sedimentation_removal_deterministic_body(
+            relative_fall_velocity, multiplicity, length_scale, timestep
+        )
+
+    def sedimentation_removal_stochastic(
+        self, *, relative_fall_velocity, multiplicity, length_scale, timestep
+    ):
+        self._sedimentation_removal_stochastic_body(
+            relative_fall_velocity, multiplicity, length_scale, timestep
+        )

PySDM/backends/numba.py

@@ -32,6 +32,7 @@ class Numba(  # pylint: disable=too-many-ancestors,duplicate-code
     methods.IsotopeMethods,
     methods.SeedingMethods,
     methods.DepositionMethods,
+    methods.SedimentationRemoval0DMethods,
 ):
     Storage = ImportedStorage
     Random = ImportedRandom

PySDM/dynamics/__init__.py

@@ -17,3 +17,4 @@
 from PySDM.dynamics.relaxed_velocity import RelaxedVelocity
 from PySDM.dynamics.seeding import Seeding
 from PySDM.dynamics.vapour_deposition_on_ice import VapourDepositionOnIce
+from PySDM.dynamics.sedimentation_removal_0d import SedimentationRemoval0D

PySDM/dynamics/sedimentation_removal_0d.py

@@ -0,0 +1,27 @@
+"""deposition removal logic for zero-dimensional environments"""
+
+from typing import Optional
+import numpy as np
+from PySDM.dynamics.impl import register_dynamic
+
+
+@register_dynamic()
+class SedimentationRemoval0D:
+    def __init__(self, *, stochastic_sedimentation_removal: Optional[bool] = True):
+        """stochastic or deterministic removal"""
+        self.stochastic_sedimentation_removal = stochastic_sedimentation_removal
+        self.particulator = None
+
+    def register(self, builder):
+        builder.request_attribute("relative fall velocity")
+        assert builder.particulator.environment.mesh.n_dims == 0
+        self.particulator = builder.particulator
+
+    def __call__(self):
+        """for stochastic removal, see, e.g., the naive scheme in Algorithm 1 in
+        [Curtis et al. 2016](https://doi.org/10.1016/j.jcp.2016.06.029)"""
+
+        self.particulator.sedimentation_removal(
+            stochastic_sedimentation_removal=self.stochastic_sedimentation_removal,
+            length_scale=np.cbrt(self.particulator.environment.mesh.dv),
+        )

PySDM/particulator.py

@@ -595,3 +595,19 @@ def thaw_instantaneous(self):
             temperature=self.environment["T"],
         )
         self.attributes.mark_updated("signed water mass")
+
+    def sedimentation_removal(self, *, stochastic_sedimentation_removal, length_scale):
+        if stochastic_sedimentation_removal:
+            self.backend.sedimentation_removal_stochastic(
+                relative_fall_velocity=self.attributes["relative fall velocity"].data,
+                multiplicity=self.attributes["multiplicity"].data,
+                length_scale=length_scale,
+                timestep=self.dt,
+            )
+        else:
+            self.backend.sedimentation_removal_deterministic(
+                relative_fall_velocity=self.attributes["relative fall velocity"].data,
+                multiplicity=self.attributes["multiplicity"].data,
+                length_scale=length_scale,
+                timestep=self.dt,
+            )

docs/bibliography.json

@@ -1047,5 +1047,12 @@
     ],
     "label": "Arabas & Pawlowska 2011",
     "title": "Adaptive method of lines for multi-component aerosol condensational growth and CCN activation"
+  },
+  "https://doi.org/10.1016/j.jcp.2016.06.029": {
+    "usages": [
+      "PySDM/dynamics/sedimentation_removal_0d.py"
+    ],
+    "label": "Curtis et al. 2016 (J. Comp. Phys. 322)",
+    "title": "Accelerated simulation of stochastic particle removal processes in particle-resolved aerosol models"
   }
 }

tests/unit_tests/dynamics/test_sedimentation_removal_0d.py

@@ -0,0 +1,106 @@
+# pylint: disable=missing-module-docstring
+from matplotlib import pyplot
+import pytest
+import numpy as np
+from PySDM.dynamics import SedimentationRemoval0D
+from PySDM.physics import si
+from PySDM.environments import Box
+from PySDM import Builder
+from PySDM.backends import ThrustRTC
+from PySDM.products import ParticleConcentration, SuperDropletCountPerGridbox, Time
+
+
+class TestSedimentationRemoval0D:  # pylint: disable=too-few-public-methods,too-many-branches,too-many-locals
+    @staticmethod
+    @pytest.mark.parametrize("stochastic", (True, False))
+    def test_convergence_wrt_dt(backend_class, stochastic, plot=False):
+        # arrange
+        dts = 0.5 * si.s, 5 * si.s
+        dvs = 1e2 * si.m**3, 1e3 * si.m**3, 1e4 * si.m**3
+        t_max = 600 * si.s
+        multiplicities = 1e5, 1e6, 1e7, 1e8
+        water_masses = 1 * si.ug, 2 * si.ug, 3 * si.ug, 4 * si.ug
+
+        if backend_class is ThrustRTC:
+            pytest.skip("TODO #1871")
+
+        # act
+        output = {}
+        for dt in dts:
+            for dv in dvs:
+                builder = Builder(
+                    n_sd=len(multiplicities),
+                    environment=Box(dv=dv, dt=dt),
+                    backend=backend_class(),
+                    dynamics=(
+                        SedimentationRemoval0D(
+                            stochastic_sedimentation_removal=stochastic
+                        ),
+                    ),
+                )
+                particulator = builder.build(
+                    attributes={
+                        "multiplicity": np.asarray(multiplicities),
+                        "signed water mass": np.asarray(water_masses),
+                    },
+                    products=(
+                        ParticleConcentration(),
+                        SuperDropletCountPerGridbox(),
+                        Time(),
+                    ),
+                )
+                key = f"{dt=} {dv=}"
+                output[key] = {name: [] for name in particulator.products}
+                while particulator.n_steps * dt <= t_max:
+                    if len(output[key]["time"]) != 0:
+                        particulator.run(steps=1)
+                    for name, product in particulator.products.items():
+                        output[key][name].append(product.get() + 0)
+
+        # plot
+        pyplot.title(f"{stochastic=}")
+        pyplot.xlabel("time [s]")
+        pyplot.ylabel("particle concentration [m$^{-3}$]")
+        for dt in dts:
+            for dv in dvs:
+                key = f"{dt=} {dv=}"
+                pyplot.plot(
+                    output[key]["time"],
+                    output[key]["particle concentration"],
+                    label=key,
+                    linewidth=4 + 3 * np.log10(dt),
+                )
+        pyplot.gca().set_yscale("log")
+        pyplot.gca().set_xlim(left=0, right=t_max)
+        pyplot.legend()
+        pyplot.grid()
+
+        if plot:
+            pyplot.show()
+        else:
+            pyplot.clf()
+
+        # assert
+        for dt in dts:
+            for dv in dvs:
+                key = f"{dt=} {dv=}"
+                particle_concentration = np.asarray(
+                    output[key]["particle concentration"]
+                )
+                assert particle_concentration[-1] == 0.0
+
+        for dv in dvs:
+            time_compare = 0
+            for dt in dts:
+                key = f"{dt=} {dv=}"
+                time_end = np.asarray(output[key]["time"])[-1]
+                assert time_end >= time_compare
+                time_compare = time_end
+
+        for dt in dts:
+            time_compare = 0
+            for dv in dvs:
+                key = f"{dt=} {dv=}"
+                time_end = np.asarray(output[key]["time"])[-1]
+                assert time_end >= time_compare
+                time_compare = time_end