switching from volume to mass condensation coordinates, handling dm_dt(r_dr_dt) in particle_shape_and_density (#1538)

Author: slayoo

.github/workflows/tests.yml

@@ -271,8 +271,8 @@ jobs:
           python -m pytest --durations=10 -v -p no:unraisableexception -We tests/examples_tests/test_tests_completeness.py
 
        # TODO #1207
-      - if: startsWith(matrix.platform, 'windows-')
+      - if: startsWith(matrix.platform, 'macos-13')
         run: python -m pytest --durations=10 -v -p no:unraisableexception -We tests/examples_tests/test_run* -k "not Rozanski_and_Sonntag_1982" --suite ${{ matrix.test-suite }}
 
-      - if: ( ! startsWith(matrix.platform, 'windows-') )
+      - if: ( ! startsWith(matrix.platform, 'macos-13') )
         run: python -m pytest --durations=10 -v -p no:unraisableexception -We tests/examples_tests/test_run* --suite ${{ matrix.test-suite }}

PySDM/backends/impl_numba/methods/condensation_methods.py

@@ -390,14 +390,18 @@ def minfun(  # pylint: disable=too-many-arguments,too-many-locals
             K,
             ventilation_factor,
         ):
-            volume = formulae.condensation_coordinate__volume(x_new)
+            if x_new > formulae.condensation_coordinate__x_max():
+                return x_old - x_new
+            mass_new = formulae.condensation_coordinate__mass(x_new)
+            volume_new = formulae.particle_shape_and_density__mass_to_volume(mass_new)
+            r_new = formulae.trivia__radius(volume_new)
             RH_eq = formulae.hygroscopicity__RH_eq(
-                formulae.trivia__radius(volume),
+                r_new,
                 temperature,
                 kappa,
                 rd3,
                 formulae.surface_tension__sigma(
-                    temperature, volume, formulae.constants.PI_4_3 * rd3, f_org
+                    temperature, volume_new, formulae.constants.PI_4_3 * rd3, f_org
                 ),
             )
             r_dr_dt = formulae.drop_growth__r_dr_dt(
@@ -410,10 +414,11 @@ def minfun(  # pylint: disable=too-many-arguments,too-many-locals
                 K,
                 ventilation_factor,
             )
+            dm_dt = formulae.particle_shape_and_density__dm_dt(r=r_new, r_dr_dt=r_dr_dt)
             return (
                 x_old
                 - x_new
-                + timestep * formulae.condensation_coordinate__dx_dt(x_new, r_dr_dt)
+                + timestep * formulae.condensation_coordinate__dx_dt(mass_new, dm_dt)
             )
 
         @numba.njit(**jit_flags)
@@ -440,15 +445,17 @@ def calculate_ml_new(  # pylint: disable=too-many-branches,too-many-arguments,to
             lambdaK = formulae.diffusion_kinetics__lambdaK(T, p)
             lambdaD = formulae.diffusion_kinetics__lambdaD(DTp, T)
             for drop in cell_idx:
+                if attributes.water_mass[drop] <= 0:
+                    continue
                 v_drop = formulae.particle_shape_and_density__mass_to_volume(
                     attributes.water_mass[drop]
                 )
-                if v_drop <= 0:
-                    continue
-                x_old = formulae.condensation_coordinate__x(v_drop)
+                x_old = formulae.condensation_coordinate__x(attributes.water_mass[drop])
                 r_old = formulae.trivia__radius(v_drop)
                 x_insane = formulae.condensation_coordinate__x(
-                    attributes.vdry[drop] / 100
+                    formulae.particle_shape_and_density__volume_to_mass(
+                        attributes.vdry[drop] / 100
+                    )
                 )
                 rd3 = attributes.vdry[drop] / formulae.constants.PI_4_3
                 sgm = formulae.surface_tension__sigma(
@@ -492,8 +499,12 @@ def calculate_ml_new(  # pylint: disable=too-many-branches,too-many-arguments,to
                         Kr,
                         ventilation_factor,
                     )
+                    mass_old = formulae.condensation_coordinate__mass(x_old)
+                    dm_dt_old = formulae.particle_shape_and_density__dm_dt(
+                        r=r_old, r_dr_dt=r_dr_dt_old
+                    )
                     dx_old = timestep * formulae.condensation_coordinate__dx_dt(
-                        x_old, r_dr_dt_old
+                        mass_old, dm_dt_old
                     )
                 else:
                     dx_old = 0.0
@@ -561,9 +572,7 @@ def calculate_ml_new(  # pylint: disable=too-many-branches,too-many-arguments,to
                     else:
                         x_new = x_old
 
-                mass_new = formulae.particle_shape_and_density__volume_to_mass(
-                    formulae.condensation_coordinate__volume(x_new)
-                )
+                mass_new = formulae.condensation_coordinate__mass(x_new)
                 mass_cr = formulae.particle_shape_and_density__volume_to_mass(
                     attributes.v_cr[drop]
                 )

PySDM/backends/impl_numba/test_helpers/scipy_ode_condensation_solver.py

@@ -92,15 +92,7 @@ def _make_solve(formulae):  # pylint: disable=too-many-statements,too-many-local
 
     @numba.njit(**{**JIT_FLAGS, **{"parallel": False}})
     def _liquid_water_mixing_ratio(n, x, m_d_mean):
-        return (
-            np.sum(
-                n
-                * jit_formulae.particle_shape_and_density__volume_to_mass(
-                    jit_formulae.condensation_coordinate__volume(x)
-                )
-            )
-            / m_d_mean
-        )
+        return np.sum(n * jit_formulae.condensation_coordinate__mass(x)) / m_d_mean
 
     @numba.njit(**{**JIT_FLAGS, **{"parallel": False}})
     def _impl(  # pylint: disable=too-many-arguments,too-many-locals
@@ -133,8 +125,10 @@ def _impl(  # pylint: disable=too-many-arguments,too-many-locals
             diffusivity=DTp,
             density=air_density,
         )
+        sum_n_dm_dt = 0
         for i, x_i in enumerate(x):
-            v = jit_formulae.condensation_coordinate__volume(x_i)
+            m = jit_formulae.condensation_coordinate__mass(x_i)
+            v = jit_formulae.particle_shape_and_density__mass_to_volume(m)
             r = jit_formulae.trivia__radius(v)
             Dr = jit_formulae.diffusion_kinetics__D(DTp, r, lambdaD)
             Kr = jit_formulae.diffusion_kinetics__K(KTp, r, lambdaK)
@@ -145,34 +139,23 @@ def _impl(  # pylint: disable=too-many-arguments,too-many-locals
                 )
             )
             sgm = jit_formulae.surface_tension__sigma(T, v, dry_volume[i], f_org[i])
-            dy_dt[idx_x + i] = jit_formulae.condensation_coordinate__dx_dt(
-                x_i,
-                jit_formulae.drop_growth__r_dr_dt(
-                    jit_formulae.hygroscopicity__RH_eq(
-                        r, T, kappa[i], dry_volume[i] / PI_4_3, sgm
-                    ),
-                    T,
-                    RH,
-                    lv,
-                    pvs,
-                    Dr,
-                    Kr,
-                    ventilation_factor,
+            r_dr_dt = jit_formulae.drop_growth__r_dr_dt(
+                jit_formulae.hygroscopicity__RH_eq(
+                    r, T, kappa[i], dry_volume[i] / PI_4_3, sgm
                 ),
+                T,
+                RH,
+                lv,
+                pvs,
+                Dr,
+                Kr,
+                ventilation_factor,
             )
-        d_water_vapour_mixing_ratio__dt = (
-            dot_water_vapour_mixing_ratio
-            - np.sum(
-                n
-                * jit_formulae.particle_shape_and_density__volume_to_mass(
-                    jit_formulae.condensation_coordinate__volume(x)
-                )
-                * dy_dt[idx_x:]
-            )
-            / m_d_mean
-        )
+            dm_dt = jit_formulae.particle_shape_and_density__dm_dt(r, r_dr_dt)
+            dy_dt[idx_x + i] = jit_formulae.condensation_coordinate__dx_dt(m, dm_dt)
+            sum_n_dm_dt += n[i] * dm_dt
         dy_dt[idx_thd] = dot_thd + jit_formulae.state_variable_triplet__dthd_dt(
-            rhod, thd, T, d_water_vapour_mixing_ratio__dt, lv
+            rhod, thd, T, dot_water_vapour_mixing_ratio - sum_n_dm_dt / m_d_mean, lv
         )
 
     @numba.njit(**{**JIT_FLAGS, **{"parallel": False}})
@@ -252,9 +235,7 @@ def solve(  # pylint: disable=too-many-arguments,too-many-locals,unused-argument
         y0 = np.empty(n_sd_in_cell + idx_x)
         y0[idx_thd] = thd
         y0[idx_x:] = jit_formulae.condensation_coordinate__x(
-            jit_formulae.particle_shape_and_density__mass_to_volume(
-                attributes.water_mass[cell_idx]
-            )
+            attributes.water_mass[cell_idx]
         )
         total_water_mixing_ratio = (
             water_vapour_mixing_ratio
@@ -301,9 +282,7 @@ def solve(  # pylint: disable=too-many-arguments,too-many-locals,unused-argument
         m_new = 0
         for i in range(n_sd_in_cell):
             attributes.water_mass[cell_idx[i]] = (
-                jit_formulae.particle_shape_and_density__volume_to_mass(
-                    jit_formulae.condensation_coordinate__volume(y1[idx_x + i])
-                )
+                jit_formulae.condensation_coordinate__mass(y1[idx_x + i])
             )
             m_new += (
                 attributes.multiplicity[cell_idx[i]]

PySDM/backends/impl_thrust_rtc/methods/condensation_methods.py

@@ -100,11 +100,15 @@ def __update_drop_masses(self):
             struct Minfun {{
                 static __device__ real_type value(real_type x_new, void* args_p) {{
                     auto args = static_cast<real_type*>(args_p);
-                    auto vol = {phys.condensation_coordinate.volume.c_inline(x="x_new")};
-                    auto r_new = {phys.trivia.radius.c_inline(volume="vol")};
+                    if (x_new > {phys.condensation_coordinate.x_max.c_inline()}) {{
+                        return {args("x_old")} - x_new;
+                    }}
+                    auto m_new = {phys.condensation_coordinate.mass.c_inline(x="x_new")};
+                    auto v_new = {phys.particle_shape_and_density.mass_to_volume.c_inline(mass="m_new")};
+                    auto r_new = {phys.trivia.radius.c_inline(volume="v_new")};
                     auto sgm = {phys.surface_tension.sigma.c_inline(
                         T=args('_T'),
-                        v_wet="vol",
+                        v_wet="v_new",
                         v_dry=f"const.PI_4_3 * {args('rd3')}",
                         f_org=args("f_org")
                     )};
@@ -125,8 +129,11 @@ def __update_drop_masses(self):
                         K=args("Kr"),
                         ventilation_factor=args("ventilation_factor"),
                     )};
+                    auto dm_dt = {phys.particle_shape_and_density.dm_dt.c_inline(
+                        r="r_new", r_dr_dt="r_dr_dt"
+                    )};
                     return {args("x_old")} - x_new + {args("dt")} * {
-                        phys.condensation_coordinate.dx_dt.c_inline(x="x_new", r_dr_dt="r_dr_dt")
+                        phys.condensation_coordinate.dx_dt.c_inline(m="m_new", dm_dt="dm_dt")
                     };
                 }}
             }};
@@ -145,9 +152,10 @@ def __update_drop_masses(self):
             auto v_old = {phys.particle_shape_and_density.mass_to_volume.c_inline(
                 mass="water_mass[i]"
             )};
-            auto x_old = {phys.condensation_coordinate.x.c_inline(volume="v_old")};
+            auto x_old = {phys.condensation_coordinate.x.c_inline(mass="water_mass[i]")};
             auto r_old = {phys.trivia.radius.c_inline(volume="v_old")};
-            auto x_insane = {phys.condensation_coordinate.x.c_inline(volume="vdry[i]/100")};
+            auto m_insane = {phys.particle_shape_and_density.volume_to_mass.c_inline(volume="vdry[i] / 100")};
+            auto x_insane = {phys.condensation_coordinate.x.c_inline(mass="m_insane")};
             auto rd3 = vdry[i] / {const.PI_4_3};
             auto sgm = {phys.surface_tension.sigma.c_inline(
                 T="_T", v_wet="v", v_dry="vdry[i]", f_org="_f_org[i]"
@@ -161,6 +169,7 @@ def __update_drop_masses(self):
             real_type qrt_re_times_cbrt_sc=0;
             real_type ventilation_factor=0;
             real_type r_dr_dt_old=0;
+            real_type dm_dt_old=0;
             real_type dx_old=0;
 
             real_type x_new = 0;
@@ -184,8 +193,9 @@ def __update_drop_masses(self):
                     RH_eq="RH_eq", T="_T", RH="_RH", lv="_lv", pvs="_pvs", D="Dr", K="Kr",
                     ventilation_factor="ventilation_factor",
                 )};
+                dm_dt_old = {phys.particle_shape_and_density.dm_dt.c_inline(r="r_old", r_dr_dt="r_dr_dt_old")};
                 dx_old = dt * {phys.condensation_coordinate.dx_dt.c_inline(
-                    x="x_old", r_dr_dt="r_dr_dt_old"
+                    m="water_mass[i]", dm_dt="dm_dt_old"
                 )};
             }}
             else {{
@@ -239,10 +249,7 @@ def __update_drop_masses(self):
                     x_new = x_old;
                 }}
             }}
-            auto v_new = {phys.condensation_coordinate.volume.c_inline(x="x_new")};
-            water_mass[i] = {phys.particle_shape_and_density.volume_to_mass.c_inline(
-                volume="v_new"
-            )};
+            water_mass[i] = {phys.condensation_coordinate.mass.c_inline(x="x_new")};
         """.replace(
                 "real_type", self._get_c_type()
             ),

PySDM/formulae.py

@@ -30,7 +30,7 @@ def __init__(  # pylint: disable=too-many-locals
         constants: Optional[dict] = None,
         seed: int = None,
         fastmath: bool = True,
-        condensation_coordinate: str = "VolumeLogarithm",
+        condensation_coordinate: str = "WaterMassLogarithm",
         saturation_vapour_pressure: str = "FlatauWalkoCotton",
         latent_heat: str = "Kirchhoff",
         hygroscopicity: str = "KappaKoehlerLeadingTerms",
@@ -203,9 +203,11 @@ def _formula(func, constants, dimensional_analysis, **kw):
     source = re.sub(r"\n\s+\):", "):", source)
     loc = {}
     for arg_name in special_params:
-        source = source.replace(
-            f"def {func.__name__}({arg_name},", f"def {func.__name__}("
-        )
+        for sep in ",", ")":
+            source = source.replace(
+                f"def {func.__name__}({arg_name}{sep}",
+                f"def {func.__name__}({')' if sep == ')' else ''}",
+            )
 
     extras = func.__extras if hasattr(func, "__extras") else {}
     exec(  # pylint:disable=exec-used
@@ -281,7 +283,8 @@ def _c_inline(fun, return_type=None, constants=None, **args):
         if stripped.startswith("//"):
             continue
         if stripped.startswith('"""'):
-            in_docstring = True
+            if not (stripped.endswith('"""') and len(stripped) >= 6):
+                in_docstring = True
             continue
         if stripped.startswith("def "):
             continue

PySDM/physics/condensation_coordinate/__init__.py

@@ -2,5 +2,5 @@
 definitions of particle-size coordinates for the condensation solver
 """
 
-from .volume import Volume
-from .volume_logarithm import VolumeLogarithm
+from .water_mass import WaterMass
+from .water_mass_logarithm import WaterMassLogarithm

PySDM/physics/condensation_coordinate/volume.py

@@ -0,0 +1,25 @@
+"""
+particle water mass as condensation coordinate (i.e. no transformation)
+"""
+
+
+class WaterMass:
+    def __init__(self, _):
+        pass
+
+    @staticmethod
+    def dx_dt(m, dm_dt):  # pylint: disable=unused-argument
+        return dm_dt
+
+    @staticmethod
+    def mass(x):
+        return x
+
+    @staticmethod
+    def x(mass):
+        return mass
+
+    @staticmethod
+    def x_max(const):
+        """1 kg droplet!"""
+        return const.ONE