From 03cbbcfb02c6a330da52f4df5b5324bc5e73974f Mon Sep 17 00:00:00 2001 From: Brady Bhalla <34150846+bradybhalla@users.noreply.github.com> Date: Thu, 14 Sep 2023 08:36:55 -0600 Subject: [PATCH] Make the relaxed velocity dynamic use a non-constant timescale (#1114) Co-authored-by: Sylwester Arabas --- PySDM/attributes/impl/mapper.py | 2 + PySDM/attributes/physics/__init__.py | 2 +- PySDM/attributes/physics/radius.py | 15 ++++ PySDM/backends/impl_numba/storage.py | 3 + PySDM/backends/impl_thrust_rtc/storage.py | 17 +++++ PySDM/dynamics/relaxed_velocity.py | 74 +++++++++++++++---- .../unit_tests/attributes/test_area_radius.py | 33 +++++++-- .../attributes/test_fall_velocity.py | 6 +- ...{test_arithmetics.py => test_basic_ops.py} | 22 +++++- .../dynamics/test_relaxed_velocity.py | 65 ++++++++++++++-- 10 files changed, 203 insertions(+), 36 deletions(-) rename tests/unit_tests/backends/storage/{test_arithmetics.py => test_basic_ops.py} (59%) diff --git a/PySDM/attributes/impl/mapper.py b/PySDM/attributes/impl/mapper.py index 4f1b420b1..1590601ad 100644 --- a/PySDM/attributes/impl/mapper.py +++ b/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, diff --git a/PySDM/attributes/physics/__init__.py b/PySDM/attributes/physics/__init__.py index bbb93f68d..fec980a51 100644 --- a/PySDM/attributes/physics/__init__.py +++ b/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 diff --git a/PySDM/attributes/physics/radius.py b/PySDM/attributes/physics/radius.py index bec331209..1e04c217e 100644 --- a/PySDM/attributes/physics/radius.py +++ b/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 diff --git a/PySDM/backends/impl_numba/storage.py b/PySDM/backends/impl_numba/storage.py index fb00b0836..735af4c7b 100644 --- a/PySDM/backends/impl_numba/storage.py +++ b/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) diff --git a/PySDM/backends/impl_thrust_rtc/storage.py b/PySDM/backends/impl_thrust_rtc/storage.py index 0e634e201..f104efa82 100644 --- a/PySDM/backends/impl_thrust_rtc/storage.py +++ b/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 diff --git a/PySDM/dynamics/relaxed_velocity.py b/PySDM/dynamics/relaxed_velocity.py index 9ad86b294..bffd92765 100644 --- a/PySDM/dynamics/relaxed_velocity.py +++ b/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") diff --git a/tests/unit_tests/attributes/test_area_radius.py b/tests/unit_tests/attributes/test_area_radius.py index 3904e26bb..c99338082 100644 --- a/tests/unit_tests/attributes/test_area_radius.py +++ b/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 diff --git a/tests/unit_tests/attributes/test_fall_velocity.py b/tests/unit_tests/attributes/test_fall_velocity.py index d6a89f557..a1811d281 100644 --- a/tests/unit_tests/attributes/test_fall_velocity.py +++ b/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 diff --git a/tests/unit_tests/backends/storage/test_arithmetics.py b/tests/unit_tests/backends/storage/test_basic_ops.py similarity index 59% rename from tests/unit_tests/backends/storage/test_arithmetics.py rename to tests/unit_tests/backends/storage/test_basic_ops.py index b758baafd..585736d6d 100644 --- a/tests/unit_tests/backends/storage/test_arithmetics.py +++ b/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)) diff --git a/tests/unit_tests/dynamics/test_relaxed_velocity.py b/tests/unit_tests/dynamics/test_relaxed_velocity.py index 2aa1a448a..1bed5beb4 100644 --- a/tests/unit_tests/dynamics/test_relaxed_velocity.py +++ b/tests/unit_tests/dynamics/test_relaxed_velocity.py @@ -34,9 +34,17 @@ def default_attributes_fixture(request): return request.param -def test_small_timescale(default_attributes, backend_class): +@pytest.fixture( + name="constant_timescale", + params=(True, False), +) +def constant_timescale_fixture(request): + return request.param + + +def test_small_timescale(default_attributes, constant_timescale, backend_class): """ - When the fall velocity is initialized to 0 and tau is very small, + When the fall velocity is initialized to 0 and relaxation is very quick, the velocity should quickly approach the terminal velocity """ @@ -46,7 +54,7 @@ def test_small_timescale(default_attributes, backend_class): builder.set_environment(Box(dt=1, dv=1)) - builder.add_dynamic(RelaxedVelocity(tau=1e-12)) + builder.add_dynamic(RelaxedVelocity(c=1e-12, constant=constant_timescale)) builder.request_attribute("relative fall velocity") builder.request_attribute("terminal velocity") @@ -65,9 +73,9 @@ def test_small_timescale(default_attributes, backend_class): ) -def test_large_timescale(default_attributes, backend_class): +def test_large_timescale(default_attributes, constant_timescale, backend_class): """ - When the fall velocity is initialized to 0 and tau is very large, + When the fall velocity is initialized to 0 and relaxation is very slow, the velocity should remain 0 """ @@ -77,7 +85,7 @@ def test_large_timescale(default_attributes, backend_class): builder.set_environment(Box(dt=1, dv=1)) - builder.add_dynamic(RelaxedVelocity(tau=1e12)) + builder.add_dynamic(RelaxedVelocity(c=1e15, constant=constant_timescale)) builder.request_attribute("relative fall velocity") builder.request_attribute("terminal velocity") @@ -96,7 +104,7 @@ def test_large_timescale(default_attributes, backend_class): ) -def test_behavior(default_attributes, backend_class): +def test_behavior(default_attributes, constant_timescale, backend_class): """ The fall velocity should approach the terminal velocity exponentially """ @@ -107,7 +115,8 @@ def test_behavior(default_attributes, backend_class): builder.set_environment(Box(dt=1, dv=1)) - builder.add_dynamic(RelaxedVelocity(tau=0.5)) + # relaxation happens too quickly unless c is high enough + builder.add_dynamic(RelaxedVelocity(c=100, constant=constant_timescale)) builder.request_attribute("relative fall velocity") builder.request_attribute("terminal velocity") @@ -138,3 +147,43 @@ def test_behavior(default_attributes, backend_class): # for an exponential decay, the ratio should be roughly constant using constant timesteps assert (np.abs(delta_v1 / delta_v2 - delta_v2 / delta_v3) < 0.01).all() + + +@pytest.mark.parametrize("c", [0.1, 10]) +def test_timescale(default_attributes, c, constant_timescale, backend_class): + """ + The non-constant timescale should be proportional to the sqrt of the radius. The + proportionality constant should be the parameter for the dynamic. + + The constant timescale should be constant. + """ + + builder = Builder( + n_sd=len(default_attributes["multiplicity"]), backend=backend_class() + ) + + builder.set_environment(Box(dt=1, dv=1)) + + sqrt_radius_attr = builder.get_attribute("square root of radius") + + dyn = RelaxedVelocity(c=c, constant=constant_timescale) + builder.add_dynamic(dyn) + + default_attributes["relative fall momentum"] = np.zeros_like( + default_attributes["multiplicity"] + ) + + particulator = builder.build(attributes=default_attributes, products=()) + + tau_storage = particulator.Storage.empty( + default_attributes["multiplicity"].shape, dtype=float + ) + dyn.calculate_tau(tau_storage, sqrt_radius_attr.get()) + + # expected_c should be whatever c was set to in the dynamic + if not constant_timescale: + expected_c = tau_storage.to_ndarray() / sqrt_radius_attr.get().to_ndarray() + else: + expected_c = tau_storage.to_ndarray() + + assert np.allclose(expected_c, c)