Add FallVelocity attribute and RelaxedVelocity dynamic (and tests) (

open-atmos#1105)

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

.zenodo.json

@@ -17,6 +17,10 @@
       "name": "Bartman, Piotr",
       "orcid": "0000-0003-0265-6428"
     },
+    {
+      "affiliation": "California Institute of Technology, Pasadena, CA, USA",
+      "name": "Bhalla, Brady"
+    },
     {
       "affiliation": "Jagiellonian University, Kraków, Poland",
       "name": "Bulenok, Oleksii",

PySDM/attributes/impl/dummy_attribute.py

@@ -1,5 +1,7 @@
 """ logic around `PySDM.attributes.impl.dummy_attribute.DummyAttribute` - parent class
 for do-nothing attributes """
+import warnings
+
 import numpy as np
 
 from .attribute import Attribute
@@ -17,8 +19,13 @@ def get(self):
         return self.data
 
 
-def make_dummy_attribute_factory(name):
+def make_dummy_attribute_factory(name, warn=False):
     def _factory(builder):
         return DummyAttribute(builder, name=name)
 
+    if warn:
+        warnings.warn(
+            f"dummy implementation used for requested attribute named '{name}'"
+        )
+
     return _factory

PySDM/attributes/impl/mapper.py

@@ -19,6 +19,7 @@
     Heat,
     Multiplicities,
     Radius,
+    RelativeFallVelocity,
     Temperature,
     TerminalVelocity,
     Volume,
@@ -32,6 +33,7 @@
     OrganicFraction,
 )
 from PySDM.attributes.physics.hygroscopicity import Kappa, KappaTimesDryVolume
+from PySDM.attributes.physics.relative_fall_velocity import RelativeFallMomentum
 from PySDM.dynamics.impl.chemistry_utils import AQUEOUS_COMPOUNDS
 from PySDM.physics.surface_tension import Constant
 
@@ -59,6 +61,16 @@
     "area": lambda _, __: Area,
     "dry radius": lambda _, __: DryRadius,
     "terminal velocity": lambda _, __: TerminalVelocity,
+    "relative fall momentum": lambda dynamics, __: (
+        RelativeFallMomentum
+        if "RelaxedVelocity" in dynamics
+        else make_dummy_attribute_factory("relative fall momentum", warn=True)
+        # note: could eventually make an attribute that calculates momentum
+        # from terminal velocity instead when no RelaxedVelocity dynamic is present
+    ),
+    "relative fall velocity": lambda dynamics, __: (
+        RelativeFallVelocity if "RelaxedVelocity" in dynamics else TerminalVelocity
+    ),
     "cell id": lambda _, __: CellID,
     "cell origin": lambda _, __: CellOrigin,
     "cooling rate": lambda _, __: CoolingRate,

PySDM/attributes/physics/__init__.py

@@ -9,6 +9,7 @@
 from .heat import Heat
 from .multiplicities import Multiplicities
 from .radius import Radius
+from .relative_fall_velocity import RelativeFallMomentum, RelativeFallVelocity
 from .temperature import Temperature
 from .terminal_velocity import TerminalVelocity
 from .volume import Volume

PySDM/attributes/physics/relative_fall_velocity.py

@@ -0,0 +1,28 @@
+"""
+Attributes for tracking droplet velocity
+"""
+
+from PySDM.attributes.impl.derived_attribute import DerivedAttribute
+from PySDM.attributes.impl.extensive_attribute import ExtensiveAttribute
+
+
+class RelativeFallMomentum(ExtensiveAttribute):
+    def __init__(self, builder):
+        super().__init__(builder, name="relative fall momentum", dtype=float)
+
+
+class RelativeFallVelocity(DerivedAttribute):
+    def __init__(self, builder):
+        self.momentum = builder.get_attribute("relative fall momentum")
+        self.volume = builder.get_attribute("volume")
+        self.rho_w = builder.formulae.constants.rho_w  # TODO #798
+
+        super().__init__(
+            builder,
+            name="relative fall velocity",
+            dependencies=(self.momentum, self.volume),
+        )
+
+    def recalculate(self):
+        self.data.ratio(self.momentum.get(), self.volume.get())
+        self.data[:] *= 1 / self.rho_w  # TODO #798

