Skip to content

[NDSL] Restore UW orchestration + interface #1168

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
FlorianDeconinck merged 2 commits into from
Sep 25, 2025
Merged

[NDSL] Restore UW orchestration + interface #1168

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
52e297a
Restore UW interface
FlorianDeconinck Sep 25, 2025
5db8cc9
Restore orchestration
FlorianDeconinck Sep 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 5 additions & 3 deletions ...gcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_uwshcu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
333 changes: 331 additions & 2 deletions ...idComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/interface/python_bridge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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,
Expand Down
Loading