diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_uwshcu.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_uwshcu.py index be9b6fe22..7394ba9d2 100644 --- a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_uwshcu.py +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_uwshcu.py @@ -20,7 +20,7 @@ ) import pyMoist.constants as constants -from ndsl import QuantityFactory, StencilFactory +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM from ndsl.dsl.typing import BoolFieldIJ, Float, FloatField, FloatFieldIJ, Int, IntField, IntFieldIJ from pyMoist.field_types import FloatField_NTracers, FloatFieldIJ_NTracers @@ -7281,7 +7281,10 @@ def __init__( UW_config: UWConfiguration, formulation: SaturationFormulation = SaturationFormulation.Staars, ) -> None: - # orchestrate(obj=self, config=stencil_factory.config.dace_config) + orchestrate( + obj=self, + config=stencil_factory.config.dace_config, + ) # Initialize the ComputeUwshcu class @@ -7637,7 +7640,6 @@ def __call__( qisub_inv: FloatField, tpert_out: FloatFieldIJ, qpert_out: FloatFieldIJ, - formulation: SaturationFormulation = SaturationFormulation.Staars, ): """ University of Washington Shallow Convection Scheme diff --git a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/interface/python_bridge.py b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/interface/python_bridge.py index 4e0409510..bbe1bdb17 100644 --- a/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/interface/python_bridge.py +++ b/GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/interface/python_bridge.py @@ -8,7 +8,7 @@ from mpi4py import MPI from ndsl.dsl.gt4py_utils import is_gpu_backend -from ndsl.dsl.typing import Float +from ndsl.dsl.typing import Float, Int from ndsl.optional_imports import cupy as cp from pyMoist.interface.cuda_profiler import CUDAProfiler, TimedCUDAProfiler from pyMoist.interface.f_py_conversion import FortranPythonConversion @@ -173,6 +173,335 @@ def aer_activation( self.f_py.python_to_fortran(nwfa, f_nwfa) self.f_py.python_to_fortran(nactl, f_nactl) + def UW_shallow_convection( + self, + dotransport: int, + k0: int, + windsrcavg: int, + qtsrchgt: float, + qtsrc_fac: float, + thlsrc_fac: float, + frc_rasn: float, + rbuoy: float, + epsvarw: float, + use_CINcin: int, + mumin1: float, + rmaxfrac: float, + PGFc: float, + dt: float, + niter_xc: int, + criqc: float, + rle: float, + cridist_opt: int, + mixscale: float, + rdrag: float, + rkm: float, + use_self_detrain: int, + detrhgt: float, + use_cumpenent: int, + rpen: float, + use_momenflx: int, + rdrop: float, + iter_cin: int, + f_pifc0_inv: cffi.FFI.CData, + f_zifc0_inv: cffi.FFI.CData, + f_pmid0_inv: cffi.FFI.CData, + f_zmid0_inv: cffi.FFI.CData, + f_kpbl_inv: cffi.FFI.CData, + f_exnmid0_inv: cffi.FFI.CData, + f_exnifc0_inv: cffi.FFI.CData, + f_dp0_inv: cffi.FFI.CData, + f_u0_inv: cffi.FFI.CData, + f_v0_inv: cffi.FFI.CData, + f_qv0_inv: cffi.FFI.CData, + f_ql0_inv: cffi.FFI.CData, + f_qi0_inv: cffi.FFI.CData, + f_t0_inv: cffi.FFI.CData, + f_frland: cffi.FFI.CData, + f_tke_inv: cffi.FFI.CData, + f_rkfre: cffi.FFI.CData, + f_cush: cffi.FFI.CData, + f_shfx: cffi.FFI.CData, + f_evap: cffi.FFI.CData, + f_cnvtr: cffi.FFI.CData, + f_CNV_Tracers: cffi.FFI.CData, + f_umf_inv: cffi.FFI.CData, + f_dcm_inv: cffi.FFI.CData, + f_qtflx_inv: cffi.FFI.CData, + f_slflx_inv: cffi.FFI.CData, + f_uflx_inv: cffi.FFI.CData, + f_vflx_inv: cffi.FFI.CData, + f_qvten_inv: cffi.FFI.CData, + f_qlten_inv: cffi.FFI.CData, + f_qiten_inv: cffi.FFI.CData, + f_tten_inv: cffi.FFI.CData, + f_uten_inv: cffi.FFI.CData, + f_vten_inv: cffi.FFI.CData, + f_qrten_inv: cffi.FFI.CData, + f_qsten_inv: cffi.FFI.CData, + f_cufrc_inv: cffi.FFI.CData, + f_fer_inv: cffi.FFI.CData, + f_fdr_inv: cffi.FFI.CData, + f_ndrop_inv: cffi.FFI.CData, + f_nice_inv: cffi.FFI.CData, + f_qldet_inv: cffi.FFI.CData, + f_qlsub_inv: cffi.FFI.CData, + f_qidet_inv: cffi.FFI.CData, + f_qisub_inv: cffi.FFI.CData, + f_tpert_out: cffi.FFI.CData, + f_qpert_out: cffi.FFI.CData, + ): + CUDAProfiler.start_cuda_profiler() + + with TimedCUDAProfiler("[UW] Fortran -> Python", self._timings): + in_pifc0_inv = self.f_py.fortran_to_python( + f_pifc0_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + in_zifc0_inv = self.f_py.fortran_to_python( + f_zifc0_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + in_pmid0_inv = self.f_py.fortran_to_python(f_pmid0_inv) + in_zmid0_inv = self.f_py.fortran_to_python(f_zmid0_inv) + in_kpbl_inv = self.f_py.fortran_to_python( + f_kpbl_inv, + [ + self.flags.npx, + self.flags.npy, + ], + ) + in_exnmid0_inv = self.f_py.fortran_to_python(f_exnmid0_inv) + in_exnifc0_inv = self.f_py.fortran_to_python( + f_exnifc0_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + in_dp0_inv = self.f_py.fortran_to_python(f_dp0_inv) + in_u0_inv = self.f_py.fortran_to_python(f_u0_inv) + in_v0_inv = self.f_py.fortran_to_python(f_v0_inv) + in_qv0_inv = self.f_py.fortran_to_python(f_qv0_inv) + in_ql0_inv = self.f_py.fortran_to_python(f_ql0_inv) + in_qi0_inv = self.f_py.fortran_to_python(f_qi0_inv) + in_t0_inv = self.f_py.fortran_to_python(f_t0_inv) + in_frland = self.f_py.fortran_to_python( + f_frland, + [ + self.flags.npx, + self.flags.npy, + ], + ) + in_tke_inv = self.f_py.fortran_to_python( + f_tke_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + in_rkfre = self.f_py.fortran_to_python( + f_rkfre, + [ + self.flags.npx, + self.flags.npy, + ], + ) + in_cush = self.f_py.fortran_to_python( + f_cush, + [ + self.flags.npx, + self.flags.npy, + ], + ) + in_shfx = self.f_py.fortran_to_python( + f_shfx, + [ + self.flags.npx, + self.flags.npy, + ], + ) + in_evap = self.f_py.fortran_to_python( + f_evap, + [ + self.flags.npx, + self.flags.npy, + ], + ) + in_cnvtr = self.f_py.fortran_to_python( + f_cnvtr, + [ + self.flags.npx, + self.flags.npy, + ], + ) + in_CNV_Tracers = self.f_py.fortran_to_python( + f_CNV_Tracers, + [ + self.flags.npx, + self.flags.npy, + self.flags.npz, + self.pymoist.UW_config.NCNST, + ], + ) + self.f_py.fortran_to_python(f_CNV_Tracers) + + out_umf_inv = self.f_py.fortran_to_python( + f_umf_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + out_dcm_inv = self.f_py.fortran_to_python(f_dcm_inv) + out_qtflx_inv = self.f_py.fortran_to_python( + f_qtflx_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + out_slflx_inv = self.f_py.fortran_to_python( + f_slflx_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + out_uflx_inv = self.f_py.fortran_to_python( + f_uflx_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + out_vflx_inv = self.f_py.fortran_to_python( + f_vflx_inv, + [self.flags.npx, self.flags.npy, self.flags.npz + 1], + ) + out_qvten_inv = self.f_py.fortran_to_python(f_qvten_inv) + out_qlten_inv = self.f_py.fortran_to_python(f_qlten_inv) + out_qiten_inv = self.f_py.fortran_to_python(f_qiten_inv) + out_tten_inv = self.f_py.fortran_to_python(f_tten_inv) + out_uten_inv = self.f_py.fortran_to_python(f_uten_inv) + out_vten_inv = self.f_py.fortran_to_python(f_vten_inv) + out_qrten_inv = self.f_py.fortran_to_python(f_qrten_inv) + out_qsten_inv = self.f_py.fortran_to_python(f_qsten_inv) + out_cufrc_inv = self.f_py.fortran_to_python(f_cufrc_inv) + out_fer_inv = self.f_py.fortran_to_python(f_fer_inv) + out_fdr_inv = self.f_py.fortran_to_python(f_fdr_inv) + out_ndrop_inv = self.f_py.fortran_to_python(f_ndrop_inv) + out_nice_inv = self.f_py.fortran_to_python(f_nice_inv) + out_qldet_inv = self.f_py.fortran_to_python(f_qldet_inv) + out_qlsub_inv = self.f_py.fortran_to_python(f_qlsub_inv) + out_qidet_inv = self.f_py.fortran_to_python(f_qidet_inv) + out_qisub_inv = self.f_py.fortran_to_python(f_qisub_inv) + out_tpert_out = self.f_py.fortran_to_python( + f_tpert_out, + [ + self.flags.npx, + self.flags.npy, + ], + ) + out_qpert_out = self.f_py.fortran_to_python( + f_qpert_out, + [ + self.flags.npx, + self.flags.npy, + ], + ) + + with TimedCUDAProfiler("[UW] Run", self._timings): + self.pymoist.UW_shallow_convection( + # Inputs + pifc0_inv=in_pifc0_inv, + zifc0_inv=in_zifc0_inv, + pmid0_inv=in_pmid0_inv, + zmid0_inv=in_zmid0_inv, + kpbl_inv=in_kpbl_inv, + exnmid0_inv=in_exnmid0_inv, + exnifc0_inv=in_exnifc0_inv, + dp0_inv=in_dp0_inv, + u0_inv=in_u0_inv, + v0_inv=in_v0_inv, + qv0_inv=in_qv0_inv, + ql0_inv=in_ql0_inv, + qi0_inv=in_qi0_inv, + t0_inv=in_t0_inv, + frland=in_frland, + tke_inv=in_tke_inv, + rkfre=in_rkfre, + cush=in_cush, + shfx=in_shfx, + evap=in_evap, + cnvtr=in_cnvtr, + CNV_Tracers=in_CNV_Tracers, + # Parameters + dotransport=Int(dotransport), + k0=Int(k0), + windsrcavg=Int(windsrcavg), + qtsrchgt=Float(qtsrchgt), + qtsrc_fac=Float(qtsrc_fac), + thlsrc_fac=Float(thlsrc_fac), + frc_rasn=Float(frc_rasn), + rbuoy=Float(rbuoy), + epsvarw=Float(epsvarw), + use_CINcin=Int(use_CINcin), + mumin1=Float(mumin1), + rmaxfrac=Float(rmaxfrac), + PGFc=Float(PGFc), + dt=Float(dt), + niter_xc=Int(niter_xc), + criqc=Float(criqc), + rle=Float(rle), + cridist_opt=Int(cridist_opt), + mixscale=Float(mixscale), + rdrag=Float(rdrag), + rkm=Float(rkm), + use_self_detrain=Int(use_self_detrain), + detrhgt=Float(detrhgt), + use_cumpenent=Int(use_cumpenent), + rpen=Float(rpen), + use_momenflx=Int(use_momenflx), + rdrop=Float(rdrop), + iter_cin=Int(iter_cin), + # Outputs + umf_inv=out_umf_inv, + dcm_inv=out_dcm_inv, + qtflx_inv=out_qtflx_inv, + slflx_inv=out_slflx_inv, + uflx_inv=out_uflx_inv, + vflx_inv=out_vflx_inv, + qvten_inv=out_qvten_inv, + qlten_inv=out_qlten_inv, + qiten_inv=out_qiten_inv, + tten_inv=out_tten_inv, + uten_inv=out_uten_inv, + vten_inv=out_vten_inv, + qrten_inv=out_qrten_inv, + qsten_inv=out_qsten_inv, + cufrc_inv=out_cufrc_inv, + fer_inv=out_fer_inv, + fdr_inv=out_fdr_inv, + ndrop_inv=out_ndrop_inv, + nice_inv=out_nice_inv, + qldet_inv=out_qldet_inv, + qlsub_inv=out_qlsub_inv, + qidet_inv=out_qidet_inv, + qisub_inv=out_qisub_inv, + tpert_out=out_tpert_out, + qpert_out=out_qpert_out, + ) + + with TimedCUDAProfiler("[GFDL_1M] Python -> Fortran", self._timings): + self.f_py.python_to_fortran(out_umf_inv, f_umf_inv) + self.f_py.python_to_fortran(out_dcm_inv, f_dcm_inv) + self.f_py.python_to_fortran(out_qtflx_inv, f_qtflx_inv) + self.f_py.python_to_fortran(out_slflx_inv, f_slflx_inv) + self.f_py.python_to_fortran(out_uflx_inv, f_uflx_inv) + self.f_py.python_to_fortran(out_vflx_inv, f_vflx_inv) + self.f_py.python_to_fortran(out_qvten_inv, f_qvten_inv) + self.f_py.python_to_fortran(out_qlten_inv, f_qlten_inv) + self.f_py.python_to_fortran(out_qiten_inv, f_qiten_inv) + self.f_py.python_to_fortran(out_tten_inv, f_tten_inv) + self.f_py.python_to_fortran(out_uten_inv, f_uten_inv) + self.f_py.python_to_fortran(out_vten_inv, f_vten_inv) + self.f_py.python_to_fortran(out_qrten_inv, f_qrten_inv) + self.f_py.python_to_fortran(out_qsten_inv, f_qsten_inv) + self.f_py.python_to_fortran(out_cufrc_inv, f_cufrc_inv) + self.f_py.python_to_fortran(out_fer_inv, f_fer_inv) + self.f_py.python_to_fortran(out_fdr_inv, f_fdr_inv) + self.f_py.python_to_fortran(out_ndrop_inv, f_ndrop_inv) + self.f_py.python_to_fortran(out_nice_inv, f_nice_inv) + self.f_py.python_to_fortran(out_qldet_inv, f_qldet_inv) + self.f_py.python_to_fortran(out_qlsub_inv, f_qlsub_inv) + self.f_py.python_to_fortran(out_qidet_inv, f_qidet_inv) + self.f_py.python_to_fortran(out_qisub_inv, f_qisub_inv) + self.f_py.python_to_fortran(out_tpert_out, f_tpert_out) + self.f_py.python_to_fortran(out_qpert_out, f_qpert_out) + WRAPPER = PYMOIST_WRAPPER() @@ -456,7 +785,7 @@ def compute_uwshcu_run( ): if not WRAPPER.ready: raise RuntimeError("[GEOS WRAPPER] Bad init, did you call init?") - WRAPPER.pymoist.UW_shallow_convection( + WRAPPER.UW_shallow_convection( dotransport=dotransport, k0=k0, windsrcavg=windsrcavg,