From 834ea1fcb2877bb503c37deea617c778d3d72f8f Mon Sep 17 00:00:00 2001 From: Sylwester Arabas Date: Mon, 10 May 2021 14:23:23 +0200 Subject: [PATCH] fix power operators in trivia; introduce some scaffolding for GPU condensation; move explicit_euler and critical volume to backend so that parcel stuff works on ThrustRTC backend; cleanups --- PySDM/attributes/physics/critical_volume.py | 20 +++---- PySDM/backends/numba/conf.py | 2 +- PySDM/backends/numba/impl/_physics_methods.py | 28 ++++++++++ .../numba/impl/condensation_methods.py | 4 +- .../thrustRTC/impl/_algorithmic_methods.py | 29 ++++++++--- .../thrustRTC/impl/_physics_methods.py | 36 +++++++++++-- .../thrustRTC/impl/condensation_methods.py | 26 ++++++++++ .../thrustRTC/test_helpers/cpp2python.py | 1 + PySDM/backends/thrustRTC/thrustRTC.py | 2 + PySDM/core.py | 52 +++++++++---------- PySDM/environments/parcel.py | 6 +-- PySDM/physics/trivia.py | 12 ++--- 12 files changed, 154 insertions(+), 64 deletions(-) create mode 100644 PySDM/backends/thrustRTC/impl/condensation_methods.py diff --git a/PySDM/attributes/physics/critical_volume.py b/PySDM/attributes/physics/critical_volume.py index 2469a6a81..9cb913be3 100644 --- a/PySDM/attributes/physics/critical_volume.py +++ b/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() + ) diff --git a/PySDM/backends/numba/conf.py b/PySDM/backends/numba/conf.py index a38ddec3a..6c631edd5 100644 --- a/PySDM/backends/numba/conf.py +++ b/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 diff --git a/PySDM/backends/numba/impl/_physics_methods.py b/PySDM/backends/numba/impl/_physics_methods.py index c6951cc88..785f120d7 100644 --- a/PySDM/backends/numba/impl/_physics_methods.py +++ b/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) \ No newline at end of file diff --git a/PySDM/backends/numba/impl/condensation_methods.py b/PySDM/backends/numba/impl/condensation_methods.py index 68c40654e..ce3b96942 100644 --- a/PySDM/backends/numba/impl/condensation_methods.py +++ b/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, diff --git a/PySDM/backends/thrustRTC/impl/_algorithmic_methods.py b/PySDM/backends/thrustRTC/impl/_algorithmic_methods.py index a18c2a76b..305464ee2 100644 --- a/PySDM/backends/thrustRTC/impl/_algorithmic_methods.py +++ b/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 diff --git a/PySDM/backends/thrustRTC/impl/_physics_methods.py b/PySDM/backends/thrustRTC/impl/_physics_methods.py index 2b1f6b62a..2d815fbda 100644 --- a/PySDM/backends/thrustRTC/impl/_physics_methods.py +++ b/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)) diff --git a/PySDM/backends/thrustRTC/impl/condensation_methods.py b/PySDM/backends/thrustRTC/impl/condensation_methods.py new file mode 100644 index 000000000..86103b9a9 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/PySDM/backends/thrustRTC/test_helpers/cpp2python.py b/PySDM/backends/thrustRTC/test_helpers/cpp2python.py index 9eea00cd7..148816aa0 100644 --- a/PySDM/backends/thrustRTC/test_helpers/cpp2python.py +++ b/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", } diff --git a/PySDM/backends/thrustRTC/thrustRTC.py b/PySDM/backends/thrustRTC/thrustRTC.py index a6ac0e86b..4119c22ef 100644 --- a/PySDM/backends/thrustRTC/thrustRTC.py +++ b/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 diff --git a/PySDM/core.py b/PySDM/core.py index 40cf98137..ddf306ee6 100644 --- a/PySDM/core.py +++ b/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): diff --git a/PySDM/environments/parcel.py b/PySDM/environments/parcel.py index 1dd9d6d62..002700a0f 100644 --- a/PySDM/environments/parcel.py +++ b/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, diff --git a/PySDM/physics/trivia.py b/PySDM/physics/trivia.py index 27e4779af..ac64910d4 100644 --- a/PySDM/physics/trivia.py +++ b/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)