Make the relaxed velocity dynamic use a non-constant timescale (#1114)

Co-authored-by: Sylwester Arabas <[email protected]>

PySDM/attributes/impl/mapper.py

@@ -20,6 +20,7 @@
     Multiplicities,
     Radius,
     RelativeFallVelocity,
+    SquareRootOfRadius,
     Temperature,
     TerminalVelocity,
     Volume,
@@ -61,6 +62,7 @@
     "kappa times dry volume": lambda _, __: KappaTimesDryVolume,
     "kappa": lambda _, __: Kappa,
     "radius": lambda _, __: Radius,
+    "square root of radius": lambda _, __: SquareRootOfRadius,
     "area": lambda _, __: Area,
     "dry radius": lambda _, __: DryRadius,
     "terminal velocity": lambda _, __: TerminalVelocity,

PySDM/attributes/physics/__init__.py

@@ -9,7 +9,7 @@
 from .equilibrium_supersaturation import EquilibriumSupersaturation
 from .heat import Heat
 from .multiplicities import Multiplicities
-from .radius import Radius
+from .radius import Radius, SquareRootOfRadius
 from .relative_fall_velocity import RelativeFallMomentum, RelativeFallVelocity
 from .temperature import Temperature
 from .terminal_velocity import TerminalVelocity

PySDM/attributes/physics/radius.py

@@ -14,3 +14,18 @@ def recalculate(self):
         self.data.idx = self.volume.data.idx
         self.data.product(self.volume.get(), 1 / self.formulae.constants.PI_4_3)
         self.data **= 1 / 3
+
+
+class SquareRootOfRadius(DerivedAttribute):
+    def __init__(self, builder):
+        self.radius = builder.get_attribute("radius")
+
+        super().__init__(
+            builder,
+            name="square root of radius",
+            dependencies=(self.radius,),
+        )
+
+    def recalculate(self):
+        self.data.fill(self.radius.data)
+        self.data **= 0.5

PySDM/backends/impl_numba/storage.py

@@ -193,3 +193,6 @@ def fill(self, other):
             self.data[:] = other.data
         else:
             self.data[:] = other
+
+    def exp(self):
+        self.data[:] = np.exp(self.data)

PySDM/backends/impl_thrust_rtc/storage.py

@@ -282,6 +282,19 @@ def divide_if_not_zero(output, divisor):
                 n=(output.shape[0]), args=(Impl.thrust((output, divisor)))
             )
 
+        __exp_body = trtc.For(
+            ("output",),
+            "i",
+            """
+                output[i] = exp(output[i]);
+            """,
+        )
+
+        @staticmethod
+        @nice_thrust(**NICE_THRUST_FLAGS)
+        def exp(output):
+            Impl.__exp_body.launch_n(output.shape[0], Impl.thrust((output,)))
+
     class Storage(StorageBase):
         FLOAT = BACKEND._get_np_dtype()
         INT = np.int64
@@ -510,4 +523,8 @@ def fill(self, other):
                 trtc.Fill(self.data, dvalue)
             return self
 
+        def exp(self):
+            Impl.exp(self)
+            return self
+
     return Storage

PySDM/dynamics/relaxed_velocity.py

@@ -3,27 +3,58 @@
 `PySDM.attributes.physics.relative_fall_velocity.RelativeFallVelocity`
 towards the terminal velocity
 """
-
-from math import exp
-
 from PySDM.attributes.impl.attribute import Attribute
 from PySDM.particulator import Particulator
 
 
 class RelaxedVelocity:  # pylint: disable=too-many-instance-attributes
     """
