fix power operators in trivia; introduce some scaffolding for GPU con…

…densation; move explicit_euler and critical volume to backend so that parcel stuff works on ThrustRTC backend; cleanups

PySDM/attributes/physics/critical_volume.py

@@ -18,16 +18,10 @@ def __init__(self, builder):
         super().__init__(builder, name='critical volume', dependencies=dependencies)
 
     def recalculate(self):
-        kappa = self.particles.dynamics['Condensation'].kappa
-        v_dry = self.v_dry.get().data
-        v_wet = self.v_wet.get().data
-        T = self.environment['T'].data
-        cell = self.cell_id.get().data
-        for i in range(len(self.data)):  # TODO #347 move to backend
-            sigma = self.formulae.surface_tension.sigma(T[cell[i]], v_wet[i], v_dry[i])
-            self.data.data[i] = self.formulae.trivia.volume(self.formulae.hygroscopicity.r_cr(
-                kp=kappa,
-                rd3=v_dry[i] / const.pi_4_3,
-                T=T[cell[i]],
-                sgm=sigma
-            ))
+        self.core.bck.critical_volume(self.data,
+            kappa=self.particles.dynamics['Condensation'].kappa,
+            v_dry=self.v_dry.get(),
+            v_wet=self.v_wet.get(),
+            T=self.environment['T'],
+            cell=self.cell_id.get()
+        )

PySDM/backends/numba/conf.py

@@ -18,5 +18,5 @@
 except numba.core.errors.UnsupportedParforsError:
     if 'CI' not in os.environ:
         warnings.warn("Numba version used does not support parallel for (32 bits?)")
-    JIT_FLAGS['parallel']=False
+    JIT_FLAGS['parallel'] = False
 

PySDM/backends/numba/impl/_physics_methods.py

@@ -14,6 +14,28 @@ def __init__(self):
         phys_T = self.formulae.state_variable_triplet.T
         phys_p = self.formulae.state_variable_triplet.p
         phys_pv = self.formulae.state_variable_triplet.pv
+        explicit_euler = self.formulae.trivia.explicit_euler
+        phys_sigma = self.formulae.surface_tension.sigma
+        phys_volume = self.formulae.trivia.volume
+        phys_r_cr = self.formulae.hygroscopicity.r_cr
+
+        @numba.njit(**{**conf.JIT_FLAGS, 'fastmath': self.formulae.fastmath})
+        def explicit_euler_body(y, dt, dy_dt):
+            for i in prange(y.shape[0]):
+                y[i] = explicit_euler(y[i], dt, dy_dt)
+        self.explicit_euler_body = explicit_euler_body
+
+        @numba.njit(**{**conf.JIT_FLAGS, 'fastmath': self.formulae.fastmath})
+        def critical_volume(v_cr, kappa, v_dry, v_wet, T, cell):
+            for i in prange(len(v_cr)):
+                sigma = phys_sigma(T[cell[i]], v_wet[i], v_dry[i])
+                v_cr[i] = phys_volume(phys_r_cr(
+                    kp=kappa,
+                    rd3=v_dry[i] / const.pi_4_3,
+                    T=T[cell[i]],
+                    sgm=sigma
+                ))
+        self.critical_volume_body = critical_volume
 
         @numba.njit(**{**conf.JIT_FLAGS, 'fastmath': self.formulae.fastmath})
         def temperature_pressure_RH_body(rhod, thd, qv, T, p, RH):
@@ -39,3 +61,9 @@ def temperature_pressure_RH(self, rhod, thd, qv, T, p, RH):
 
     def terminal_velocity(self, values, radius, k1, k2, k3, r1, r2):
         self.terminal_velocity_body(values, radius, k1, k2, k3, r1, r2)
+
+    def explicit_euler(self, y, dt, dy_dt):
+        self.explicit_euler_body(y.data, dt, dy_dt)
+
+    def critical_volume(self, v_cr, kappa, v_dry, v_wet, T, cell):
+        self.critical_volume_body(v_cr.data, kappa, v_dry.data, v_wet.data, T.data, cell.data)

PySDM/backends/numba/impl/condensation_methods.py

