Introduce Rogers & Yau notation for terms associated with mass (Fd) and heat (Fk) diffusion in drop growth formula; avoid recomputing them within condensation implicit-Euler root-finding problem; var naming refactors in condensation solver (#1620)

Author: AgnieszkaZaba

PySDM/backends/impl_numba/methods/condensation_methods.py

@@ -188,7 +188,7 @@ def make_adapt_substeps(
         n_substeps_min = math.ceil(timestep / dt_range[1])
 
         @numba.njit(**jit_flags)
-        def adapt_substeps(args, n_substeps, thd, rtol_thd):
+        def adapt_substeps(step_impl_args, n_substeps, thd, rtol_thd):
             n_substeps = np.maximum(n_substeps_min, n_substeps // multiplier)
             success = False
             for burnout in range(fuse + 1):
@@ -202,15 +202,15 @@ def adapt_substeps(args, n_substeps, thd, rtol_thd):
                         ),
                         return_value=(0, False),
                     )
-                thd_new_long, success = step_fake(args, timestep, n_substeps)
+                thd_new_long, success = step_fake(step_impl_args, timestep, n_substeps)
                 if success:
                     break
                 n_substeps *= multiplier
             for burnout in range(fuse + 1):
                 if burnout == fuse:
                     return warn("burnout (short)", __file__, return_value=(0, False))
                 thd_new_short, success = step_fake(
-                    args, timestep, n_substeps * multiplier
+                    step_impl_args, timestep, n_substeps * multiplier
                 )
                 if not success:
                     return warn("short failed", __file__, return_value=(0, False))
@@ -230,18 +230,18 @@ def adapt_substeps(args, n_substeps, thd, rtol_thd):
     @staticmethod
     def make_step_fake(jit_flags, step_impl):
         @numba.njit(**jit_flags)
-        def step_fake(args, dt, n_substeps):
+        def step_fake(step_impl_args, dt, n_substeps):
             dt /= n_substeps
-            _, thd_new, _, _, _, _, success = step_impl(*args, dt, 1, True)
+            _, thd_new, _, _, _, _, success = step_impl(*step_impl_args, dt, 1, True)
             return thd_new, success
 
         return step_fake
 
     @staticmethod
     def make_step(jit_flags, step_impl):
         @numba.njit(**jit_flags)
-        def step(args, dt, n_substeps):
-            return step_impl(*args, dt, n_substeps, False)
+        def step(step_impl_args, dt, n_substeps):
+            return step_impl(*step_impl_args, dt, n_substeps, False)
 
         return step
 
@@ -293,6 +293,7 @@ def step_impl(  # pylint: disable=too-many-arguments,too-many-locals
                 lv = formulae.latent_heat_vapourisation__lv(T)
                 pvs = formulae.saturation_vapour_pressure__pvs_water(T)
                 DTp = formulae.diffusion_thermics__D(T, p)
+                KTp = formulae.diffusion_thermics__K(T, p)
                 RH = pv / pvs
                 Sc = formulae.trivia__air_schmidt_number(
                     dynamic_viscosity=air_dynamic_viscosity,
@@ -317,7 +318,7 @@ def step_impl(  # pylint: disable=too-many-arguments,too-many-locals
                     lv,
                     pvs,
                     DTp,
-                    formulae.diffusion_thermics__K(T, p),
+                    KTp,
                     rtol_x,
                 )
                 dml_dt = (ml_new - ml_old) / timestep
@@ -376,21 +377,12 @@ def make_calculate_ml_new(  # pylint: disable=too-many-statements
     ):
         @numba.njit(**jit_flags)
         def minfun(  # pylint: disable=too-many-arguments,too-many-locals
-            x_new,
-            x_old,
-            timestep,
-            kappa,
-            f_org,
-            rd3,
-            temperature,
-            RH,
-            lv,
-            pvs,
-            D,
-            K,
-            mass_ventilation_factor,
-            heat_ventilation_factor,
+            x_new, x_old, timestep, kappa, f_org, rd3, temperature, RH, Fk, Fd
         ):
+            """
+            root finding problem for the implicit-in-x Euler step
+            neglecting dependence of `Fk` and `Fd` on particle size
+            """
             if x_new > formulae.diffusion_coordinate__x_max():
                 return x_old - x_new
             mass_new = formulae.diffusion_coordinate__mass(x_new)
@@ -405,15 +397,7 @@ def minfun(  # pylint: disable=too-many-arguments,too-many-locals
                     temperature, volume_new, formulae.constants.PI_4_3 * rd3, f_org
                 ),
             )
-            r_dr_dt = formulae.drop_growth__r_dr_dt(
-                RH_eq,
-                temperature,
-                RH,
-                lv,
-                pvs,
-                D * mass_ventilation_factor,
-                K * heat_ventilation_factor,
-            )
+            r_dr_dt = formulae.drop_growth__r_dr_dt(RH_eq=RH_eq, RH=RH, Fk=Fk, Fd=Fd)
             dm_dt = formulae.particle_shape_and_density__dm_dt(r=r_new, r_dr_dt=r_dr_dt)
             return (
                 x_old
@@ -476,29 +460,25 @@ def calculate_ml_new(  # pylint: disable=too-many-branches,too-many-arguments,to
                         )
                     )
                     heat_ventilation_factor = mass_ventilation_factor  # TODO #1588
-                    args = (
+                    Fk = formulae.drop_growth__Fk(
+                        T=T, K=Kr * heat_ventilation_factor, lv=lv
+                    )
+                    Fd = formulae.drop_growth__Fd(
+                        T=T, D=Dr * mass_ventilation_factor, pvs=pvs
+                    )
+                    minfun_args = (
                         x_old,
                         timestep,
                         attributes.kappa[drop],
                         attributes.f_org[drop],
                         rd3,
                         T,
                         RH,
-                        lv,
-                        pvs,
-                        Dr,
-                        Kr,
-                        mass_ventilation_factor,
-                        heat_ventilation_factor,
+                        Fk,
+                        Fd,
                     )
                     r_dr_dt_old = formulae.drop_growth__r_dr_dt(
-                        RH_eq,
-                        T,
-                        RH,
-                        lv,
-                        pvs,
-                        mass_ventilation_factor * Dr,
-                        heat_ventilation_factor * Kr,
+                        RH_eq=RH_eq, RH=RH, Fk=Fk, Fd=Fd
                     )
                     mass_old = formulae.diffusion_coordinate__mass(x_old)
                     dm_dt_old = formulae.particle_shape_and_density__dm_dt(
@@ -514,8 +494,8 @@ def calculate_ml_new(  # pylint: disable=too-many-branches,too-many-arguments,to
                 else:
                     a = x_old
                     b = max(x_insane, a + dx_old)
-                    fa = minfun(a, *args)
-                    fb = minfun(b, *args)
+                    fa = minfun(a, *minfun_args)
+                    fb = minfun(b, *minfun_args)
 
                     counter = 0
                     while not fa * fb < 0:
@@ -545,7 +525,7 @@ def calculate_ml_new(  # pylint: disable=too-many-branches,too-many-arguments,to
                             success = False
                             break
                         b = max(x_insane, a + math.ldexp(dx_old, counter))
-                        fb = minfun(b, *args)
+                        fb = minfun(b, *minfun_args)
 
                     if not success:
                         break
@@ -556,7 +536,7 @@ def calculate_ml_new(  # pylint: disable=too-many-branches,too-many-arguments,to
 
                         x_new, iters_taken = toms748_solve(
                             minfun,
-                            args,
+                            minfun_args,
                             a,
                             b,
                             fa,
@@ -673,7 +653,7 @@ def solve(  # pylint: disable=too-many-arguments,too-many-locals
             timestep,
             n_substeps,
         ):
-            args = (
+            step_impl_args = (
                 attributes,
                 cell_idx,
                 thd,
@@ -689,7 +669,9 @@ def solve(  # pylint: disable=too-many-arguments,too-many-locals
             )
             success = True
             if adaptive:
-                n_substeps, success = adapt_substeps(args, n_substeps, thd, rtols.thd)
+                n_substeps, success = adapt_substeps(
+                    step_impl_args, n_substeps, thd, rtols.thd
+                )
             if success:
                 (
                     water_vapour_mixing_ratio,
@@ -699,7 +681,7 @@ def solve(  # pylint: disable=too-many-arguments,too-many-locals
                     n_ripening,
                     RH_max,
                     success,
-                ) = step(args, timestep, n_substeps)
+                ) = step(step_impl_args, timestep, n_substeps)
             else:
                 n_activating, n_deactivating, n_ripening, RH_max = -1, -1, -1, -1
             return (

PySDM/backends/impl_numba/methods/deposition_methods.py

@@ -83,19 +83,24 @@ def body(  # pylint: disable=too-many-arguments
                     saturation_ratio_ice = (
                         current_relative_humidity[cid] / current_water_activity[cid]
                     )
-
                     if saturation_ratio_ice == 1:
                         continue
-
+                    Fk = formulae.drop_growth__Fk(
+                        T=temperature,
+                        K=thermal_conductivity * heat_ventilation_factor,
+                        lv=latent_heat_sub,
+                    )
+                    Fd = formulae.drop_growth__Fd(
+                        T=temperature,
+                        D=diffusion_coefficient * mass_ventilation_factor,
+                        pvs=pvs_ice,
+                    )
                     howell_factor_x_diffcoef_x_rhovsice_x_icess = (
                         formulae.drop_growth__r_dr_dt(
                             RH_eq=1,
-                            T=temperature,
                             RH=saturation_ratio_ice,
-                            lv=latent_heat_sub,
-                            pvs=pvs_ice,
-                            D=mass_ventilation_factor * diffusion_coefficient,
-                            K=heat_ventilation_factor * thermal_conductivity,
+                            Fk=Fk,
+                            Fd=Fd,
                         )
                         * formulae.constants.rho_w
                     )

PySDM/backends/impl_numba/test_helpers/scipy_ode_condensation_solver.py

@@ -140,16 +140,23 @@ def _impl(  # pylint: disable=too-many-arguments,too-many-locals
             )
             heat_ventilation_factor = mass_ventilation_factor  # TODO #1588
             sgm = jit_formulae.surface_tension__sigma(T, v, dry_volume[i], f_org[i])
+            Fk = jit_formulae.drop_growth__Fk(
+                T=T,
+                lv=lv,
+                K=heat_ventilation_factor * Kr,
+            )
+            Fd = jit_formulae.drop_growth__Fd(
+                T=T,
+                pvs=pvs,
+                D=mass_ventilation_factor * Dr,
+            )
             r_dr_dt = jit_formulae.drop_growth__r_dr_dt(
-                jit_formulae.hygroscopicity__RH_eq(
+                RH_eq=jit_formulae.hygroscopicity__RH_eq(
                     r, T, kappa[i], dry_volume[i] / PI_4_3, sgm
                 ),
-                T,
-                RH,
-                lv,
-                pvs,
-                mass_ventilation_factor * Dr,
-                heat_ventilation_factor * Kr,
+                RH=RH,
+                Fk=Fk,
+                Fd=Fd,
             )
             dm_dt = jit_formulae.particle_shape_and_density__dm_dt(r, r_dr_dt)
             dy_dt[idx_x + i] = jit_formulae.diffusion_coordinate__dx_dt(m, dm_dt)