-    A dynamic which updates the fall momentum according to a relaxation timescale `tau`
+    A dynamic which updates the fall momentum according to a relaxation timescale
+    proportional to the sqrt of the droplet radius.
+
+    Should be added first in order to ensure the correct attributes are selected.
     """
 
-    def __init__(self, tau):
-        self.tau: float = tau
+    def __init__(self, c: float = 8, constant: bool = False):
+        """
+        Parameters:
+            - constant: use a constant relaxation timescale for all droplets
+            - c: relaxation timescale if `constant`, otherwise the proportionality constant
+        """
+        # the default value of c is a very rough estimate
+        self.c: float = c
+        self.constant = constant
 
         self.particulator = None
         self.fall_momentum_attr = None
         self.terminal_vel_attr = None
         self.volume_attr = None
+        self.sqrt_radius_attr = None
         self.rho_w = None  # TODO #798 - we plan to use masses instead of volumes soon
-        self.tmp_data = None
+
+        self.tmp_momentum_diff = None
+        self.tmp_tau = None
+        self.tmp_scale = None
+
+        self.tmp_tau_init = False
+
+    def calculate_tau(self, output, sqrt_radius_storage):
+        """
+        Calculates the relaxation timescale.
+        """
+        output.fill(self.c)
+        if not self.constant:
+            output *= sqrt_radius_storage
+
+    def calculate_scale_factor(self, output, tau_storage):
+        output.fill(-self.particulator.dt)
+        output /= tau_storage
+        output.exp()
+        output *= -1
+        output += 1
+
+    def create_storage(self, n):
+        return self.particulator.Storage.empty((n,), dtype=float)
 
     def register(self, builder):
         self.particulator: Particulator = builder.particulator
@@ -33,21 +64,32 @@ def register(self, builder):
         )
         self.terminal_vel_attr: Attribute = builder.get_attribute("terminal velocity")
         self.volume_attr: Attribute = builder.get_attribute("volume")
+        self.sqrt_radius_attr: Attribute = builder.get_attribute(
+            "square root of radius"
+        )
 
         self.rho_w: float = builder.formulae.constants.rho_w  # TODO #798
 
-        self.tmp_data = self.particulator.Storage.empty(
-            (self.particulator.n_sd,), dtype=float
-        )
+        self.tmp_momentum_diff = self.create_storage(self.particulator.n_sd)
+        self.tmp_tau = self.create_storage(self.particulator.n_sd)
+        self.tmp_scale = self.create_storage(self.particulator.n_sd)
 
     def __call__(self):
-        self.tmp_data.product(self.terminal_vel_attr.get(), self.volume_attr.get())
-        self.tmp_data *= (
-            self.rho_w  # TODO #798
+        # calculate momentum difference
+        self.tmp_momentum_diff.product(
+            self.terminal_vel_attr.get(), self.volume_attr.get()
+        )
+        self.tmp_momentum_diff *= (
+            self.rho_w
         )  # TODO #798 - we plan to use masses instead of volumes soon
-        self.tmp_data -= self.fall_momentum_attr.get()
-        self.tmp_data *= 1 - exp(-self.particulator.dt / self.tau)
+        self.tmp_momentum_diff -= self.fall_momentum_attr.get()
+
+        if not self.tmp_tau_init or not self.constant:
+            self.tmp_tau_init = True
+            self.calculate_tau(self.tmp_tau, self.sqrt_radius_attr.get())
 
-        self.fall_momentum_attr.data += self.tmp_data
+        self.calculate_scale_factor(self.tmp_scale, self.tmp_tau)
 
+        self.tmp_momentum_diff *= self.tmp_scale
+        self.fall_momentum_attr.data += self.tmp_momentum_diff
         self.particulator.attributes.mark_updated("relative fall momentum")

tests/unit_tests/attributes/test_area_radius.py

@@ -3,18 +3,17 @@
 import pytest
 
 from PySDM import Builder
-from PySDM.backends import CPU
 from PySDM.environments import Box
 
 
 @pytest.mark.parametrize("volume", (np.asarray([44, 666]),))
-def test_radius(volume):
+def test_radius(volume, backend_class):
     # arrange
-    builder = Builder(backend=CPU(), n_sd=volume.size)
+    builder = Builder(backend=backend_class(), n_sd=volume.size)
     builder.set_environment(Box(dt=None, dv=None))
     builder.request_attribute("radius")
     particulator = builder.build(
-        attributes={"volume": [volume], "multiplicity": np.ones_like(volume)}
+        attributes={"volume": volume, "multiplicity": np.ones_like(volume)}
     )
 
     # act
@@ -26,13 +25,33 @@ def test_radius(volume):
 
 
 @pytest.mark.parametrize("volume", (np.asarray([44, 666]),))
-def test_area(volume):
+def test_sqrt_radius(volume, backend_class):
     # arrange
-    builder = Builder(backend=CPU(), n_sd=volume.size)
+    builder = Builder(backend=backend_class(), n_sd=volume.size)
+    builder.set_environment(Box(dt=None, dv=None))
+    builder.request_attribute("radius")
+    builder.request_attribute("square root of radius")
+    particulator = builder.build(
+        attributes={"volume": volume, "multiplicity": np.ones_like(volume)}
+    )
+
+    # act
+    radius_actual = particulator.attributes["radius"].to_ndarray()
+    sqrt_radius_actual = particulator.attributes["square root of radius"].to_ndarray()
+
+    # assert
+    sqrt_radius_expected = np.sqrt(radius_actual)
+    np.testing.assert_allclose(sqrt_radius_actual, sqrt_radius_expected)
+
+
+@pytest.mark.parametrize("volume", (np.asarray([44, 666]),))
+def test_area(volume, backend_class):
+    # arrange
+    builder = Builder(backend=backend_class(), n_sd=volume.size)
     builder.set_environment(Box(dt=None, dv=None))
     builder.request_attribute("area")
     particulator = builder.build(
-        attributes={"volume": [volume], "multiplicity": np.ones_like(volume)}
+        attributes={"volume": volume, "multiplicity": np.ones_like(volume)}
     )
 
     # act

tests/unit_tests/attributes/test_fall_velocity.py

@@ -64,7 +64,7 @@ def test_fall_velocity_calculation(default_attributes, backend_class):
 
     # needed to use relative fall velocity instead of terminal
     # velocity behind the scenes
-    builder.add_dynamic(RelaxedVelocity(tau=1))
+    builder.add_dynamic(RelaxedVelocity())
 
     builder.request_attribute("relative fall velocity")
 
@@ -91,7 +91,7 @@ def test_conservation_of_momentum(default_attributes, backend_class):
     builder.set_environment(Box(dt=1, dv=1))
 
     # add and remove relaxed velocity to prevent warning
-    builder.add_dynamic(RelaxedVelocity(tau=1))
+    builder.add_dynamic(RelaxedVelocity())
 
     builder.request_attribute("relative fall momentum")
 
@@ -139,7 +139,7 @@ def test_attribute_selection(backend_class):
 
     builder = Builder(n_sd=1, backend=backend_class())
     builder.set_environment(Box(dt=1, dv=1))
-    builder.add_dynamic(RelaxedVelocity(tau=1))
+    builder.add_dynamic(RelaxedVelocity())
     builder.request_attribute("relative fall velocity")
 
     # with RelaxedVelocity, the builder should use RelativeFallVelocity

tests/unit_tests/backends/storage/test_arithmetics.py → tests/unit_tests/backends/storage/test_basic_ops.py

@@ -3,7 +3,7 @@
 import pytest
 
 
-class TestArithmetics:  # pylint: disable=too-few-public-methods
+class TestBasicOps:
     @staticmethod
     @pytest.mark.parametrize(
         "output, addend, expected",
@@ -24,3 +24,23 @@ def test_addition(backend_class, output, addend, expected):
 
         # Assert
         np.testing.assert_array_equal(output.to_ndarray(), expected)
+
+    @staticmethod
+    @pytest.mark.parametrize(
+        "data",
+        (
+            [1.0],
+            [2.0, 3, 4],
+            [-1, np.nan, np.inf],
+        ),
+    )
+    def test_exp(backend_class, data):
+        # Arrange
+        backend = backend_class(double_precision=True)
+        output = backend.Storage.from_ndarray(np.asarray(data))
+
+        # Act
+        output.exp()
+
+        # Assert
+        np.testing.assert_array_equal(output.to_ndarray(), np.exp(data))