@@ -217,10 +217,10 @@ def calculate_ml_new(dt, fake, T, p, RH, v, v_cr, n, vdry, cell_idx, kappa, lv,
 
         return calculate_ml_new
 
-    def make_condensation_solver(self, dt, dt_range, adaptive=True):
+    def make_condensation_solver(self, dt, dt_range, adaptive):
         return CondensationMethods.make_condensation_solver_impl(
             fastmath=self.formulae.fastmath,
-            phys_pvs_C = self.formulae.saturation_vapour_pressure.pvs_Celsius,
+            phys_pvs_C=self.formulae.saturation_vapour_pressure.pvs_Celsius,
             phys_lv=self.formulae.latent_heat.lv,
             phys_r_dr_dt=self.formulae.drop_growth.r_dr_dt,
             phys_RH_eq=self.formulae.hygroscopicity.RH_eq,

PySDM/backends/thrustRTC/impl/_algorithmic_methods.py

@@ -186,12 +186,24 @@ def compute_gamma(gamma, rand, n, cell_id,
     @staticmethod
     @nice_thrust(**NICE_THRUST_FLAGS)
     def condensation(
-            solver,
-            n_cell, cell_start_arg,
-            v, particle_temperatures, n, vdry, idx, rhod, thd, qv, dv, prhod, pthd, pqv, kappa,
-            rtol_x, rtol_thd, dt, substeps, cell_order
+        solver,
+        n_cell, cell_start_arg,
+        v, v_cr, n, vdry, idx, rhod, thd, qv, dv, prhod, pthd, pqv, kappa,
+        rtol_x, rtol_thd, dt, counters, cell_order, RH_max, success
     ):
-        raise NotImplementedError()
+        if n_cell != 1:
+            raise NotImplementedError()
+        cell_id = 0
+        cell_idx = None
+
+        n_substeps = counters['n_substeps']
+        dv_mean = dv
+        dthd_dt = (pthd[cell_id] - thd[cell_id]) / dt
+        dqv_dt = (pqv[cell_id] - qv[cell_id]) / dt
+        rhod_mean = (prhod[cell_id] + rhod[cell_id]) / 2
+        m_d = rhod_mean * dv_mean
+
+        solver(v, v_cr, n, vdry, cell_idx, kappa, thd, qv, dthd_dt, dqv_dt, m_d, rhod_mean, rtol_x, rtol_thd, dt, n_substeps)
 
     __flag_precipitated_body = trtc.For(['idx', 'n_sd', 'n_dims', 'healthy', 'cell_origin', 'position_in_cell',
                                          'volume', 'n'], "i", '''
@@ -407,7 +419,8 @@ def _sort_by_cell_id_and_update_cell_start(cell_id, cell_idx, cell_start, idx):
         # TODO #330
         n_sd = cell_id.shape[0]
         trtc.Fill(cell_start.data, trtc.DVInt64(n_sd))
-        AlgorithmicMethods.___sort_by_cell_id_and_update_cell_start_body.launch_n(len(idx) - 1,
-                                                                                  [cell_id.data, cell_start.data,
-                                                                                   idx.data])
+        if len(idx) > 1:
+            AlgorithmicMethods.___sort_by_cell_id_and_update_cell_start_body.launch_n(len(idx) - 1,
+                                                                                      [cell_id.data, cell_start.data,
+                                                                                       idx.data])
         return idx

PySDM/backends/thrustRTC/impl/_physics_methods.py

@@ -11,12 +11,36 @@
 
 class PhysicsMethods:
     def __init__(self):
+        phys = self.formulae
+
         self._temperature_pressure_RH_body = trtc.For(["rhod", "thd", "qv", "T", "p", "RH"], "i", f'''
-            T[i] = {c_inline(self.formulae.state_variable_triplet.T, rhod="rhod[i]", thd="thd[i]")};
-            p[i] = {c_inline(self.formulae.state_variable_triplet.p, rhod="rhod[i]", T="T[i]", qv="qv[i]")};
-            RH[i] = {c_inline(self.formulae.state_variable_triplet.pv, p="p[i]", qv="qv[i]")} / {c_inline(self.formulae.saturation_vapour_pressure.pvs_Celsius, T="T[i] - const.T0")};
+            T[i] = {c_inline(phys.state_variable_triplet.T, rhod="rhod[i]", thd="thd[i]")};
+            p[i] = {c_inline(phys.state_variable_triplet.p, rhod="rhod[i]", T="T[i]", qv="qv[i]")};
+            RH[i] = {c_inline(phys.state_variable_triplet.pv, p="p[i]", qv="qv[i]")} / {c_inline(phys.saturation_vapour_pressure.pvs_Celsius, T="T[i] - const.T0")};
+        ''')
+
+        self.__explicit_euler_body = trtc.For(("y", "dt", "dy_dt"), "i", f'''
+            y[i] = {c_inline(phys.trivia.explicit_euler, y="y[i]", dt="dt", dy_dt="dy_dt")};
+        ''')
+
+        self.__critical_volume_body = trtc.For(("v_cr", "kappa", "v_dry", "v_wet", "T", "cell"), "i", f'''
+            auto sigma = {c_inline(phys.surface_tension.sigma, T="T[cell[i]]", v_wet="v_wet[i]", v_dry="v_dry[i]")};
+            auto r_cr = {c_inline(phys.hygroscopicity.r_cr,
+                kp="kappa",
+                rd3="v_dry[i] / const.pi_4_3",
+                T="T[cell[i]]",
+                sgm="sigma"
+            )};
+            v_cr[i] = {c_inline(phys.trivia.volume, radius="r_cr")};
         ''')
 
+    @nice_thrust(**NICE_THRUST_FLAGS)
+    def critical_volume(self, v_cr, kappa, v_dry, v_wet, T, cell):
+        kappa = PrecisionResolver.get_floating_point(kappa)
+        self.__critical_volume_body.launch_n(
+            v_cr.shape[0], (v_cr.data, kappa, v_dry.data, v_wet.data, T.data, cell.data)
+        )
+
     @nice_thrust(**NICE_THRUST_FLAGS)
     def temperature_pressure_RH(self, rhod, thd, qv, T, p, RH):
         self._temperature_pressure_RH_body.launch_n(
@@ -45,3 +69,9 @@ def terminal_velocity(values, radius, k1, k2, k3, r1, r2):
         r1 = PrecisionResolver.get_floating_point(r1)
         r2 = PrecisionResolver.get_floating_point(r2)
         PhysicsMethods.__terminal_velocity_body.launch_n(values.size(), [values, radius, k1, k2, k3, r1, r2])
+
+    @nice_thrust(**NICE_THRUST_FLAGS)
+    def explicit_euler(self, y, dt, dy_dt):
+        dt = PrecisionResolver.get_floating_point(dt)
+        dy_dt = PrecisionResolver.get_floating_point(dy_dt)
+        self.__explicit_euler_body.launch_n(y.shape[0], (y.data, dt, dy_dt))

PySDM/backends/thrustRTC/impl/condensation_methods.py

@@ -0,0 +1,26 @@
+from PySDM.physics import constants as const
+
+
+class CondensationMethods:
+    def make_condensation_solver(self, dt, dt_range, adaptive):
+
+        def calculate_ml_old(v, n, cell_idx):
+            result = 0
+            for drop in cell_idx:
+                result += n[drop] * v[drop] * const.rho_w
+            return result
+
+        def step(args, dt, n_substep):
+            step_iml(*args, dt, n_substep, fake=False)
+
+        def step_iml(v, v_cr, n, vdry, cell_idx, kappa, thd, qv, dthd_dt, dqv_dt, m_d, rhod_mean,
+                  rtol_x, rtol_thd, dt, n_substep, fake):
+            ml_old = calculate_ml_old(v, n, cell_idx)
+
+        def solve(v, v_cr, n, vdry, cell_idx, kappa, thd, qv, dthd_dt, dqv_dt, m_d, rhod_mean,
+                  rtol_x, rtol_thd, dt, n_substeps):
+            args = (v, v_cr, n, vdry, cell_idx, kappa, thd, qv, dthd_dt, dqv_dt, m_d, rhod_mean, rtol_x)
+
+            step(args, dt, n_substeps)
+
+        return solve

PySDM/backends/thrustRTC/test_helpers/cpp2python.py

@@ -23,6 +23,7 @@
     "ceil": "np.ceil",
     "exp(": "np.exp(",
     "power(": "np.power(",
+    "sqrt(": "np.sqrt(",
     "return": "continue",
 }
 

PySDM/backends/thrustRTC/thrustRTC.py

@@ -8,6 +8,7 @@
 from PySDM.backends.thrustRTC.impl._pair_methods import PairMethods
 from PySDM.backends.thrustRTC.impl._index_methods import IndexMethods
 from PySDM.backends.thrustRTC.impl._physics_methods import PhysicsMethods
+from PySDM.backends.thrustRTC.impl.condensation_methods import CondensationMethods
 from PySDM.backends.thrustRTC.storage import Storage as ImportedStorage
 from PySDM.backends.thrustRTC.random import Random as ImportedRandom
 
@@ -17,6 +18,7 @@ class ThrustRTC(
     PairMethods,
     IndexMethods,
     PhysicsMethods,
+    CondensationMethods
 ):
     ENABLE = True
     Storage = ImportedStorage

PySDM/core.py

@@ -87,35 +87,31 @@ def update_TpRH(self):
         # TODO #443: mark_updated
 
     def condensation(self, kappa, rtol_x, rtol_thd, counters, RH_max, success, cell_order):
-        particle_heat = \
-            self.particles["heat"] if self.particles.has_attribute("heat") else \
-            self.null
-
         self.backend.condensation(
-                solver=self.condensation_solver,
-                n_cell=self.mesh.n_cell,
-                cell_start_arg=self.particles.cell_start,
-                v=self.particles["volume"],
-                n=self.particles['n'],
-                vdry=self.particles["dry volume"],
-                idx=self.particles._Particles__idx,
-                rhod=self.env["rhod"],
-                thd=self.env["thd"],
-                qv=self.env["qv"],
-                dv=self.env.dv,
-                prhod=self.env.get_predicted("rhod"),
-                pthd=self.env.get_predicted("thd"),
-                pqv=self.env.get_predicted("qv"),
-                kappa=kappa,
-                rtol_x=rtol_x,
-                rtol_thd=rtol_thd,
-                v_cr=self.particles["critical volume"],
-                dt=self.dt,
-                counters=counters,
-                cell_order=cell_order,
-                RH_max=RH_max,
-                success=success
-            )
+            solver=self.condensation_solver,
+            n_cell=self.mesh.n_cell,
+            cell_start_arg=self.particles.cell_start,
+            v=self.particles["volume"],
+            n=self.particles['n'],
+            vdry=self.particles["dry volume"],
+            idx=self.particles._Particles__idx,
+            rhod=self.env["rhod"],
+            thd=self.env["thd"],
+            qv=self.env["qv"],
+            dv=self.env.dv,
+            prhod=self.env.get_predicted("rhod"),
+            pthd=self.env.get_predicted("thd"),
+            pqv=self.env.get_predicted("qv"),
+            kappa=kappa,
+            rtol_x=rtol_x,
+            rtol_thd=rtol_thd,
+            v_cr=self.particles["critical volume"],
+            dt=self.dt,
+            counters=counters,
+            cell_order=cell_order,
+            RH_max=RH_max,
+            success=success
+        )
 
     def run(self, steps):
         for _ in range(steps):

PySDM/environments/parcel.py

@@ -85,9 +85,9 @@ def advance_parcel_vars(self):
         drho_dz = self.formulae.hydrostatics.drho_dz(self.g, p, T, qv, lv, dql_dz=dql_dz)
         drhod_dz = drho_dz
 
-        self.core.formulae.trivia.explicit_euler(self._tmp['t'].data, dt, 1)
-        self.core.formulae.trivia.explicit_euler(self._tmp['z'].data, dt, dz_dt)
-        self.core.formulae.trivia.explicit_euler(self._tmp['rhod'].data, dt, dz_dt * drhod_dz)
+        self.core.bck.explicit_euler(self._tmp['t'], dt, 1)
+        self.core.bck.explicit_euler(self._tmp['z'], dt, dz_dt)
+        self.core.bck.explicit_euler(self._tmp['rhod'], dt, dz_dt * drhod_dz)
 
         self.mesh.dv = self.formulae.trivia.volume_of_density_mass(
             (self._tmp['rhod'][0] + self["rhod"][0]) / 2,

PySDM/physics/trivia.py

@@ -1,5 +1,5 @@
 from PySDM.physics import constants as const
-from numpy import abs, log10, exp
+from numpy import abs, log10, exp, power
 
 
 class Trivia:
@@ -9,15 +9,15 @@ def volume_of_density_mass(rho, m):
 
     @staticmethod
     def radius(volume):
-        return (volume / const.pi_4_3)**(1/3)
+        return power(volume / const.pi_4_3, 1./3)
 
     @staticmethod
     def volume(radius):
-        return const.pi_4_3 * radius**3
+        return const.pi_4_3 * power(radius, 3)
 
     @staticmethod
     def explicit_euler(y, dt, dy_dt):
-        y += dt * dy_dt
+        return y + dt * dy_dt
 
     @staticmethod
     def within_tolerance(error_estimate, value, rtol):
@@ -29,7 +29,7 @@ def H2pH(H):
 
     @staticmethod
     def pH2H(pH):
-        return 10**(-pH) * 1e3
+        return power(10, -pH) * 1e3
 
     @staticmethod
     def vant_hoff(K, dH, T, *, T_0):
@@ -57,4 +57,4 @@ def p_d(p, qv):
 
     @staticmethod
     def th_std(p, T):
-        return T * (const.p1000 / p)**(const.Rd_over_c_pd)
+        return T * power(const.p1000 / p, const.Rd_over_c_pd)