PySDM/dynamics/__init__.py

@@ -10,3 +10,4 @@
 from PySDM.dynamics.displacement import Displacement
 from PySDM.dynamics.eulerian_advection import EulerianAdvection
 from PySDM.dynamics.freezing import Freezing
+from PySDM.dynamics.relaxed_velocity import RelaxedVelocity

PySDM/dynamics/collisions/breakup_fragmentations/lowlist82.py

@@ -22,7 +22,7 @@ def register(self, builder):
         self.const = self.particulator.formulae.constants
         builder.request_attribute("radius")
         builder.request_attribute("volume")
-        builder.request_attribute("terminal velocity")
+        builder.request_attribute("relative fall velocity")
         for key in ("Sc", "St", "tmp", "tmp2", "CKE", "We", "W2", "ds", "dl", "dcoal"):
             self.arrays[key] = self.particulator.PairwiseStorage.empty(
                 self.particulator.n_sd // 2, dtype=float
@@ -64,7 +64,7 @@ def __call__(self, nf, frag_size, u01, is_first_in_pair):
 
         self.arrays["tmp"].sum(self.particulator.attributes["volume"], is_first_in_pair)
         self.arrays["tmp2"].distance(
-            self.particulator.attributes["terminal velocity"], is_first_in_pair
+            self.particulator.attributes["relative fall velocity"], is_first_in_pair
         )
         self.arrays["tmp2"] **= 2
         self.arrays["CKE"].multiply(

PySDM/dynamics/collisions/breakup_fragmentations/straub2010.py

@@ -27,7 +27,7 @@ def register(self, builder):
         self.const = self.particulator.formulae.constants
         builder.request_attribute("radius")
         builder.request_attribute("volume")
-        builder.request_attribute("terminal velocity")
+        builder.request_attribute("relative fall velocity")
         for key in ("Sc", "tmp", "tmp2", "CKE", "We", "gam", "CW", "ds"):
             self.arrays[key] = self.particulator.PairwiseStorage.empty(
                 self.particulator.n_sd // 2, dtype=float
@@ -53,7 +53,7 @@ def __call__(self, nf, frag_size, u01, is_first_in_pair):
             self.const.PI * self.const.sgm_w * (6 / self.const.PI) ** (2 / 3)
         )
         self.arrays["tmp2"].distance(
-            self.particulator.attributes["terminal velocity"], is_first_in_pair
+            self.particulator.attributes["relative fall velocity"], is_first_in_pair
         )
         self.arrays["tmp2"] **= 2
         self.arrays["CKE"].multiply(

PySDM/dynamics/collisions/coalescence_efficiencies/_gravitational.py

@@ -11,7 +11,7 @@ def __init__(self):
     def register(self, builder):
         self.particulator = builder.particulator
         builder.request_attribute("radius")
-        builder.request_attribute("terminal velocity")
+        builder.request_attribute("relative fall velocity")
         self.pair_tmp = self.particulator.PairwiseStorage.empty(
             self.particulator.n_sd // 2, dtype=float
         )

PySDM/dynamics/collisions/coalescence_efficiencies/lowlist1982.py

@@ -29,7 +29,7 @@ def register(self, builder):
         self.const = self.particulator.formulae.constants
         builder.request_attribute("radius")
         builder.request_attribute("volume")
-        builder.request_attribute("terminal velocity")
+        builder.request_attribute("relative fall velocity")
         for key in ("Sc", "St", "dS", "tmp", "tmp2", "CKE", "Et", "ds", "dl"):
             self.arrays[key] = self.particulator.PairwiseStorage.empty(
                 self.particulator.n_sd // 2, dtype=float
@@ -64,7 +64,7 @@ def __call__(self, output, is_first_in_pair):
 
         self.arrays["tmp"].sum(self.particulator.attributes["volume"], is_first_in_pair)
         self.arrays["tmp2"].distance(
-            self.particulator.attributes["terminal velocity"], is_first_in_pair
+            self.particulator.attributes["relative fall velocity"], is_first_in_pair
         )
         self.arrays["tmp2"] **= 2
         self.arrays["CKE"].multiply(

PySDM/dynamics/collisions/coalescence_efficiencies/straub2010.py

@@ -17,7 +17,7 @@ def register(self, builder):
         self.particulator = builder.particulator
         self.const = self.particulator.formulae.constants
         builder.request_attribute("volume")
-        builder.request_attribute("terminal velocity")
+        builder.request_attribute("relative fall velocity")
         for key in ("Sc", "tmp", "tmp2", "We"):
             self.arrays[key] = self.particulator.PairwiseStorage.empty(
                 self.particulator.n_sd // 2, dtype=float
@@ -30,7 +30,7 @@ def __call__(self, output, is_first_in_pair):
         self.arrays["tmp"] *= 2
 
         self.arrays["tmp2"].distance(
-            self.particulator.attributes["terminal velocity"], is_first_in_pair
+            self.particulator.attributes["relative fall velocity"], is_first_in_pair
         )
         self.arrays["tmp2"] **= 2
         self.arrays["We"].multiply(

PySDM/dynamics/collisions/collision_kernels/electric.py

@@ -8,5 +8,5 @@
 class Electric(Parameterized):  # pylint: disable=too-few-public-methods
     def __init__(self):
         super().__init__(
-            (1, 1, -7, 1.78, -20.5, 1.73, 0.26, 1.47, 1, 0.82, -0.003, 4.4, 8)
+            (1, 1, -7, 1.78, -20.5, 1.73, 0.26, 1.47, 1, 0.82, -0.003, 4.4, 8),
         )

PySDM/dynamics/collisions/collision_kernels/geometric.py

@@ -16,6 +16,6 @@ def __call__(self, output, is_first_in_pair):
         output **= 2
         output *= const.PI * self.collection_efficiency
         self.pair_tmp.distance(
-            self.particulator.attributes["terminal velocity"], is_first_in_pair
+            self.particulator.attributes["relative fall velocity"], is_first_in_pair
         )
         output *= self.pair_tmp

PySDM/dynamics/collisions/collision_kernels/hydrodynamic.py

@@ -8,4 +8,6 @@
 
 class Hydrodynamic(Parameterized):  # pylint: disable=too-few-public-methods
     def __init__(self):
-        super().__init__((1, 1, -27, 1.65, -58, 1.9, 15, 1.13, 16.7, 1, 0.004, 4, 8))
+        super().__init__(
+            (1, 1, -27, 1.65, -58, 1.9, 15, 1.13, 16.7, 1, 0.004, 4, 8),
+        )

PySDM/dynamics/collisions/collision_kernels/impl/gravitational.py

@@ -9,7 +9,7 @@ def __init__(self):
     def register(self, builder):
         self.particulator = builder.particulator
         builder.request_attribute("radius")
-        builder.request_attribute("terminal velocity")
+        builder.request_attribute("relative fall velocity")
         self.pair_tmp = self.particulator.PairwiseStorage.empty(
             self.particulator.n_sd // 2, dtype=float
         )

PySDM/dynamics/collisions/collision_kernels/impl/parameterized.py

@@ -26,6 +26,6 @@ def __call__(self, output, is_first_in_pair):
         output *= self.pair_tmp
 
         self.pair_tmp.distance(
-            self.particulator.attributes["terminal velocity"], is_first_in_pair
+            self.particulator.attributes["relative fall velocity"], is_first_in_pair
         )
         output *= self.pair_tmp

PySDM/dynamics/displacement.py

@@ -20,7 +20,7 @@ def __init__(
         precipitation_counting_level_index: int = 0,
         adaptive=DEFAULTS.adaptive,
         rtol=DEFAULTS.rtol,
-    ):
+    ):  # pylint: disable=too-many-arguments
         self.particulator = None
         self.enable_sedimentation = enable_sedimentation
         self.dimension = None
@@ -36,7 +36,7 @@ def __init__(
         self._n_substeps = 1
 
     def register(self, builder):
-        builder.request_attribute("terminal velocity")
+        builder.request_attribute("relative fall velocity")
         self.particulator = builder.particulator
         self.dimension = len(builder.particulator.environment.mesh.grid)
         self.grid = self.particulator.Storage.from_ndarray(
@@ -130,7 +130,7 @@ def calculate_displacement(
             dt = self.particulator.dt / self._n_substeps
             dt_over_dz = dt / self.particulator.mesh.dz
             displacement_z *= 1 / dt_over_dz
-            displacement_z -= self.particulator.attributes["terminal velocity"]
+            displacement_z -= self.particulator.attributes["relative fall velocity"]
             displacement_z *= dt_over_dz
 
     @staticmethod

PySDM/dynamics/relaxed_velocity.py

@@ -0,0 +1,53 @@
+"""
+A dynamic which relaxes
+`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`
+    """
+
+    def __init__(self, tau):
+        self.tau: float = tau
+
+        self.particulator = None
+        self.fall_momentum_attr = None
+        self.terminal_vel_attr = None
+        self.volume_attr = None
+        self.rho_w = None  # TODO #798 - we plan to use masses instead of volumes soon
+        self.tmp_data = None
+
+    def register(self, builder):
+        self.particulator: Particulator = builder.particulator
+
+        self.fall_momentum_attr: Attribute = builder.get_attribute(
+            "relative fall momentum"
+        )
+        self.terminal_vel_attr: Attribute = builder.get_attribute("terminal velocity")
+        self.volume_attr: Attribute = builder.get_attribute("volume")
+
+        self.rho_w: float = builder.formulae.constants.rho_w  # TODO #798
+
+        self.tmp_data = self.particulator.Storage.empty(
+            (self.particulator.n_sd,), dtype=float
+        )
+
+    def __call__(self):
+        self.tmp_data.product(self.terminal_vel_attr.get(), self.volume_attr.get())
+        self.tmp_data *= (
+            self.rho_w  # TODO #798
+        )  # 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.fall_momentum_attr.data += self.tmp_data
+
+        self.particulator.attributes.mark_updated("relative fall momentum")

PySDM/initialisation/__init__.py

@@ -4,3 +4,4 @@
 """
 from .discretise_multiplicities import discretise_multiplicities
 from .equilibrate_wet_radii import equilibrate_wet_radii
+from .init_fall_momenta import init_fall_momenta

PySDM/initialisation/init_fall_momenta.py

@@ -0,0 +1,47 @@
+"""
+Initialize the `PySDM.attributes.physics.relative_fall_velocity.RelativeFallMomentum`
+of droplets
+"""
+
+import numpy as np
+
+from PySDM.backends import CPU
+from PySDM.dynamics.terminal_velocity import Interpolation
+from PySDM.formulae import Formulae
+from PySDM.particulator import Particulator
+
+
+def init_fall_momenta(
+    volume: np.ndarray,
+    rho_w: float,  # TODO #798 - we plan to use masses instead of volumes soon
+    zero: bool = False,
+    terminal_velocity_approx=Interpolation,
+):
+    """
+    Calculate default values of the
+    `PySDM.attributes.physics.relative_fall_velocity.RelativeFallMomentum` attribute
+    (needed when using
+    `PySDM.attributes.physics.relative_fall_velocity.RelativeFallVelocity` attribute)
+
+    Parameters:
+        - volume: a numpy array of superdroplet volumes
+        - rho_w: the density of water (generally found in formulae.constants)
+
+    Returns:
+        - a numpy array of initial momentum values
+    """
+    if zero:
+        return np.zeros_like(volume)
+
+    particulator = Particulator(0, CPU(Formulae()))
+
+    approximation = terminal_velocity_approx(particulator=particulator)
+
+    radii_arr = particulator.formulae.trivia.radius(volume)
+    radii = particulator.Storage.from_ndarray(radii_arr)
+
+    output = particulator.Storage.empty((len(volume),), dtype=float)
+
+    approximation(output=output, radius=radii)
+
+    return output.to_ndarray() * volume * rho_w  # TODO #798