coalescence-only single-column test case ("deJong_Azimi" based on S&H…

…); power-law terminal velocity; initialisation of particles in part of the domain in 1D kinematic env (#1221)

Co-authored-by: Sylwester Arabas <sylwester.arabas@uj.edu.pl>
Co-authored-by: Sylwester Arabas <sylwester.arabas@agh.edu.pl>

PySDM/backends/impl_numba/methods/terminal_velocity_methods.py

@@ -46,3 +46,23 @@ def terminal_velocity(self, *, values, radius, k1, k2, k3, r1, r2):
         self.terminal_velocity_body(
             values=values, radius=radius, k1=k1, k2=k2, k3=k3, r1=r1, r2=r2
         )
+
+    @cached_property
+    def power_series_body(self):
+        @numba.njit(**{**conf.JIT_FLAGS, "fastmath": self.formulae.fastmath})
+        def power_series_body(*, values, radius, num_terms, prefactors, powers):
+            for i in numba.prange(len(values)):  # pylint: disable=not-an-iterable
+                values[i] = 0.0
+                for j in range(num_terms):
+                    values[i] = values[i] + prefactors[j] * radius[i] ** (powers[j] * 3)
+
+        return power_series_body
+
+    def power_series(self, *, values, radius, num_terms, prefactors, powers):
+        self.power_series_body(
+            values=values,
+            radius=radius,
+            num_terms=num_terms,
+            prefactors=prefactors,
+            powers=powers,
+        )

PySDM/backends/impl_thrust_rtc/methods/terminal_velocity_methods.py

@@ -131,3 +131,27 @@ def interpolation(self, *, output, radius, factor, b, c):
         self.__interpolation_body.launch_n(
             len(radius), (output.data, radius.data, factor_device, b.data, c.data)
         )
+
+    @cached_property
+    def __power_series_body(self):
+        # TODO #599 r<0
+        return trtc.For(
+            ("output", "radius", "num_terms", "prefactors", "powers"),
+            "i",
+            """
+            output[i] = 0.0
+            int kmax = num_terms
+            for (auto k = 0; k < kmax; k+=1) {
+                auto sumterm = prefactors[k] * pow(radius[i], 3*powers[k])
+                output[i] = output[i] + sumterm
+            """,
+        )
+
+    @nice_thrust(**NICE_THRUST_FLAGS)
+    def power_series(self, *, values, radius, num_terms, prefactors, powers):
+        prefactors = self._get_floating_point(prefactors)
+        powers = self._get_floating_point(powers)
+        num_terms = self._get_floating_point(num_terms)
+        self.__power_series_body.launch_n(
+            values.size(), (values, radius, num_terms, prefactors, powers)
+        )

PySDM/dynamics/terminal_velocity/__init__.py

@@ -3,4 +3,5 @@
 """
 
 from PySDM.dynamics.terminal_velocity.gunn_and_kinzer import GunnKinzer1949
+from PySDM.dynamics.terminal_velocity.power_series import PowerSeries
 from PySDM.dynamics.terminal_velocity.rogers_and_yau import RogersYau

PySDM/dynamics/terminal_velocity/power_series.py

@@ -0,0 +1,30 @@
+"""
+Power series expression
+"""
+
+import numpy as np
+
+from PySDM.physics import constants as const
+
+
+class PowerSeries:  # pylint: disable=too-few-public-methods
+    def __init__(self, particulator, *, prefactors=None, powers=None):
+        si = const.si
+        self.particulator = particulator
+        prefactors = prefactors or [2.0e-1 * si.m / si.s / np.sqrt(si.m)]
+        self.prefactors = np.array(prefactors)
+        powers = powers or [1 / 6]
+        self.powers = np.array(powers)
+        for i, p in enumerate(self.powers):
+            self.prefactors[i] *= (4 / 3 * const.PI) ** (p)
+            self.prefactors[i] /= (1 * si.um**3) ** (p)
+        assert len(self.prefactors) == len(self.powers)
+
+    def __call__(self, output, radius):
+        self.particulator.backend.power_series(
+            values=output.data,
+            radius=radius.data,
+            num_terms=len(self.powers),
+            prefactors=self.prefactors,
+            powers=self.powers,
+        )

PySDM/environments/kinematic_1d.py

@@ -32,7 +32,13 @@ def get_thd(self) -> np.ndarray:
         return self.thd0
 
     def init_attributes(
-        self, *, spatial_discretisation, spectral_discretisation, kappa
+        self,
+        *,
+        spatial_discretisation,
+        spectral_discretisation,
+        kappa,
+        z_part=None,
+        collisions_only=False
     ):
         super().sync()
         self.notify()
@@ -43,31 +49,40 @@ def init_attributes(
                 backend=self.particulator.backend,
                 grid=self.mesh.grid,
                 n_sd=self.particulator.n_sd,
+                z_part=z_part,
             )
             (
                 attributes["cell id"],
                 attributes["cell origin"],
                 attributes["position in cell"],
             ) = self.mesh.cellular_attributes(positions)
 
-            r_dry, n_per_kg = spectral_discretisation.sample(
-                backend=self.particulator.backend, n_sd=self.particulator.n_sd
-            )
-            attributes["dry volume"] = self.formulae.trivia.volume(radius=r_dry)
-            attributes["kappa times dry volume"] = attributes["dry volume"] * kappa
-            r_wet = equilibrate_wet_radii(
-                r_dry=r_dry,
-                environment=self,
-                cell_id=attributes["cell id"],
-                kappa_times_dry_volume=attributes["kappa times dry volume"],
-            )
+            if collisions_only:
+                v_wet, n_per_kg = spectral_discretisation.sample(
+                    backend=self.particulator.backend, n_sd=self.particulator.n_sd
+                )
+                # attributes["dry volume"] = v_wet
+                attributes["volume"] = v_wet
+                # attributes["kappa times dry volume"] = attributes["dry volume"] * kappa
+            else:
+                r_dry, n_per_kg = spectral_discretisation.sample(
+                    backend=self.particulator.backend, n_sd=self.particulator.n_sd
+                )
+                attributes["dry volume"] = self.formulae.trivia.volume(radius=r_dry)
+                attributes["kappa times dry volume"] = attributes["dry volume"] * kappa
+                r_wet = equilibrate_wet_radii(
+                    r_dry=r_dry,
+                    environment=self,
+                    cell_id=attributes["cell id"],
+                    kappa_times_dry_volume=attributes["kappa times dry volume"],
+                )
+                attributes["volume"] = self.formulae.trivia.volume(radius=r_wet)
 
             rhod = self["rhod"].to_ndarray()
             cell_id = attributes["cell id"]
             domain_volume = np.prod(np.array(self.mesh.size))
 
         attributes["multiplicity"] = n_per_kg * rhod[cell_id] * domain_volume
-        attributes["volume"] = self.formulae.trivia.volume(radius=r_wet)
 
         return attributes
 

PySDM/formulae.py

@@ -19,7 +19,7 @@
 
 from PySDM import physics
 from PySDM.backends.impl_numba import conf
-from PySDM.dynamics.terminal_velocity import GunnKinzer1949, RogersYau
+from PySDM.dynamics.terminal_velocity import GunnKinzer1949, PowerSeries, RogersYau
 from PySDM.dynamics.terminal_velocity.gunn_and_kinzer import TpDependent
 
 
@@ -124,6 +124,7 @@ def __init__(  # pylint: disable=too-many-locals
             "GunnKinzer1949": GunnKinzer1949,
             "RogersYau": RogersYau,
             "TpDependent": TpDependent,
+            "PowerSeries": PowerSeries,
         }[terminal_velocity]
 
     def __str__(self):

PySDM/initialisation/sampling/spatial_sampling.py

@@ -8,14 +8,23 @@
 
 class Pseudorandom:  # pylint: disable=too-few-public-methods
     @staticmethod
-    def sample(*, backend, grid, n_sd):
+    def sample(*, backend, grid, n_sd, z_part=None):
         dimension = len(grid)
         n_elements = dimension * n_sd
 
         storage = backend.Storage.empty(n_elements, dtype=float)
         backend.Random(seed=backend.formulae.seed, size=n_elements)(storage)
         positions = storage.to_ndarray().reshape(dimension, n_sd)
 
-        for dim in range(dimension):
-            positions[dim, :] *= grid[dim]
+        if z_part is None:
+            for dim in range(dimension):
+                positions[dim, :] *= grid[dim]
+        else:
+            assert dimension == 1
+            iz_min = int(grid[0] * z_part[0])
+            iz_max = int(grid[0] * z_part[1])
+            for dim in range(dimension):
+                positions[dim, :] *= iz_max - iz_min
+                positions[dim, :] += iz_min
+
         return positions

PySDM/products/size_spectral/__init__.py

@@ -4,6 +4,7 @@
     RadiusFirstMoment,
     RadiusSixthMoment,
     VolumeFirstMoment,
+    VolumeSecondMoment,
     ZerothMoment,
 )
 from .effective_radius import EffectiveRadius

PySDM/products/size_spectral/arbitrary_moment.py

@@ -32,6 +32,10 @@ def _impl(self, **kwargs):
     rank=1, attr="volume", attr_unit="m^3"
 )
 
+VolumeSecondMoment = make_arbitrary_moment_product(
+    rank=2, attr="volume", attr_unit="m^3"
+)
+
 RadiusSixthMoment = make_arbitrary_moment_product(rank=6, attr="radius", attr_unit="m")
 
 RadiusFirstMoment = make_arbitrary_moment_product(rank=1, attr="radius", attr_unit="m")

examples/PySDM_examples/Shipway_and_Hill_2012/settings.py

@@ -1,4 +1,4 @@
-from typing import Iterable
+from typing import Iterable, Optional
 
 import numpy as np
 from numdifftools import Derivative
@@ -9,6 +9,7 @@
 from PySDM.dynamics import condensation
 from PySDM.initialisation import spectra
 from PySDM.physics import si
+from PySDM.dynamics.collisions.collision_kernels import Geometric
 
 
 class Settings:
@@ -41,10 +42,13 @@ def __init__(
         dt: float = 1 * si.s,
         dz: float = 25 * si.m,
         z_max: float = 3000 * si.m,
+        z_part: Optional[tuple] = None,
         t_max: float = 60 * si.minutes,
         precip: bool = True,
+        enable_condensation: bool = True,
         formulae: Formulae = None,
-        save_spec_and_attr_times=()
+        save_spec_and_attr_times=(),
+        collision_kernel=None
     ):
         self.formulae = formulae or Formulae()
         self.n_sd_per_gridbox = n_sd_per_gridbox
@@ -55,9 +59,11 @@ def __init__(
         self.dt = dt
         self.dz = dz
         self.precip = precip
-
+        self.enable_condensation = enable_condensation
+        self.z_part = z_part
         self.z_max = z_max
         self.t_max = t_max
+        self.collision_kernel = collision_kernel or Geometric(collection_efficiency=1)
 
         t_1 = 600 * si.s
         self.rho_times_w = lambda t: (