From 75c46aedaae73a109d34697ec4cab943aba53dcb Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 11 Sep 2024 08:04:55 -0700 Subject: [PATCH 001/252] Added GEOS modifications in moist_cv where GEOS applies the 'max' function to tracer values. Also, r_vir is casted to Float so that it can run with environment variable PACE_CONSTANTS=32 --- pyFV3/stencils/moist_cv.py | 10 ++++++---- .../savepoint/translate/translate_moistcvpluspt_2d.py | 4 ++-- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index c571a24e..d22f0a26 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -20,9 +20,10 @@ def set_cappa(qvapor, cvm, r_vir): @gtscript.function def moist_cvm(qvapor, gz, ql, qs): + # CK : GEOS applies the "max" function to tracer values cvm = ( - (1.0 - (qvapor + gz)) * constants.CV_AIR - + qvapor * constants.CV_VAP + (1.0 - (max(qvapor,0.0) + gz)) * constants.CV_AIR + + max(qvapor,0.0) * constants.CV_VAP + ql * constants.C_LIQ + qs * constants.C_ICE ) @@ -38,8 +39,9 @@ def moist_cv_nwat6_fn( qice: FloatField, qgraupel: FloatField, ): - ql = qliquid + qrain - qs = qice + qsnow + qgraupel + # CK : GEOS applies the "max" function to tracer values + ql = max(qliquid, 0.0) + max(qrain, 0.0) + qs = max(qice,0.0) + max(qsnow,0.0) + max(qgraupel,0.0) gz = ql + qs cvm = moist_cvm(qvapor, gz, ql, qs) return cvm, gz diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d.py b/tests/savepoint/translate/translate_moistcvpluspt_2d.py index c50fda2e..6e526ea3 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d.py @@ -1,7 +1,7 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField +from ndsl.dsl.typing import FloatField, Float from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j from pyFV3.stencils import moist_cv @@ -18,7 +18,7 @@ def moist_pt( cappa: FloatField, delp: FloatField, delz: FloatField, - r_vir: float, + r_vir: Float, ): with computation(PARALLEL), interval(...): cvm, gz, q_con, cappa, pt = moist_cv.moist_pt_func( From 73c4ffadcc891bafe0286c56b84020e7ae7641b6 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 20 Sep 2024 14:22:33 -0700 Subject: [PATCH 002/252] Added translate test for fixing excessive W, which is implemented in GEOS LagrangianToEulerian but not in GFDL's version --- pyFV3/stencils/remapping.py | 36 +++- .../translate/translate_pn2_pk_delp.py | 128 ++++++++++++ .../translate_w_fix_consrv_moment.py | 190 ++++++++++++++++++ 3 files changed, 345 insertions(+), 9 deletions(-) create mode 100644 tests/savepoint/translate/translate_pn2_pk_delp.py create mode 100644 tests/savepoint/translate/translate_w_fix_consrv_moment.py diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index a14c38a1..901fccc3 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -100,6 +100,7 @@ def moist_cv_pt_pressure( ps: FloatFieldIJ, pn2: FloatField, peln: FloatField, + remap_t: bool, r_vir: Float, ): """ @@ -126,11 +127,12 @@ def moist_cv_pt_pressure( pn2 (out): peln (in): """ - from __externals__ import hydrostatic, kord_tm + from __externals__ import hydrostatic#, kord_tm # moist_cv.moist_pt with computation(PARALLEL), interval(0, -1): - if __INLINED(kord_tm < 0): + # if __INLINED(kord_tm < 0): + if(remap_t): cvm, gz, q_con, cappa, pt = moist_pt_func( qvapor, qliquid, @@ -145,9 +147,10 @@ def moist_cv_pt_pressure( delz, r_vir, ) - # delz_adjust - if __INLINED(not hydrostatic): - delz = -delz / delp + # NOTE : GEOS does not perform the delz computation at this location + # # delz_adjust + # if __INLINED(not hydrostatic): + # delz = -delz / delp # pressure_updates with computation(FORWARD): with interval(-1, None): @@ -164,9 +167,11 @@ def moist_cv_pt_pressure( pn2 = peln with computation(BACKWARD), interval(0, -1): dp2 = pe2[0, 0, 1] - pe2 + + # NOTE : GEOS doesn't perform the delp calcuation at this location # copy_stencil - with computation(PARALLEL), interval(0, -1): - delp = dp2 + # with computation(PARALLEL), interval(0, -1): + # delp = dp2 def pn2_pk_delp( @@ -186,7 +191,9 @@ def pn2_pk_delp( pk (out): """ with computation(PARALLEL), interval(...): - delp = dp2 + # NOTE : GEOS doesn't perform the delp calcuation at this location + # Also, in moist_cv_pt_pressure, the below calculation is also done + # delp = dp2 pn2 = log(pe2) pk = exp(akap * pn2) @@ -375,6 +382,13 @@ def __init__( self._do_sat_adjust = config.do_sat_adj + self._remap_t = False + + # NOTE: In GEOS, remap_t is set to True in general + # Add in the "remap_option" check later + if(True): + self._remap_t = True + self.kmp = grid_indexing.domain[2] - 1 for k in range(pfull.shape[0]): if pfull.view[k] > 10.0e2: @@ -387,7 +401,8 @@ def __init__( self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( moist_cv_pt_pressure, - externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, + # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, + externals={"hydrostatic": hydrostatic}, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) @@ -612,6 +627,7 @@ def __call__( ps, self._pn2, peln, + self._remap_t, zvir, ) @@ -625,6 +641,8 @@ def __call__( self._map_single_w(w, self._pe1, self._pe2, qs=wsd) self._map_single_delz(delz, self._pe1, self._pe2) + # W_limiter routine will go here + self._undo_delz_adjust_and_copy_peln(delp, delz, peln, self._pe0, self._pn2) # if do_omega: # NOTE untested # pe3 = copy(omga, origin=(grid_indexing.isc, grid_indexing.jsc, 1)) diff --git a/tests/savepoint/translate/translate_pn2_pk_delp.py b/tests/savepoint/translate/translate_pn2_pk_delp.py new file mode 100644 index 00000000..543c83ee --- /dev/null +++ b/tests/savepoint/translate/translate_pn2_pk_delp.py @@ -0,0 +1,128 @@ +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + +from ndsl import StencilFactory +from ndsl.dsl.typing import FloatField, Float +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.remapping import pn2_pk_delp + + + +class testClass: + """ + Class to test with DaCe orchestration. test class is MoistCVPlusPt_2d + """ + + def __init__( + self, + stencil_factory: StencilFactory, + grid, + ): + self._pn2_pk_delp = stencil_factory.from_origin_domain( + func=pn2_pk_delp, + origin=(3,3,1), + domain=(24,24,71), + ) + + def __call__( + self, + dp2: FloatField, + delp: FloatField, + pe2: FloatField, + pn2: FloatField, + pk: FloatField, + akap: Float, + ): + self._pn2_pk_delp( + dp2, + delp, + pe2, + pn2, + pk, + akap + ) + + +class TranslatePN2_PK_DelP(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.compute_func = testClass(self.stencil_factory, self.grid) # type: ignore + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "pe2_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pn2_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1,}, + "pk_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1,}, + + } + self.in_vars["parameters"] = [ + "akap", + ] + + self.out_vars = { + "pe2_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1,}, + "pn2_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1,}, + "pk_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1,}, + } + self._dp2 = self.quantity_factory._numpy.zeros( + ( + 31, + 31, + 73, + ), dtype=Float, + ) + + self._delp = self.quantity_factory._numpy.zeros( + ( + 31, + 31, + 73, + ), dtype=Float, + ) + + + def compute_from_storage(self, inputs): + + # print("delp shape = ", self._delp.shape) + # print("inputs[pe2] shape = ", inputs["pe2_"].shape) + # print("inputs[pe2_][:,3,0] = ", inputs["pe2_"][:,3,0]) + # print('self.grid.is_ = ', self.grid.is_) + # print('self.grid.ie = ', self.grid.ie) + # print('self.grid.js = ', self.grid.js) + # print('self.grid.je = ', self.grid.je) + # print('self.storage_vars() = ', self.storage_vars()) + # self.make_storage_data_input_vars(inputs) + # exit(1) + self.compute_func(self._dp2, + self._delp, + inputs["pe2_"], + inputs["pn2_"], + inputs["pk_"], + inputs["akap"]) + return inputs diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py new file mode 100644 index 00000000..97194e49 --- /dev/null +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -0,0 +1,190 @@ +from gt4py.cartesian.gtscript import PARALLEL, FORWARD, BACKWARD, computation, interval + +from ndsl import StencilFactory, QuantityFactory +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Float, BoolFieldIJ +from ndsl.stencils.testing import TranslateFortranData2Py + +def W_fix_consrv_moment( + w: FloatField, + w2: FloatField, + dp2: FloatField, + gz: FloatFieldIJ, + w_max: Float, + w_min: Float, + compute_performed: BoolFieldIJ, +): + """ + Args: + w (in/out): + w2 (in?): + dp2(in): + w_max(in): + w_min(in): + compute_performed: (Internal Temporary), + """ + + with computation(PARALLEL), interval(...): + w2 = w + + with computation(FORWARD): + with interval(0,1): + compute_performed = False + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + compute_performed = True + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + compute_performed = True + with interval(1,-1): + if(compute_performed): + w2 = w2 + gz / dp2 + compute_performed = False + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + compute_performed = True + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + compute_performed = True + + with computation(BACKWARD): + with interval(-1,None): + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + with interval(1,-1): + w2 = w2 + gz / dp2 + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + +class testClass: + """ + Class to test with DaCe orchestration. test class is MoistCVPlusPt_2d + """ + + def __init__( + self, + stencil_factory: StencilFactory, + grid, + ): + + + self._w_fix_consrv_moment = stencil_factory.from_origin_domain( + func=W_fix_consrv_moment, + origin=(3,3,0), + domain=(24,1,72), + ) + + def __call__( + self, + w: FloatField, + w2: FloatField, + dp2: FloatField, + gz: FloatFieldIJ, + w_max: Float, + w_min: Float, + ): + self._w_fix_consrv_moment( + w, + w2, + dp2, + gz, + w_max, + w_min, + ) + + +class TranslateW_fix_consrv_moment(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.compute_func = testClass(self.stencil_factory, self.grid) # type: ignore + self.quantity_factory = grid.quantity_factory + + print("Running W_fix_consrv_moment translate test") + print("npz = ", grid.npz) + self.in_vars["data_vars"] = { + "w": { + "kend": grid.npz-1, + }, + "w2": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1,}, + "dp2_W": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1,}, + + } + + self.in_vars["parameters"] = ["w_max", "w_min"] + + self.out_vars = { + "w": { + "kend": grid.npz-1, + }, + "w2": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1,}, + "dp2_W": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1,}, + } + self._gz = self.quantity_factory._numpy.zeros( + ( + 31, + 31, + ), dtype=Float, + ) + + self._compute_performed = self.quantity_factory._numpy.zeros( + ( + 31, + 31, + ), dtype=bool, + ) + + def compute_from_storage(self, inputs): + + # print("inputs[w] shape = ", inputs["w"].shape) + # print("inputs[w2] shape = ", inputs["w2"].shape) + # print('self.grid.is_ = ', self.grid.is_) + # print('self.grid.ie = ', self.grid.ie) + # print('self.grid.js = ', self.grid.js) + # print('self.grid.je = ', self.grid.je) + # print('self.storage_vars() = ', self.storage_vars()) + # self.make_storage_data_input_vars(inputs) + # exit(1) + self.compute_func( + inputs["w"], + inputs["w2"], + inputs["dp2_W"], + self._gz, + inputs["w_max"], + inputs["w_min"], + self._compute_performed + ) + + for i in range(3,27): + for k in range(72): + temp = inputs["w2"][i,3,k] + inputs["w2"][i,:,k] = temp + return inputs From 59cec19f6b87afa77cd475a5c02b21f432acfd91 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 24 Sep 2024 14:02:29 -0700 Subject: [PATCH 003/252] Added the "compute_performed" test to the BACKWARD portion of W_fix_consrv_moment. Note that the current serialized results from tinyBC does not test W_fix_consrv_moment since the code never goes into the branches based on if(w2>w_max) or if(w2 < w_min). --- .../translate/translate_w_fix_consrv_moment.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py index 97194e49..3e7086ef 100644 --- a/tests/savepoint/translate/translate_w_fix_consrv_moment.py +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -52,20 +52,27 @@ def W_fix_consrv_moment( with computation(BACKWARD): with interval(-1,None): + compute_performed = False if(w2 > w_max): gz = (w2 - w_max) * dp2 w2 = w_max + compute_performed = True elif(w2 < w_min): gz = (w2 - w_min) * dp2 w2 = w_min + compute_performed = True with interval(1,-1): - w2 = w2 + gz / dp2 + if(compute_performed): + w2 = w2 + gz / dp2 + compute_performed = False if(w2 > w_max): gz = (w2 - w_max) * dp2 w2 = w_max + compute_performed = True elif(w2 < w_min): gz = (w2 - w_min) * dp2 w2 = w_min + compute_performed = True class testClass: """ @@ -93,6 +100,7 @@ def __call__( gz: FloatFieldIJ, w_max: Float, w_min: Float, + compute_performed_bool: BoolFieldIJ ): self._w_fix_consrv_moment( w, @@ -101,6 +109,7 @@ def __call__( gz, w_max, w_min, + compute_performed_bool, ) From 39bc1896b5ec0004c8acc3cfbfc5d33f3b12a6e4 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 24 Sep 2024 14:21:19 -0700 Subject: [PATCH 004/252] Added remaining computations performed in the fix for w's momentum. Note again that TinyBC serialized data does not test the W_fix_consrv_moment adequetely. --- .../savepoint/translate/translate_w_fix_consrv_moment.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py index 3e7086ef..f92e3391 100644 --- a/tests/savepoint/translate/translate_w_fix_consrv_moment.py +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -74,6 +74,14 @@ def W_fix_consrv_moment( w2 = w_min compute_performed = True + with computation(FORWARD), interval(0,1): + if(w2 > (w_max * 2.0)): + w2 = w_max * 2.0 + elif(w2 < (w_min * 2.0)): + w2 = w_min * 2.0 + + with computation(PARALLEL), interval(...): + w = w2 class testClass: """ Class to test with DaCe orchestration. test class is MoistCVPlusPt_2d From fcab87e76f87407f7f332212957f1f315746fad6 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 11 Oct 2024 07:59:53 -0700 Subject: [PATCH 005/252] Added ported additions to map_scalar that exist in GEOS dycore but not in the GFDL dycore. The additions are in map_single.py. translate_lagrangian_contribution.py tests these new additions, and currently, the translate tests only seem to pass with the rank 0 and rank 4 dataset, but it's stuck when executed with the other ranks. --- pyFV3/stencils/map_single.py | 148 ++++++++--- .../translate_lagrangian_contribution.py | 247 ++++++++++++++++++ .../translate_w_fix_consrv_moment.py | 58 ++-- 3 files changed, 376 insertions(+), 77 deletions(-) create mode 100644 tests/savepoint/translate/translate_lagrangian_contribution.py diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index a768f86a..b5b768df 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -4,7 +4,7 @@ from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntFieldIJ # noqa: F401 +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntField, IntFieldIJ, BoolFieldIJ # noqa: F401 from ndsl.stencils.basic_operations import copy_defn from pyFV3.stencils.remap_profile import RemapProfile @@ -17,6 +17,11 @@ def set_dp(dp1: FloatField, pe1: FloatField, lev: IntFieldIJ): def lagrangian_contributions( + km: int, + not_exit_loop: BoolFieldIJ, + INDEX_LM1: IntField, + INDEX_LP0: IntField, + # q2: FloatField, q: FloatField, pe1: FloatField, pe2: FloatField, @@ -26,10 +31,16 @@ def lagrangian_contributions( q4_4: FloatField, dp1: FloatField, lev: IntFieldIJ, + K0: FloatField, ): """ Args: - q (out): + km (in): + not_exit_loop (in/temp): + LM1 (in/temp): + LP0 (in/temp): + *NOT NEEDED * -> q2 (in/temp): + q (in/out): pe1 (in): pe2 (in): q4_1 (in): @@ -39,45 +50,101 @@ def lagrangian_contributions( dp1 (in): lev (inout): """ + + # This computation creates a IntField that allows for "absolute" references in the k-dimension + # for q and pe1. + + # INDEX_LM1 and INDEX_LP0 is initialized such that if it's plugged into "q" (ex: q[0,0,INDEX_LM1]), the k level in + # q is "k = 0". + + # For example, during the stencil computation at k = 2, INDEX_LM1[i,j,2] = -2 + with computation(FORWARD): + with interval(0,1): + INDEX_LM1 = 0 + INDEX_LP0 = 0 + with interval(1,None): + INDEX_LM1 = INDEX_LM1[0,0,-1] - 1 + INDEX_LP0 = INDEX_LP0[0,0,-1] - 1 + # TODO: Can we make lev a 2D temporary? with computation(FORWARD), interval(...): - pl = (pe2 - pe1[0, 0, lev]) / dp1[0, 0, lev] - if pe2[0, 0, 1] <= pe1[0, 0, lev + 1]: - pr = (pe2[0, 0, 1] - pe1[0, 0, lev]) / dp1[0, 0, lev] - q = ( - q4_2[0, 0, lev] - + 0.5 - * (q4_4[0, 0, lev] + q4_3[0, 0, lev] - q4_2[0, 0, lev]) - * (pr + pl) - - q4_4[0, 0, lev] * 1.0 / 3.0 * (pr * (pr + pl) + pl * pl) - ) + LM1 = 1 + LP0 = 1 + not_exit_loop = True + while(LP0 <= km and not_exit_loop): + if(pe1[0,0,INDEX_LP0] < pe2): + LP0 = LP0 + 1 + INDEX_LP0 = INDEX_LP0 + 1 + else: + not_exit_loop = False + + LM1 = max(LP0-1,1) + INDEX_LM1 = INDEX_LM1 + (LM1 - 1) + LP0 = min(LP0, km) + + + + if(LP0 == 1): + INDEX_LP0 = INDEX_LM1 + elif(LP0 <= km): + INDEX_LP0 = INDEX_LM1+1 else: - qsum = (pe1[0, 0, lev + 1] - pe2) * ( - q4_2[0, 0, lev] - + 0.5 - * (q4_4[0, 0, lev] + q4_3[0, 0, lev] - q4_2[0, 0, lev]) - * (1.0 + pl) - - q4_4[0, 0, lev] * 1.0 / 3.0 * (1.0 + pl * (1.0 + pl)) - ) - lev = lev + 1 - while pe1[0, 0, lev + 1] < pe2[0, 0, 1]: - qsum += dp1[0, 0, lev] * q4_1[0, 0, lev] + INDEX_LP0 = INDEX_LM1 + + if(LM1 == 1 and LP0 == 1): + q_temp = q[0,0,INDEX_LM1] + (q[0,0,INDEX_LM1+1] - q[0,0,INDEX_LM1]) * (pe2 - pe1[0,0,INDEX_LM1]) \ + / (pe1[0,0,INDEX_LM1+1] - pe1[0,0,INDEX_LM1]) + + elif(LM1 == km and LP0 == km): + q_temp = q[0,0,INDEX_LM1] + (q[0,0,INDEX_LM1] - q[0,0,INDEX_LM1-1]) * (pe2 - pe1[0,0,INDEX_LM1]) \ + / (pe1[0,0,INDEX_LM1] - pe1[0,0,INDEX_LM1-1]) + + elif(LM1 == 1 or LP0 == km): + q_temp = q[0,0,INDEX_LP0] + (q[0,0,INDEX_LM1] - q[0,0,INDEX_LP0]) * (pe2 - pe1[0,0,INDEX_LP0]) \ + / (pe1[0,0,INDEX_LM1] - pe1[0,0,INDEX_LP0]) + + else: + pl = (pe2 - pe1[0, 0, lev]) / dp1[0, 0, lev] + if pe2[0, 0, 1] <= pe1[0, 0, lev + 1]: + pr = (pe2[0, 0, 1] - pe1[0, 0, lev]) / dp1[0, 0, lev] + q_temp = ( + q4_2[0, 0, lev] + + 0.5 + * (q4_4[0, 0, lev] + q4_3[0, 0, lev] - q4_2[0, 0, lev]) + * (pr + pl) + - q4_4[0, 0, lev] * 1.0 / 3.0 * (pr * (pr + pl) + pl * pl) + ) + else: + qsum = (pe1[0, 0, lev + 1] - pe2) * ( + q4_2[0, 0, lev] + + 0.5 + * (q4_4[0, 0, lev] + q4_3[0, 0, lev] - q4_2[0, 0, lev]) + * (1.0 + pl) + - q4_4[0, 0, lev] * 1.0 / 3.0 * (1.0 + pl * (1.0 + pl)) + ) lev = lev + 1 - dp = pe2[0, 0, 1] - pe1[0, 0, lev] - esl = dp / dp1[0, 0, lev] - qsum += dp * ( - q4_2[0, 0, lev] - + 0.5 - * esl - * ( - q4_3[0, 0, lev] - - q4_2[0, 0, lev] - + q4_4[0, 0, lev] * (1.0 - (2.0 / 3.0) * esl) + while pe1[0, 0, lev + 1] < pe2[0, 0, 1]: + qsum += dp1[0, 0, lev] * q4_1[0, 0, lev] + lev = lev + 1 + dp = pe2[0, 0, 1] - pe1[0, 0, lev] + esl = dp / dp1[0, 0, lev] + qsum += dp * ( + q4_2[0, 0, lev] + + 0.5 + * esl + * ( + q4_3[0, 0, lev] + - q4_2[0, 0, lev] + + q4_4[0, 0, lev] * (1.0 - (2.0 / 3.0) * esl) + ) ) - ) - q = qsum / (pe2[0, 0, 1] - pe2) - lev = lev - 1 + q_temp = qsum / (pe2[0, 0, 1] - pe2) + + lev = lev - 1 + q = q_temp + + K0 = lev class MapSingle: """ @@ -142,6 +209,12 @@ def make_quantity(): compute_dims=dims, ) + self._INDEX_LM1 = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=int) + self._INDEX_LP0 = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=int) + self._km = grid_indexing.domain[2] + self._not_exit_loop = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=bool) + # self._q_temp = make_quantity() + @property def i_extent(self): return self._extents[0] @@ -193,6 +266,11 @@ def __call__( qmin, ) self._lagrangian_contributions( + self._km, + self._not_exit_loop, + self._INDEX_LM1, + self._INDEX_LP0, + # self._q_temp, q1, pe1, pe2, diff --git a/tests/savepoint/translate/translate_lagrangian_contribution.py b/tests/savepoint/translate/translate_lagrangian_contribution.py new file mode 100644 index 00000000..f68268b5 --- /dev/null +++ b/tests/savepoint/translate/translate_lagrangian_contribution.py @@ -0,0 +1,247 @@ +from ndsl import StencilFactory, QuantityFactory +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Float, BoolFieldIJ, IntField, IntFieldIJ, Float +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid +from gt4py.cartesian.gtscript import PARALLEL, FORWARD, BACKWARD, computation, interval +from pyFV3.stencils.map_single import lagrangian_contributions + +class test_Lagragian_Contribution: + def __init__( + self, + stencil_factory: StencilFactory, + grid: Grid, + ): + print("In test_Lagragian_Contribution") + + self._lagrangian_contributions = stencil_factory.from_origin_domain( + func=lagrangian_contributions, + origin=(3,3,0), + domain=(24,1,72), + ) + + def __call__( + self, + km: int, + not_exit_loop: BoolFieldIJ, + INDEX_LM1: IntField, + INDEX_LP0: IntField, + # q2: FloatField, + q: FloatField, + pe1: FloatField, + pe2: FloatField, + q4_1: FloatField, + q4_2: FloatField, + q4_3: FloatField, + q4_4: FloatField, + dp1: FloatField, + lev: IntFieldIJ, + K0: IntField, + ): + self._lagrangian_contributions( + km, + not_exit_loop, + INDEX_LM1, + INDEX_LP0, + # q2, + q, + pe1, + pe2, + q4_1, + q4_2, + q4_3, + q4_4, + dp1, + lev, + K0, + ) + +class TranslateLagrangian_Contribution(TranslateFortranData2Py): + def __init__(self, grid: Grid, namelist, stencil_factory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.compute_func = test_Lagragian_Contribution(self.stencil_factory, self.grid) # type: ignore + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "q1": { + "kend": grid.npz-1, + }, + + "pe1_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "q4_1": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_2": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_3": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_4": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "dp1_":{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + } + } + + self.out_vars = { + "q1": { + "kend": grid.npz-1, + }, + + # "q2": { + # "kend": grid.npz-1, + # }, + + # "pe1_": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + # }, + # "pe2_": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + # }, + # "q4_1": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + # "q4_2": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + # "q4_3": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + # "q4_4": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + # "dp1_":{ + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # } + } + + def compute_from_storage(self, inputs): + + # print('self._not_exit_loop shape: ',self._not_exit_loop.shape) + # print('self._INDEX_LM1 shape: ', self._INDEX_LM1.shape) + # print('self._INDEX_LP0 shape: ', self._INDEX_LP0.shape)) + # print('inputs["q1"] shape: ', inputs["q1"].shape) + # print('inputs["pe1"] shape:',inputs["pe1"].shape) + # print('inputs["pe2"] shape:',inputs["pe2"].shape) + # print('inputs["q4_1"] shape:',inputs["q4_1"].shape) + # print('inputs["q4_2"] shape:',inputs["q4_2"].shape) + # print('inputs["q4_3"] shape:',inputs["q4_3"].shape) + # print('inputs["q4_4"] shape:',inputs["q4_4"].shape) + # print('inputs["dp1_"] shape:',inputs["dp1_"].shape) + # print('self._lev shape:', self._lev.shape) + + self._not_exit_loop = self.quantity_factory.zeros( + [X_DIM, Y_DIM], + units="", + dtype=bool, + ) + + self._INDEX_LM1 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="", + dtype=int, + ) + + self._INDEX_LP0 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="", + dtype=int, + ) + + self._lev = self.quantity_factory.zeros( + [X_DIM, Y_DIM], + units="", + dtype=int, + ) + + self._K0 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="", + dtype=Float, + ) + + self.compute_func( + self.grid.npz, + self._not_exit_loop, + self._INDEX_LM1, + self._INDEX_LP0, + inputs["q1"], + inputs["pe1_"], + inputs["pe2_"], + inputs["q4_1"], + inputs["q4_2"], + inputs["q4_3"], + inputs["q4_4"], + inputs["dp1_"], + self._lev, + self._K0 + ) + + # print("k0[3,3,:] = ", self._K0.data[3,3,:]) + + return inputs \ No newline at end of file diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py index f92e3391..1a07efe6 100644 --- a/tests/savepoint/translate/translate_w_fix_consrv_moment.py +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -2,6 +2,7 @@ from ndsl import StencilFactory, QuantityFactory from ndsl.dsl.typing import FloatField, FloatFieldIJ, Float, BoolFieldIJ +from ndsl.stencils.testing.grid import Grid from ndsl.stencils.testing import TranslateFortranData2Py def W_fix_consrv_moment( @@ -83,16 +84,14 @@ def W_fix_consrv_moment( with computation(PARALLEL), interval(...): w = w2 class testClass: - """ - Class to test with DaCe orchestration. test class is MoistCVPlusPt_2d - """ def __init__( self, stencil_factory: StencilFactory, - grid, + grid: Grid, ): - + print("In testClass") + print("Grid.nid Grid.njd Grid.npz", grid.nid, grid.njd, grid.npz) self._w_fix_consrv_moment = stencil_factory.from_origin_domain( func=W_fix_consrv_moment, @@ -122,7 +121,7 @@ def __call__( class TranslateW_fix_consrv_moment(TranslateFortranData2Py): - def __init__(self, grid, namelist, stencil_factory): + def __init__(self, grid: Grid, namelist, stencil_factory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid @@ -130,21 +129,13 @@ def __init__(self, grid, namelist, stencil_factory): self.quantity_factory = grid.quantity_factory print("Running W_fix_consrv_moment translate test") - print("npz = ", grid.npz) + self.in_vars["data_vars"] = { "w": { "kend": grid.npz-1, }, - "w2": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1,}, - "dp2_W": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1,}, + "w2": grid.compute_dict(), + "dp2_W": grid.compute_dict(), } @@ -154,42 +145,25 @@ def __init__(self, grid, namelist, stencil_factory): "w": { "kend": grid.npz-1, }, - "w2": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1,}, - "dp2_W": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1,}, + "w2": grid.compute_dict(), + "dp2_W": grid.compute_dict(), } self._gz = self.quantity_factory._numpy.zeros( ( - 31, - 31, + grid.nid, + grid.njd, ), dtype=Float, ) self._compute_performed = self.quantity_factory._numpy.zeros( ( - 31, - 31, + grid.nid, + grid.njd, ), dtype=bool, ) def compute_from_storage(self, inputs): - # print("inputs[w] shape = ", inputs["w"].shape) - # print("inputs[w2] shape = ", inputs["w2"].shape) - # print('self.grid.is_ = ', self.grid.is_) - # print('self.grid.ie = ', self.grid.ie) - # print('self.grid.js = ', self.grid.js) - # print('self.grid.je = ', self.grid.je) - # print('self.storage_vars() = ', self.storage_vars()) - # self.make_storage_data_input_vars(inputs) - # exit(1) self.compute_func( inputs["w"], inputs["w2"], @@ -200,8 +174,8 @@ def compute_from_storage(self, inputs): self._compute_performed ) - for i in range(3,27): - for k in range(72): + for i in range(self.grid.is_,self.grid.ie+1): + for k in range(self.grid.npz): temp = inputs["w2"][i,3,k] inputs["w2"][i,:,k] = temp return inputs From f39f87dc2d20a393abb670be6cd76e264caf1106 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 14 Oct 2024 12:59:50 -0700 Subject: [PATCH 006/252] Corrected else branch by adding if branch contained within Fortran code. Ported code passes translate tests with C24 dataset. --- pyFV3/stencils/map_single.py | 10 +-- .../translate_lagrangian_contribution.py | 80 ------------------- 2 files changed, 3 insertions(+), 87 deletions(-) diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index b5b768df..6361ce88 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -21,7 +21,6 @@ def lagrangian_contributions( not_exit_loop: BoolFieldIJ, INDEX_LM1: IntField, INDEX_LP0: IntField, - # q2: FloatField, q: FloatField, pe1: FloatField, pe2: FloatField, @@ -31,7 +30,6 @@ def lagrangian_contributions( q4_4: FloatField, dp1: FloatField, lev: IntFieldIJ, - K0: FloatField, ): """ Args: @@ -82,8 +80,6 @@ def lagrangian_contributions( INDEX_LM1 = INDEX_LM1 + (LM1 - 1) LP0 = min(LP0, km) - - if(LP0 == 1): INDEX_LP0 = INDEX_LM1 elif(LP0 <= km): @@ -104,6 +100,8 @@ def lagrangian_contributions( / (pe1[0,0,INDEX_LM1] - pe1[0,0,INDEX_LP0]) else: + while(pe2 < pe1[0,0,lev] or pe2 > pe1[0,0,lev+1]): + lev = lev + 1 pl = (pe2 - pe1[0, 0, lev]) / dp1[0, 0, lev] if pe2[0, 0, 1] <= pe1[0, 0, lev + 1]: pr = (pe2[0, 0, 1] - pe1[0, 0, lev]) / dp1[0, 0, lev] @@ -140,12 +138,10 @@ def lagrangian_contributions( ) q_temp = qsum / (pe2[0, 0, 1] - pe2) - lev = lev - 1 + lev = lev - 1 q = q_temp - K0 = lev - class MapSingle: """ Fortran name is map_single, test classes are Map1_PPM_2d, Map_Scalar_2d diff --git a/tests/savepoint/translate/translate_lagrangian_contribution.py b/tests/savepoint/translate/translate_lagrangian_contribution.py index f68268b5..e5ba3646 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution.py @@ -26,7 +26,6 @@ def __call__( not_exit_loop: BoolFieldIJ, INDEX_LM1: IntField, INDEX_LP0: IntField, - # q2: FloatField, q: FloatField, pe1: FloatField, pe2: FloatField, @@ -36,14 +35,12 @@ def __call__( q4_4: FloatField, dp1: FloatField, lev: IntFieldIJ, - K0: IntField, ): self._lagrangian_contributions( km, not_exit_loop, INDEX_LM1, INDEX_LP0, - # q2, q, pe1, pe2, @@ -53,7 +50,6 @@ def __call__( q4_4, dp1, lev, - K0, ) class TranslateLagrangian_Contribution(TranslateFortranData2Py): @@ -124,77 +120,10 @@ def __init__(self, grid: Grid, namelist, stencil_factory): "q1": { "kend": grid.npz-1, }, - - # "q2": { - # "kend": grid.npz-1, - # }, - - # "pe1_": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz - # }, - # "pe2_": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz - # }, - # "q4_1": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, - # "q4_2": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, - # "q4_3": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, - # "q4_4": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, - # "dp1_":{ - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # } } def compute_from_storage(self, inputs): - # print('self._not_exit_loop shape: ',self._not_exit_loop.shape) - # print('self._INDEX_LM1 shape: ', self._INDEX_LM1.shape) - # print('self._INDEX_LP0 shape: ', self._INDEX_LP0.shape)) - # print('inputs["q1"] shape: ', inputs["q1"].shape) - # print('inputs["pe1"] shape:',inputs["pe1"].shape) - # print('inputs["pe2"] shape:',inputs["pe2"].shape) - # print('inputs["q4_1"] shape:',inputs["q4_1"].shape) - # print('inputs["q4_2"] shape:',inputs["q4_2"].shape) - # print('inputs["q4_3"] shape:',inputs["q4_3"].shape) - # print('inputs["q4_4"] shape:',inputs["q4_4"].shape) - # print('inputs["dp1_"] shape:',inputs["dp1_"].shape) - # print('self._lev shape:', self._lev.shape) - self._not_exit_loop = self.quantity_factory.zeros( [X_DIM, Y_DIM], units="", @@ -219,12 +148,6 @@ def compute_from_storage(self, inputs): dtype=int, ) - self._K0 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="", - dtype=Float, - ) - self.compute_func( self.grid.npz, self._not_exit_loop, @@ -239,9 +162,6 @@ def compute_from_storage(self, inputs): inputs["q4_4"], inputs["dp1_"], self._lev, - self._K0 ) - # print("k0[3,3,:] = ", self._K0.data[3,3,:]) - return inputs \ No newline at end of file From 48ece2fd3d6656495a6bec5925138abd0941ff4a Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 14 Oct 2024 13:07:37 -0700 Subject: [PATCH 007/252] Modified __init.py to add in translate test import statements --- tests/savepoint/translate/__init__.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 651b2551..58deb657 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -85,3 +85,6 @@ from .translate_xtp_u import TranslateXTP_U from .translate_yppm import TranslateYPPM, TranslateYPPM_2 from .translate_ytp_v import TranslateYTP_V + +from .translate_w_fix_consrv_moment import TranslateW_fix_consrv_moment +from .translate_lagrangian_contribution import TranslateLagrangian_Contribution From c5a3e5476ecf29ebbe02542d79632a62e253d2c4 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 15 Oct 2024 09:05:21 -0700 Subject: [PATCH 008/252] Very minor edit in the comments to map_single.py. The important note is that the code passes the Translate tests after I found and corrected a Fortran data serialization bug that prevented the complete pressure (pe1 and pe2) dataset to the translate test. --- pyFV3/stencils/map_single.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index 6361ce88..b6e055c8 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -37,7 +37,6 @@ def lagrangian_contributions( not_exit_loop (in/temp): LM1 (in/temp): LP0 (in/temp): - *NOT NEEDED * -> q2 (in/temp): q (in/out): pe1 (in): pe2 (in): From 6b27c2ab76714746a099a8264be12dfe8d8bbcd6 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 21 Oct 2024 16:34:29 -0700 Subject: [PATCH 009/252] Starting assembly of fv_mapz stencils. Initial indication that init_pe and part of moist_cv_pt_pressure are passing translate tests --- pyFV3/stencils/remapping.py | 76 ++++---- tests/savepoint/translate/__init__.py | 1 + .../translate_lagrangian_contribution.py | 2 +- .../translate/translate_remapping_GEOS.py | 181 ++++++++++++++++++ 4 files changed, 222 insertions(+), 38 deletions(-) create mode 100644 tests/savepoint/translate/translate_remapping_GEOS.py diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 901fccc3..9ab61cca 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -92,14 +92,14 @@ def moist_cv_pt_pressure( cappa: FloatField, delp: FloatField, delz: FloatField, - pe: FloatField, - pe2: FloatField, - ak: FloatFieldK, - bk: FloatFieldK, - dp2: FloatField, - ps: FloatFieldIJ, - pn2: FloatField, - peln: FloatField, + # pe: FloatField, + # pe2: FloatField, + # ak: FloatFieldK, + # bk: FloatFieldK, + # dp2: FloatField, + # ps: FloatFieldIJ, + # pn2: FloatField, + # peln: FloatField, remap_t: bool, r_vir: Float, ): @@ -118,14 +118,14 @@ def moist_cv_pt_pressure( cappa (out): delp (inout): delz (inout): - pe (in): - pe2 (inout): - ak (in): - bk (in): - dp2 (out): - ps (out): - pn2 (out): - peln (in): + # pe (in): + # pe2 (inout): + # ak (in): + # bk (in): + # dp2 (out): + # ps (out): + # pn2 (out): + # peln (in): """ from __externals__ import hydrostatic#, kord_tm @@ -151,27 +151,29 @@ def moist_cv_pt_pressure( # # delz_adjust # if __INLINED(not hydrostatic): # delz = -delz / delp - # pressure_updates - with computation(FORWARD): - with interval(-1, None): - ps = pe - with computation(PARALLEL): - with interval(0, 1): - pn2 = peln - # TODO: refactor the pe2 = ptop assignment from - # previous stencil into this one, and remove - # pe2 from the other stencil - with interval(1, -1): - pe2 = ak + bk * ps - with interval(-1, None): - pn2 = peln - with computation(BACKWARD), interval(0, -1): - dp2 = pe2[0, 0, 1] - pe2 - - # NOTE : GEOS doesn't perform the delp calcuation at this location - # copy_stencil - # with computation(PARALLEL), interval(0, -1): - # delp = dp2 + + ### Comment out for now ### + # # pressure_updates + # with computation(FORWARD): + # with interval(-1, None): + # ps = pe + # with computation(PARALLEL): + # with interval(0, 1): + # pn2 = peln + # # TODO: refactor the pe2 = ptop assignment from + # # previous stencil into this one, and remove + # # pe2 from the other stencil + # with interval(1, -1): + # pe2 = ak + bk * ps + # with interval(-1, None): + # pn2 = peln + # with computation(BACKWARD), interval(0, -1): + # dp2 = pe2[0, 0, 1] - pe2 + + # # NOTE : GEOS doesn't perform the delp calcuation at this location + # # copy_stencil + # # with computation(PARALLEL), interval(0, -1): + # # delp = dp2 def pn2_pk_delp( diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 58deb657..5361fa00 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -88,3 +88,4 @@ from .translate_w_fix_consrv_moment import TranslateW_fix_consrv_moment from .translate_lagrangian_contribution import TranslateLagrangian_Contribution +from .translate_remapping_GEOS import TranslateRemapping_GEOS \ No newline at end of file diff --git a/tests/savepoint/translate/translate_lagrangian_contribution.py b/tests/savepoint/translate/translate_lagrangian_contribution.py index e5ba3646..4013507b 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution.py @@ -113,7 +113,7 @@ def __init__(self, grid: Grid, namelist, stencil_factory): "jstart": grid.js, "jend": grid.je, "kend": grid.npz-1, - } + }, } self.out_vars = { diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py new file mode 100644 index 00000000..b195b2ba --- /dev/null +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -0,0 +1,181 @@ +import ndsl.dsl.gt4py_utils as utils +from ndsl import Namelist, StencilFactory +from pyFV3 import DynamicalCoreConfig +from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure +from pyFV3.stencils import moist_cv +from ndsl.stencils.testing import pad_field_in_j +from pyFV3.testing import TranslateDycoreFortranData2Py +from ndsl.constants import ( + X_DIM, + X_INTERFACE_DIM, + Y_DIM, + Y_INTERFACE_DIM, + Z_DIM, + Z_INTERFACE_DIM, +) +from ndsl.dsl.typing import Float + +class TranslateRemapping_GEOS(TranslateDycoreFortranData2Py): + def __init__( + self, + grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, namelist, stencil_factory) + + self.in_vars["data_vars"] = { + + "pe_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + # "kaxis": 1, + }, + "pe1_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + # "kaxis": 1, + }, + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + # "kaxis": 1, + }, + # Note that tracers are i x k shaped. + # Setting "axis" as 1 enables the translate test to read them properly + "qvapor": { + "axis": 1 + }, + "qliquid": { + "axis": 1 + }, + "qice": { + "axis": 1 + }, + "qrain": { + "axis": 1 + }, + "qsnow": { + "axis": 1 + }, + "qgraupel": { + "axis": 1 + }, + "delp": {}, + "delz": {}, + "q_con": {}, + "pt": {}, + "cappa": {}, + } + self.in_vars["parameters"] = [ + "ptop", + "r_vir", + # "remap_t", # For some reason, translate test can't accept a logical variable + # "akap", + # "zvir", + # "last_step", + # "consv_te", + # "mdt", + # "nq", + ] + self.out_vars = { + "pe1_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + # "kaxis": 1, + }, + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + # "kaxis": 1, + }, + # "delp": {}, + # "delz": {}, + # "q_con": {}, + # "pt": {}, + "cappa": {}, + } + + self.stencil_factory = stencil_factory + # self.namelist found in TranslateDycoreFortranData2Py + # self.namelist = DynamicalCoreConfig.from_namelist(namelist) + config=DynamicalCoreConfig.from_namelist(self.namelist).remapping + + hydrostatic = config.hydrostatic + if hydrostatic: + raise NotImplementedError("Hydrostatic is not implemented") + + grid_indexing = stencil_factory.grid_indexing + self._domain_jextra = ( + grid_indexing.domain[0], + grid_indexing.domain[1] + 1, + grid_indexing.domain[2] + 1, + ) + + self._init_pe = stencil_factory.from_origin_domain( + init_pe, + # origin=(3,3,0), + origin=grid_indexing.origin_compute(), + domain=(24,1,73), + ) + + self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( + moist_cv_pt_pressure, + # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, + externals={"hydrostatic": hydrostatic}, + origin=grid_indexing.origin_compute(), + # domain=grid_indexing.domain_compute(add=(0, 0, 1)), + domain=(grid.nic, 1, grid.npz+1), # Note : Many intervals go from (0,-1) in this stencil + ) + + def compute_from_storage(self, inputs): + + # print("inputs[qvapor].data.shape() 1 = ", inputs["qvapor"].data.shape) + + # Replicates tracer values in I along the J direction + for name, value in inputs.items(): + if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: + inputs[name] = self.make_storage_data( + pad_field_in_j( + value, self.grid.njd, backend=self.stencil_factory.backend + ) + ) + # print("inputs[qvapor].data.shape() 2 = ", inputs["qvapor"].data.shape) + self._init_pe( + inputs["pe_"], + inputs["pe1_"], + inputs["pe2_"], + inputs["ptop"], + ) + + self._moist_cv_pt_pressure( + inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + inputs["q_con"], + inputs["pt"], + inputs["cappa"], + inputs["delp"], + inputs["delz"], + True, + Float(inputs["r_vir"]), + ) + return inputs From 027f3117c9099439b3dfb265186c85ed38740a65 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 23 Oct 2024 08:41:27 -0700 Subject: [PATCH 010/252] Assembled translate tests that tests init_pe and moist_cv_pt_pressure. Current belief is that these stencils pass the C24 translate tests with a limited number of unique errors for pt and cappa that are assumed to be due to rounding from constants. --- pyFV3/stencils/remapping.py | 67 +++++++------- .../translate/translate_remapping_GEOS.py | 90 ++++++++++++++----- 2 files changed, 101 insertions(+), 56 deletions(-) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 9ab61cca..b179ed09 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -92,14 +92,14 @@ def moist_cv_pt_pressure( cappa: FloatField, delp: FloatField, delz: FloatField, - # pe: FloatField, - # pe2: FloatField, - # ak: FloatFieldK, - # bk: FloatFieldK, - # dp2: FloatField, - # ps: FloatFieldIJ, - # pn2: FloatField, - # peln: FloatField, + pe: FloatField, + pe2: FloatField, + ak: FloatFieldK, + bk: FloatFieldK, + dp2: FloatField, + ps: FloatFieldIJ, + pn2: FloatField, + peln: FloatField, remap_t: bool, r_vir: Float, ): @@ -118,14 +118,14 @@ def moist_cv_pt_pressure( cappa (out): delp (inout): delz (inout): - # pe (in): - # pe2 (inout): - # ak (in): - # bk (in): - # dp2 (out): - # ps (out): - # pn2 (out): - # peln (in): + pe (in): + pe2 (inout): + ak (in): + bk (in): + dp2 (out): + ps (out): + pn2 (out): + peln (in): """ from __externals__ import hydrostatic#, kord_tm @@ -151,24 +151,23 @@ def moist_cv_pt_pressure( # # delz_adjust # if __INLINED(not hydrostatic): # delz = -delz / delp - - ### Comment out for now ### - # # pressure_updates - # with computation(FORWARD): - # with interval(-1, None): - # ps = pe - # with computation(PARALLEL): - # with interval(0, 1): - # pn2 = peln - # # TODO: refactor the pe2 = ptop assignment from - # # previous stencil into this one, and remove - # # pe2 from the other stencil - # with interval(1, -1): - # pe2 = ak + bk * ps - # with interval(-1, None): - # pn2 = peln - # with computation(BACKWARD), interval(0, -1): - # dp2 = pe2[0, 0, 1] - pe2 + + # pressure_updates + with computation(FORWARD): + with interval(-1, None): + ps = pe + with computation(PARALLEL): + with interval(0, 1): + pn2 = peln + # TODO: refactor the pe2 = ptop assignment from + # previous stencil into this one, and remove + # pe2 from the other stencil + with interval(1, -1): + pe2 = ak + bk * ps + with interval(-1, None): + pn2 = peln + with computation(BACKWARD), interval(0, -1): + dp2 = pe2[0, 0, 1] - pe2 # # NOTE : GEOS doesn't perform the delp calcuation at this location # # copy_stencil diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index b195b2ba..416743a0 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -3,7 +3,7 @@ from pyFV3 import DynamicalCoreConfig from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure from pyFV3.stencils import moist_cv -from ndsl.stencils.testing import pad_field_in_j +from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py from ndsl.constants import ( X_DIM, @@ -18,7 +18,7 @@ class TranslateRemapping_GEOS(TranslateDycoreFortranData2Py): def __init__( self, - grid, + grid: Grid, namelist: Namelist, stencil_factory: StencilFactory, ): @@ -52,22 +52,22 @@ def __init__( }, # Note that tracers are i x k shaped. # Setting "axis" as 1 enables the translate test to read them properly - "qvapor": { + "qvapor_": { "axis": 1 }, - "qliquid": { + "qliquid_": { "axis": 1 }, - "qice": { + "qice_": { "axis": 1 }, - "qrain": { + "qrain_": { "axis": 1 }, - "qsnow": { + "qsnow_": { "axis": 1 }, - "qgraupel": { + "qgraupel_": { "axis": 1 }, "delp": {}, @@ -75,6 +75,24 @@ def __init__( "q_con": {}, "pt": {}, "cappa": {}, + "ps": {}, + "pn2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, + "peln_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, + "ak": {}, + "bk": {}, + "dp2_3d": grid.compute_dict() } self.in_vars["parameters"] = [ "ptop", @@ -104,14 +122,25 @@ def __init__( "kend": grid.npz + 1, # "kaxis": 1, }, - # "delp": {}, - # "delz": {}, - # "q_con": {}, - # "pt": {}, + "delp": {}, + "delz": {}, + "q_con": {}, + "pt": {}, "cappa": {}, + "ps": {}, + "dp2_3d": grid.compute_dict(), + "pn2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, } self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + # self.namelist found in TranslateDycoreFortranData2Py # self.namelist = DynamicalCoreConfig.from_namelist(namelist) config=DynamicalCoreConfig.from_namelist(self.namelist).remapping @@ -127,11 +156,22 @@ def __init__( grid_indexing.domain[2] + 1, ) + # self._pe1 = self.quantity_factory.zeros( + # [X_DIM, Y_DIM, Z_INTERFACE_DIM], + # units="Pa", + # dtype=Float, + # ) + # self._pe2 = self.quantity_factory.zeros( + # [X_DIM, Y_DIM, Z_INTERFACE_DIM], + # units="Pa", + # dtype=Float, + # ) + self._init_pe = stencil_factory.from_origin_domain( init_pe, # origin=(3,3,0), origin=grid_indexing.origin_compute(), - domain=(24,1,73), + domain=(grid.nic,1,73), ) self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( @@ -143,9 +183,7 @@ def __init__( domain=(grid.nic, 1, grid.npz+1), # Note : Many intervals go from (0,-1) in this stencil ) - def compute_from_storage(self, inputs): - - # print("inputs[qvapor].data.shape() 1 = ", inputs["qvapor"].data.shape) + def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction for name, value in inputs.items(): @@ -164,17 +202,25 @@ def compute_from_storage(self, inputs): ) self._moist_cv_pt_pressure( - inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], + inputs["qvapor_"], + inputs["qliquid_"], + inputs["qrain_"], + inputs["qsnow_"], + inputs["qice_"], + inputs["qgraupel_"], inputs["q_con"], inputs["pt"], inputs["cappa"], inputs["delp"], inputs["delz"], + inputs["pe_"], + inputs["pe2_"], + inputs["ak"], + inputs["bk"], + inputs["dp2_3d"], + inputs["ps"], + inputs["pn2_3d"], + inputs["peln_3d"], True, Float(inputs["r_vir"]), ) From 69ac4397d31922608d5239d0667149a173f986ed Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 23 Oct 2024 10:49:49 -0700 Subject: [PATCH 011/252] Added call to pn2_pk_delp into Remapping_GEOS translate test and verified with C24 dataset. With numpy backend, the variables pn2 and pk have unique errors on the order of 10 to 20 per 1000 errors, and the errors are 7 orders of magnitude down from the solution it represents. With the dace cpu backend, both pn2 and pk verify --- pyFV3/stencils/remapping.py | 4 +-- .../translate/translate_remapping_GEOS.py | 33 +++++++++++++++++-- 2 files changed, 32 insertions(+), 5 deletions(-) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index b179ed09..6843a59a 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -176,8 +176,8 @@ def moist_cv_pt_pressure( def pn2_pk_delp( - dp2: FloatField, - delp: FloatField, + # dp2: FloatField, + # delp: FloatField, pe2: FloatField, pn2: FloatField, pk: FloatField, diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 416743a0..fde25872 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,7 +1,7 @@ import ndsl.dsl.gt4py_utils as utils from ndsl import Namelist, StencilFactory from pyFV3 import DynamicalCoreConfig -from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure +from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure, pn2_pk_delp from pyFV3.stencils import moist_cv from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py @@ -92,13 +92,20 @@ def __init__( }, "ak": {}, "bk": {}, - "dp2_3d": grid.compute_dict() + "dp2_3d": grid.compute_dict(), + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + } } self.in_vars["parameters"] = [ "ptop", "r_vir", # "remap_t", # For some reason, translate test can't accept a logical variable - # "akap", + "akap", # "zvir", # "last_step", # "consv_te", @@ -136,6 +143,13 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + } } self.stencil_factory = stencil_factory @@ -183,6 +197,12 @@ def __init__( domain=(grid.nic, 1, grid.npz+1), # Note : Many intervals go from (0,-1) in this stencil ) + self._pn2_pk_delp = stencil_factory.from_origin_domain( + pn2_pk_delp, + origin=grid_indexing.origin_compute(), + domain=(grid.nic, 1, grid.npz+1), + ) + def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction @@ -224,4 +244,11 @@ def compute_from_storage(self, inputs): True, Float(inputs["r_vir"]), ) + + self._pn2_pk_delp( + inputs["pe2_"], + inputs["pn2_3d"], + inputs["pk"], + Float(inputs["akap"]), + ) return inputs From 6076680f6fa2307c11a5c3bd3288c75fd185dc86 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 23 Oct 2024 11:01:14 -0700 Subject: [PATCH 012/252] Adding fv_mapz.F90.SER file that contains data serialization performed for translate tests --- fv_mapz.F90.SER | 4266 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 4266 insertions(+) create mode 100644 fv_mapz.F90.SER diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER new file mode 100644 index 00000000..e9e8f588 --- /dev/null +++ b/fv_mapz.F90.SER @@ -0,0 +1,4266 @@ +!*********************************************************************** +!* GNU Lesser General Public License +!* +!* This file is part of the FV3 dynamical core. +!* +!* The FV3 dynamical core is free software: you can redistribute it +!* and/or modify it under the terms of the +!* GNU Lesser General Public License as published by the +!* Free Software Foundation, either version 3 of the License, or +!* (at your option) any later version. +!* +!* The FV3 dynamical core is distributed in the hope that it will be +!* useful, but WITHOUT ANYWARRANTY; without even the implied warranty +!* of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. +!* See the GNU General Public License for more details. +!* +!* You should have received a copy of the GNU Lesser General Public +!* License along with the FV3 dynamical core. +!* If not, see . +!*********************************************************************** + +!>@brief The module 'fv_mapz' contains the vertical mapping routines \cite lin2004vertically +!>@note April 12, 2012 -SJL: This revision may actually produce rounding level differences +!! due to the elimination of KS to compute pressure level for remapping. + +module fv_mapz_mod + +! Modules Included: +! +! +! +! +! +!
Module NameFunctions Included
+! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +!
constants_modradius, pi=>pi_8, rvgas, rdgas, grav, hlv, hlf, cp_air, cp_vapor
field_manager_modMODEL_ATMOS
fv_arrays_modfv_grid_type
fv_cmp_modqs_init, fv_sat_adj
fv_fill_modfillz
fv_grid_utils_modg_sum, ptop_min
fv_mp_modis_master
fv_timing_modtiming_on, timing_off
fv_tracer2d_modtracer_2d, tracer_2d_1L, tracer_2d_nested
mpp_mod/td> +! NOTE, mpp_error, get_unit, mpp_root_pe, mpp_pe
mpp_domains_mod/td> +! mpp_update_domains, domain2d
tracer_manager_modget_tracer_index
+ !$ser verbatim USE m_serialize, ONLY: fs_is_serialization_on + use constants_mod, only: radius, pi=>pi_8, rvgas, rdgas, grav, hlv, hlf, hls, cp_air, cp_vapor + use tracer_manager_mod,only: get_tracer_index + use field_manager_mod, only: MODEL_ATMOS + use fv_grid_utils_mod, only: g_sum, ptop_min + use fv_fill_mod, only: fillz + use mpp_domains_mod, only: mpp_update_domains, domain2d, mpp_global_sum, BITWISE_EFP_SUM, BITWISE_EXACT_SUM + use mpp_mod, only: NOTE, mpp_error, get_unit, mpp_root_pe, mpp_pe + use fv_arrays_mod, only: fv_grid_type, fv_flags_type + use fv_timing_mod, only: timing_on, timing_off + use fv_mp_mod, only: is_master + use fv_cmp_mod, only: qs_init, fv_sat_adj + + implicit none + real, parameter:: consv_min= 0.001 !< below which no correction applies + real, parameter:: te_min= -1.e25 + real, parameter:: t_min= 184. !< below which applies stricter constraint + real, parameter:: r2=1./2., r0=0.0 + real, parameter:: r3 = 1./3., r23 = 2./3., r12 = 1./12. + real, parameter:: cv_vap = 3.*rvgas !< 1384.5 + real, parameter:: cv_air = cp_air - rdgas !< = rdgas * (7/2-1) = 2.5*rdgas=717.68 +! real, parameter:: c_ice = 2106. !< heat capacity of ice at 0.C + real, parameter:: c_ice = 1972. !< heat capacity of ice at -15.C + real, parameter:: c_liq = 4.1855e+3 !< GFS: heat capacity of water at 0C +! real, parameter:: c_liq = 4218. !< ECMWF-IFS + real, parameter:: cp_vap = cp_vapor !< 1846. + real, parameter:: tice = 273.16 + + logical, parameter :: w_limiter = .true. + real, parameter :: w_max = 90. + real, parameter :: w_min = -60. + + real(kind=8) :: E_Flux = 0. + private + + public compute_total_energy, Lagrangian_to_Eulerian, moist_cv, moist_cp, & + rst_remap, mappm, E_Flux, mapn_tracer, map1_q2 + +!---- version number ----- + character(len=128) :: version = '$Id$' + character(len=128) :: tagname = '$Name$' + +contains + +!>@brief The subroutine 'Lagrangian_to_Eulerian' remaps deformed Lagrangian layers back to the reference Eulerian coordinate. +!>@details It also includes the entry point for calling fast microphysical processes. This is typically calle on the k_split loop. + subroutine Lagrangian_to_Eulerian(last_step, consv, ps, pe, delp, pkz, pk, & + mdt, pdt, km, is,ie,js,je, isd,ied,jsd,jed, & + nq, nwat, sphum, q_con, u, v, w, delz, pt, q, hs, r_vir, cp, & + akap, cappa, kord_mt, kord_wz, kord_tr, kord_tm, peln, te0_2d, & + ng, ua, va, omga, te, ws, fill, reproduce_sum, out_dt, dtdt, & + ptop, ak, bk, pfull, flagstruct, gridstruct, domain, do_sat_adj, & + hydrostatic, hybrid_z, do_omega, adiabatic, do_adiabatic_init, & + mfx, mfy, cx, cy, remap_option, gmao_remap) + logical, intent(in):: last_step + real, intent(in):: mdt !< remap time step + real, intent(in):: pdt !< phys time step + integer, intent(in):: km + integer, intent(in):: nq !< number of tracers (including h2o) + integer, intent(in):: nwat + integer, intent(in):: sphum !< index for water vapor (specific humidity) + integer, intent(in):: ng + integer, intent(in):: is,ie,isd,ied !< starting & ending X-Dir index + integer, intent(in):: js,je,jsd,jed !< starting & ending Y-Dir index + integer, intent(in):: kord_mt !< Mapping order for the vector winds + integer, intent(in):: kord_wz !< Mapping order/option for w + integer, intent(in):: kord_tr(nq) !< Mapping order for tracers + integer, intent(in):: kord_tm !< Mapping order for thermodynamics + + real, intent(in):: consv !< factor for TE conservation + real, intent(in):: r_vir + real, intent(in):: cp + real, intent(in):: akap + real, intent(in):: hs(isd:ied,jsd:jed) !< surface geopotential + real(kind=8), intent(inout):: te0_2d(is:ie,js:je) + real, intent(in):: ws(is:ie,js:je) + + logical, intent(in):: do_sat_adj + logical, intent(in):: fill !< fill negative tracers + logical, intent(in):: reproduce_sum + logical, intent(in):: do_omega, adiabatic, do_adiabatic_init + real, intent(in) :: ptop + real, intent(in) :: ak(km+1) + real, intent(in) :: bk(km+1) + real, intent(in):: pfull(km) + type(fv_grid_type), intent(IN), target :: gridstruct + type(fv_flags_type), intent(INOUT) :: flagstruct + type(domain2d), intent(INOUT) :: domain + + ! INPUT/OUTPUT + real, intent(inout):: pk(is:ie,js:je,km+1) !< pe to the kappa + real, intent(inout):: q(isd:ied,jsd:jed,km,*) + real, intent(inout):: delp(isd:ied,jsd:jed,km) !< pressure thickness + real, intent(inout):: pe(is-1:ie+1,km+1,js-1:je+1) !< pressure at layer edges + real, intent(inout):: ps(isd:ied,jsd:jed) !< surface pressure + + ! u-wind will be ghosted one latitude to the north upon exit + real, intent(inout):: u(isd:ied ,jsd:jed+1,km) !< u-wind (m/s) + real, intent(inout):: v(isd:ied+1,jsd:jed ,km) !< v-wind (m/s) + real, intent(inout):: w(isd: ,jsd: ,1:) !< vertical velocity (m/s) + real, intent(inout):: pt(isd:ied ,jsd:jed ,km) !< cp*virtual potential temperature + !< as input; output: temperature + real, intent(inout), dimension(isd:,jsd:,1:)::delz, q_con, cappa + logical, intent(in):: hydrostatic + logical, intent(in):: hybrid_z + logical, intent(in):: out_dt + + real, intent(inout):: ua(isd:ied,jsd:jed,km) !< u-wind (m/s) on physics grid + real, intent(inout):: va(isd:ied,jsd:jed,km) !< v-wind (m/s) on physics grid + real, intent(inout):: omga(isd:ied,jsd:jed,km) !< vertical press. velocity (pascal/sec) + real, intent(inout):: peln(is:ie,km+1,js:je) !< log(pe) + real, intent(inout):: dtdt(is:ie,js:je,km) + real, intent(out):: pkz(is:ie,js:je,km) !< layer-mean pk for converting t to pt + real, intent(out):: te(isd:ied,jsd:jed,km) + ! Mass fluxes + real, optional, intent(inout):: mfx(is:ie+1,js:je ,km) ! X-dir Mass Flux + real, optional, intent(inout):: mfy(is:ie ,js:je+1,km) ! Y-dir Mass Flux + ! Courant numbers + real, optional, intent(inout):: cx(is:ie+1, jsd:jed,km) + real, optional, intent(inout):: cy(isd:ied ,js:je+1,km) + + integer, intent(in):: remap_option, gmao_remap + + ! !DESCRIPTION: + ! + ! !REVISION HISTORY: + ! SJL 03.11.04: Initial version for partial remapping + ! + !----------------------------------------------------------------------- + real(kind=8), dimension(is:ie,js:je):: te_2d, zsum0, zsum1 + real, dimension(is:ie,js:je):: dpln + real, dimension(is:ie,km) :: q2, dp2, w2 + real, dimension(is:ie,km+1):: pe1, pe2, pk1, pk2, pn1, pn2, phis + real, dimension(is:ie+1,km+1):: pe0, pe3 + real, dimension(is:ie):: gz, cvm + real(kind=8):: tesum, zsum, dtmp + real :: rcp, rg, tmp, tpe, rrg, bkh, k1k, dlnp + logical:: fast_mp_consv + integer:: i,j,k + integer:: nt, liq_wat, ice_wat, rainwat, snowwat, cld_amt, graupel, iq, n, kmp, kp, k_next + logical:: remap_t, remap_pt, remap_te + + real, dimension(is:ie, js:je, km+1) :: pe_3d, pe1_3d, pe2_3d, pn2_3d, pk2_3d, peln_3d + + real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d + +!$ser verbatim real :: w_max_, w_min_ + + !$ser verbatim integer:: mode, abskord, iep1, iedp1, jedp1, js2d, o3mr, sgs_tke + !$ser verbatim real :: qmin + !$ser verbatim qmin = 0.0 + !$ser verbatim iep1=ie+1 + !$ser verbatim iedp1=ied+1 + !$ser verbatim jedp1=jed+1 + !$ser verbatim js2d=js + !$ser verbatim pe_3d = 0.0 + !$ser verbatim pe1_3d = 0.0 + !$ser verbatim pe2_3d = 0.0 + !$ser verbatim peln_3d = 0.0 + !$ser verbatim dp2_3d = 0.0 + !$ser verbatim pn2_3d = 0.0 + !$ser verbatim pk2_3d = pk + + remap_t = .false. + remap_pt = .false. + remap_te = .false. + select case (remap_option) + case(0) + remap_t = .true. + case(1) + remap_pt = .true. + case(2) + remap_te = .true. + case default + print*, ' INVALID REMAP_OPTION ' + stop + end select + + select case (gmao_remap) + case(0) + ! use GFDL schemes + case(1) + ! GMAO linear remap + case(2) + ! GMAO quadratic remap + case(3) + ! GMAO cubic remap + case default + print*, ' INVALID GMAO_REMAP' + stop + end select + + if (is_master() .and. flagstruct%fv_debug) then + print*, '' + select case (remap_option) + case(0) + print*, ' REMAPPING T in logP' + case(1) + print*, ' REMAPPING PT in P' + case(2) + print*, ' REMAPPING TE in logP' + end select + + print*, '' + select case (gmao_remap) + case(0) + print*, ' Using GFDL schemes' + case(1) + print*, ' Using GMAO linear scheme' + case(2) + print*, ' Using GMAO quadratic scheme' + case(3) + print*, ' Using GMAO cubic scheme' + end select + + ! Total eergy conservation + print*, '' + print*, ' REMAPPING CONSV: ', consv + print*, ' REMAPPING CONSV_MIN: ', consv_min + print*, '' + endif + + k1k = rdgas/cv_air ! akap / (1.-akap) = rg/Cv=0.4 + rg = rdgas + rcp = 1./ cp + rrg = -rdgas/grav + +#ifdef MAPL_MODE + select case(nwat) + case(1) + liq_wat = -1 + ice_wat = -1 + rainwat = -1 + snowwat = -1 + graupel = -1 + cld_amt = -1 + case(3) + liq_wat = 2 + ice_wat = 3 + rainwat = -1 + snowwat = -1 + graupel = -1 + cld_amt = -1 + case(6:7) + liq_wat = 2 + ice_wat = 3 + rainwat = 4 + snowwat = 5 + graupel = 6 + cld_amt = 7 + end select +#else + liq_wat = get_tracer_index (MODEL_ATMOS, 'liq_wat') + ice_wat = get_tracer_index (MODEL_ATMOS, 'ice_wat') + rainwat = get_tracer_index (MODEL_ATMOS, 'rainwat') + snowwat = get_tracer_index (MODEL_ATMOS, 'snowwat') + graupel = get_tracer_index (MODEL_ATMOS, 'graupel') + cld_amt = get_tracer_index (MODEL_ATMOS, 'cld_amt') + !$ser verbatim o3mr = get_tracer_index (MODEL_ATMOS, 'o3mr') + !$ser verbatim sgs_tke = get_tracer_index (MODEL_ATMOS, 'sgs_tke') +#endif + + if ( do_sat_adj .and. nwat>=6 ) then + print*,'CODE NOT TESTED HERE 1' + fast_mp_consv = (.not.do_adiabatic_init) .and. consv>consv_min + do k=1,km + kmp = k + if ( pfull(k) > 10.E2 ) exit + enddo + call qs_init(kmp) + endif + +!$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & +!$OMP remap_pt,remap_te,mfy,mfx,cx,cy,hydrostatic, & +!$OMP pt,pk,rg,peln,q,nwat,liq_wat,rainwat,ice_wat,snowwat, & +#ifdef SERIALIZE +!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, cld_amt, mode,qmin, abskord,iep1, iedp1, jedp1, js2d, o3mr, sgs_tke, & +#endif +!$OMP graupel,q_con, sphum,cappa,r_vir,rcp,cp,k1k,delp, & +!$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & +!$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & +!$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,pe_3d, pe1_3d, pe2_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, peln_3d) & +!$OMP private(gz,cvm,bkh,dp2, & +!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_) + do 1000 j=js,je+1 + +!$ser verbatim if(j == js) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-In +!$ser data pe_=pe_3d pe1_=pe1_3d pe2_=pe2_3d ptop=ptop +!$ser data qvapor_=q(:,j,:,sphum) +!$ser data qliquid_=q(:,j,:,liq_wat) +!$ser data qice_=q(:,j,:,ice_wat) +!$ser data qrain_=q(:,j,:,rainwat) +!$ser data qsnow_=q(:,j,:,snowwat) +!$ser data qgraupel_=q(:,j,:,graupel) +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data r_vir=r_vir +!$ser data remap_t=remap_t +!$ser data ps=ps +!$ser data peln_3d=peln_3d +!$ser data pn2_3d=pn2_3d +!$ser data ak=ak +!$ser data bk=bk +!$ser data dp2_3d=dp2_3d +!$ser data pk=pk +!$ser data akap=akap +!$ser verbatim endif + + do k=1,km+1 + do i=is,ie + pe1(i,k) = pe(i,k,j) + enddo + enddo + + do i=is,ie + pe2(i, 1) = ptop + pe2(i,km+1) = pe(i,km+1,j) + enddo + + if ( j /= (je+1) ) then + + if (remap_t) then + ! Remap T in logP + ! Note: pt at this stage is Theta_v + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 2' + ! Transform virtual pt to virtual Temp + do k=1,km + do i=is,ie + pt(i,j,k) = pt(i,j,k)*(pk(i,j,k+1)-pk(i,j,k))/(akap*(peln(i,k+1,j)-peln(i,k,j))) + enddo + enddo + else + ! Transform "density pt" to "density temp" + !$ser verbatim if(j == js2d) then + !$ser savepoint MoistCVPlusPt_2d-In + !$ser data qvapor_js=q(:,j,:,sphum) qliquid_js=q(:,j,:,liq_wat) qice_js=q(:,j,:,ice_wat) qrain_js=q(:,j,:,rainwat) qsnow_js=q(:,j,:,snowwat) qgraupel_js=q(:,j,:,graupel) qcld_js=q(:,j,:,cld_amt) gz1d=gz cvm=cvm r_vir=r_vir cappa=cappa rrg=rrg delp=delp delz=delz pt=pt k1k=k1k j_2d=js2d q_con=q_con + !$ser verbatim endif + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + !$ser verbatim q_con(i,j,k) = gz(i) + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pt(i,j,k) = pt(i,j,k)*exp(cappa(i,j,k)/(1.-cappa(i,j,k))*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#else + do i=is,ie + pt(i,j,k) = pt(i,j,k)*exp(k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#endif + enddo + !$ser verbatim if(j == js2d) then + !$ser savepoint MoistCVPlusPt_2d-Out + !$ser data gz1d=gz cvm=cvm pt=pt cappa=cappa q_con=q_con + !$ser verbatim endif + endif ! hydro test + elseif (remap_pt) then + print*,'CODE NOT TESTED HERE 3' + ! Remap PT in P + ! pt is already virtual PT + elseif (remap_te) then + print*,'CODE NOT TESTED HERE 4' + ! Remap TE in logP + ! Transform virtual pt to total energy + if ( hydrostatic ) then + call pkez(km, is, ie, js, je, j, pe, pk, akap, peln, pkz, ptop) + do i=is,ie + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) + cp_air*pt(i,j,k)*(pk(i,j,k+1)-pk(i,j,k)) + enddo + enddo + do k=1,km+1 + do i=is,ie + phis(i,k) = phis(i,k) * pe(i,k,j) + enddo + enddo + ! Compute cp*T + KE + do k=1,km + do i=is,ie + te(i,j,k) = 0.25*gridstruct%rsin2(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + cp_air*pt(i,j,k)*pkz(i,j,k) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) + enddo + enddo + else + ! TE using 3D winds (pt is virtual potential temperature): + do i=is,ie + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do k=1,km+1 + do i=is,ie + phis(i,k) = phis(i,k) * pe(i,k,j) + enddo + enddo + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pkz(i,j,k) = exp(cappa(i,j,k)/(1.-cappa(i,j,k))*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + ! TE = KE + Cv*T_v + PE + te(i,j,k) = 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + cvm(i)*pt(i,j,k)*pkz(i,j,k) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) + enddo +#else + do i=is,ie + pkz(i,j,k) = exp(k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + ! TE = KE + Cv*T_v + PE + te(i,j,k) = 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + cv_air*pt(i,j,k)*pkz(i,j,k) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) + enddo +#endif + enddo + endif ! hydro test + endif + + ! update ps + do i=is,ie + ps(i,j) = pe1(i,km+1) + enddo + ! + ! Hybrid sigma-P coordinate: + ! + do k=2,km + do i=is,ie + pe2(i,k) = ak(k) + bk(k)*pe(i,km+1,j) + enddo + enddo + do k=1,km + do i=is,ie + dp2(i,k) = pe2(i,k+1) - pe2(i,k) + enddo + enddo + !------------------ + ! Compute p**Kappa + !------------------ + do k=1,km+1 + do i=is,ie + pk1(i,k) = pk(i,j,k) + pn1(i,k) = peln(i,k,j) + enddo + enddo + + do i=is,ie + pn2(i, 1) = peln(i, 1,j) + pn2(i,km+1) = peln(i,km+1,j) + pk2(i, 1) = pk1(i, 1) + pk2(i,km+1) = pk1(i,km+1) + enddo + + do k=2,km + do i=is,ie + pn2(i,k) = log(pe2(i,k)) + pk2(i,k) = exp(akap*pn2(i,k)) + enddo + enddo + +!$ser verbatim if(j == js) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe1_=pe1_3d pe2_=pe2_3d +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data dp2_3d=dp2_3d +!$ser data pn2_3d=pn2_3d +!$ser data pk=pk2_3d +!$ser verbatim endif + + if (remap_te) then + print*,'CODE NOT TESTED HERE 5' + !---------------------------------- + ! map TE in log P + !---------------------------------- + if ( gmao_remap > 0 ) then + call map1_gmao (km, pe1, te, & + km, pe2, te, & + is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.true.) + else + !$ser verbatim if(j == js2d) then + !$ser savepoint MapScalar_2d-In + !$ser verbatim mode=1 + !$ser data peln=peln pt=pt gz1d=gz pn2=pn2 mode=mode j_2d=js2d + !$ser verbatim endif + call map_scalar(km, pn1, te, gz, & + km, pn2, te, & + is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), te_min) + !$ser verbatim if(j == js2d) then + !$ser savepoint MapScalar_2d-Out + !$ser data pt=pt + !$ser verbatim endif + endif + else + !---------------------------------- + ! map T or PT in log P + !---------------------------------- + if ( gmao_remap > 0 ) then + print*,'CODE NOT TESTED HERE 6' + call map1_gmao (km, pe1, pt, & + km, pe2, pt, & + is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.false.) + else + !$ser verbatim if(j == js2d) then + !$ser savepoint MapScalar_2d-In + !$ser verbatim mode=1 + !$ser data peln=peln pt=pt gz1d=gz pn2=pn2 mode=mode j_2d=js2d kord_tm=kord_tm + !$ser verbatim endif + call map_scalar(km, pn1, pt, gz, & + km, pn2, pt, & + is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), t_min) + !$ser verbatim if(j == js2d) then + !$ser savepoint MapScalar_2d-Out + !$ser data pt=pt + !$ser verbatim endif + endif + endif + + !---------------- + ! Map constituents + !---------------- + if( nq > 5 ) then + !$ser verbatim if(j == js2d) then + !$ser savepoint MapN_Tracer_2d-In + !$ser data j_2d=js2d nq=nq qtracers=q(:,:,:,1:nq) pe1=pe1 pe2=pe2 dp2=dp2 q_min=qmin + !$ser verbatim endif + call mapn_tracer(nq, km, pe1, pe2, q, dp2, kord_tr, j, & + is, ie, isd, ied, jsd, jed, 0., fill) + !$ser verbatim if(j == js2d) then + !$ser savepoint MapN_Tracer_2d-Out + !$ser data qtracers=q(:,:,:,1:nq) + !$ser verbatim endif + elseif ( nq > 0 ) then + print*,'CODE NOT TESTED HERE 7' + ! Remap one tracer at a time + do iq=1,nq + call map1_q2(km, pe1, q(isd,jsd,1,iq), & + km, pe2, q2, dp2, & + is, ie, 0, kord_tr(iq), j, isd, ied, jsd, jed, 0.) + if (fill) call fillz(ie-is+1, km, 1, q2, dp2) + do k=1,km + do i=is,ie + q(i,j,k,iq) = q2(i,k) + enddo + enddo + enddo + endif + + if ( .not. hydrostatic ) then + ! Remap vertical wind: + call map1_ppm (km, pe1, w, ws(is,j), & + km, pe2, w, & + is, ie, j, isd, ied, jsd, jed, -2, kord_wz) + ! Remap delz for hybrid sigma-p coordinate + do k=1,km + do i=is,ie + delz(i,j,k) = -delz(i,j,k) / delp(i,j,k) ! ="specific volume"/grav + enddo + enddo + call map1_ppm (km, pe1, delz, gz, & + km, pe2, delz, & + is, ie, j, isd, ied, jsd, jed, 1, abs(kord_wz)) + do k=1,km + do i=is,ie + delz(i,j,k) = -delz(i,j,k)*dp2(i,k) + enddo + enddo + + !Fix excessive w - momentum conserving --- sjl + ! gz(:) used here as a temporary array + if ( w_limiter ) then +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,:,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim w_max_ = w_max +!$ser verbatim w_min_ = w_min + +!$ser savepoint W_fix_consrv_moment-In +!$ser data w=w w2=w2_3d dp2_W=dp2_3d w_max=w_max_ w_min=w_min_ +!$ser verbatim endif + do k=1,km + do i=is,ie + w2(i,k) = w(i,j,k) + enddo + enddo + do k=1, km-1 + do i=is,ie + if ( w2(i,k) > w_max ) then + if(j == js2d) then + print*,"ENTERED LOOP 1" + endif + gz(i) = (w2(i,k)-w_max) * dp2(i,k) + w2(i,k ) = w_max + w2(i,k+1) = w2(i,k+1) + gz(i)/dp2(i,k+1) + !print*, ' W_LIMITER down: ', i,j,k, w2(i,k:k+1), w(i,j,k:k+1) + elseif ( w2(i,k) < w_min ) then + if(j == js2d) then + print*,"ENTERED LOOP 2" + endif + gz(i) = (w2(i,k)-w_min) * dp2(i,k) + w2(i,k ) = w_min + w2(i,k+1) = w2(i,k+1) + gz(i)/dp2(i,k+1) + !print*, ' W_LIMITER down: ', i,j,k, w2(i,k:k+1), w(i,j,k:k+1) + endif + enddo + enddo + do k=km, 2, -1 + do i=is,ie + if ( w2(i,k) > w_max ) then + if(j == js2d) then + print*,"ENTERED LOOP 3" + endif + gz(i) = (w2(i,k)-w_max) * dp2(i,k) + w2(i,k ) = w_max + w2(i,k-1) = w2(i,k-1) + gz(i)/dp2(i,k-1) + !print*, ' W_LIMITER up: ', i,j,k, w2(i,k-1:k), w(i,j,k-1:k) + elseif ( w2(i,k) < w_min ) then + if(j == js2d) then + print*,"ENTERED LOOP 4" + endif + gz(i) = (w2(i,k)-w_min) * dp2(i,k) + w2(i,k ) = w_min + w2(i,k-1) = w2(i,k-1) + gz(i)/dp2(i,k-1) + !print*, ' W_LIMITER up: ', i,j,k, w2(i,k-1:k), w(i,j,k-1:k) + endif + enddo + enddo + do i=is,ie + if (w2(i,1) > w_max*2. ) then + if(j == js2d) then + print*,'ENTERED LOOP 5' + endif + w2(i,1) = w_max*2 ! sink out of the top of the domain + !print*, ' W_LIMITER top limited: ', i,j,1, w2(i,1), w(i,j,1) + elseif (w2(i,1) < w_min*2. ) then + if(j == js2d) then + print*,'ENTERED LOOP 6' + endif + w2(i,1) = w_min*2. + !print*, ' W_LIMITER top limited: ', i,j,1, w2(i,1), w(i,j,1) + endif + enddo + do k=1,km + do i=is,ie + w(i,j,k) = w2(i,k) + enddo + enddo +!$ser verbatim if(j == js2d) then +!$ser savepoint W_fix_consrv_moment-Out +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim w2_3d(i,:,k) = w2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser data w=w w2=w2_3d dp2_W=dp2_3d +!$ser verbatim endif + endif + endif + + endif !(j < je+1) + + !------ + ! map u + !------ + do i=is,ie+1 + pe0(i,1) = ptop + enddo + + do k=2,km+1 + do i=is,ie + pe0(i,k) = 0.5*(pe(i,k,j-1)+pe(i,k,j)) + enddo + enddo + + do k=1,km+1 + bkh = 0.5*bk(k) + do i=is,ie + pe3(i,k) = ak(k) + bkh*(pe(i,km+1,j-1)+pe(i,km+1,j)) + enddo + enddo + + call map1_ppm( km, pe0(is:ie,:), u, gz, & + km, pe3(is:ie,:), u, & + is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) + if (present(mfy)) then + ! print*, 'mfy present' + call map1_ppm( km, pe0(is:ie,:), mfy, gz, & + km, pe3(is:ie,:), mfy, & + is, ie, j, is, ie, js, je+1, -1, kord_mt) + endif + if (present(cy)) then + ! print*, 'cy present' + call map1_ppm( km, pe0(is:ie,:), cy, gz, & + km, pe3(is:ie,:), cy, & + is, ie, j, isd, ied, js, je+1, -1, kord_mt) + endif + + !------ + ! map v + !------ + if (j < je+1) then + do k=2,km+1 + do i=is,ie+1 + pe0(i,k) = 0.5*(pe(i-1,k, j)+pe(i,k, j)) + enddo + enddo + + do k=1,km+1 + bkh = 0.5*bk(k) + do i=is,ie+1 + pe3(i,k) = ak(k) + bkh*(pe(i-1,km+1,j)+pe(i,km+1,j)) + enddo + enddo + + call map1_ppm (km, pe0, v, gz, & + km, pe3, v, is, ie+1, & + j, isd, ied+1, jsd, jed, -1, kord_mt) + if (present(mfx)) then + ! print*, 'mfx present' + call map1_ppm (km, pe0, mfx, gz, & + km, pe3, mfx, is, ie+1, & + j, is, ie+1, js, je, -1, kord_mt) + endif + if (present(cx)) then + ! print*, 'cx present' + call map1_ppm (km, pe0, cx, gz, & + km, pe3, cx, is, ie+1, & + j, is, ie+1, jsd, jed, -1, kord_mt) + endif + endif ! (j < je+1) + +1000 continue + +! Update pressure variables and get new pkz, T_v, and omega + +!$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & +!$OMP remap_pt,remap_te,mfy,mfx,cx,cy,hydrostatic, & +!$OMP pt,pk,rg,peln,q,nwat,liq_wat,rainwat,ice_wat,snowwat, & +!$OMP graupel,sphum,cappa,r_vir,rcp,cp,k1k,delp, & +!$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & +!$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, & +!$OMP hs,w,ws,kord_wz,do_omega,omga,rrg,kord_mt) & +!$OMP private(gz,cvm,kp,k_next,bkh,dp2, & +!$OMP pe2,pe3,pk2,pn2,phis,tpe,dlnp,tmp) + do 2000 j=js,je + + !---------- + ! Update pe + !---------- + do i=is,ie + pe2(i, 1) = ptop + pe2(i,km+1) = pe(i,km+1,j) + enddo + do k=2,km + do i=is,ie + pe2(i,k) = ak(k) + bk(k)*pe(i,km+1,j) + enddo + enddo + do k=1,km+1 + do i=is,ie + pe(i,k,j) = pe2(i,k) + enddo + enddo + + !---------- + ! Update pk + !---------- + do i=is,ie + pn2(i, 1) = peln(i, 1,j) + pn2(i,km+1) = peln(i,km+1,j) + pk2(i, 1) = pk(i,j, 1) + pk2(i,km+1) = pk(i,j,km+1) + enddo + do k=2,km + do i=is,ie + pn2(i,k) = log(pe2(i,k)) + pk2(i,k) = exp(akap*pn2(i,k)) + enddo + enddo + do k=1,km+1 + do i=is,ie + pk(i,j,k) = pk2(i,k) + enddo + enddo + + !------------ + ! update delp + !------------ + do k=1,km + do i=is,ie + delp(i,j,k) = pe2(i,k+1) - pe2(i,k) + enddo + enddo + + !------------ + ! update logP + !------------ + do k=1,km+1 + do i=is,ie + peln(i,k,j) = pn2(i,k) + enddo + enddo + + !--------------------- + ! Compute pkz and T_v + !--------------------- + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 8' + do k=1,km + do i=is,ie + pkz(i,j,k) = (pk(i,j,k+1)-pk(i,j,k))/(akap*(peln(i,k+1,j)-peln(i,k,j))) + enddo + enddo + if (.not.remap_t) then + if (remap_te) then + ! Get updated T_v (store in pt) + do i=is,ie + gz(i) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + tpe = te(i,j,k) - gz(i) - 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j) ) + dlnp = rg*(peln(i,k+1,j) - peln(i,k,j)) + tmp = tpe / (cp - pe(i,k,j)*dlnp/delp(i,j,k)) + pt(i,j,k) = tmp + gz(i) = gz(i) + dlnp*tmp + enddo + enddo ! end k-loop + else + ! Make pt T_v + do k=1,km + do i=is,ie + pt(i,j,k) = pt(i,j,k)*pkz(i,j,k) + enddo + enddo + endif + endif + else + if (remap_te) then + print*,'CODE NOT TESTED HERE 9' + ! Invert TE using 3D winds to get pt (virtual temperature) and pkz: + do i=is,ie + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do k=1,km+1 + do i=is,ie + phis(i,k) = phis(i,k) * pe(i,k,j) + enddo + enddo + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + tpe = te(i,j,k) - & + ( 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) ) + pt(i,j,k) = tpe/cvm(i) + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#else + do i=is,ie + tpe = te(i,j,k) - & + ( 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j))) ) + pt(i,j,k) = tpe/cv_air + pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#endif + enddo + endif + if (remap_t) then + ! print*,'CODE EXECUTED' +! Note: pt at this stage is T_v or T_m + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#else + do i=is,ie + pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#endif + enddo + endif + if (remap_pt) then + print*,'CODE NOT TESTED HERE 10' +! Note: pt at this stage is Theta_v + do k=1,km + do i=is,ie + pkz(i,j,k) = exp( k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k)) ) + ! Make pt T_v + pt(i,j,k) = pt(i,j,k)*pkz(i,j,k) + enddo + enddo + endif + endif + +! Interpolate omega/pe3 (defined at peln) to remapped cell center (dp2) + if ( do_omega ) then +! Copy omega field to pe3 + print*,'CODE NOT TESTED HERE 11' + do i=is,ie + pe3(i,1) = 0. + enddo + do k=2,km+1 + do i=is,ie + pe3(i,k) = omga(i,j,k-1) + enddo + enddo + do k=1,km + do i=is,ie + dp2(i,k) = 0.5*(peln(i,k,j) + peln(i,k+1,j)) + enddo + enddo + do i=is,ie + k_next = 1 + do n=1,km + kp = k_next + do k=kp,km + if( dp2(i,n) <= peln(i,k+1,j) .and. dp2(i,n) >= peln(i,k,j) ) then + omga(i,j,n) = pe3(i,k) + (pe3(i,k+1) - pe3(i,k)) * & + (dp2(i,n)-peln(i,k,j)) / (peln(i,k+1,j)-peln(i,k,j) ) + k_next = k + exit + endif + enddo + enddo + enddo + endif ! end do_omega + +2000 continue + +! Do total energy conservation and fast saturation adjustment as requested +! and fill new PT (Theta_V) for next k_split step or export dry T + +!$OMP parallel default(none) shared(is,ie,js,je,km,kmp,ptop,u,v,pe,isd,ied,jsd,jed,kord_mt, & +!$OMP remap_t,remap_pt,remap_te, & +!$OMP te_2d,te,delp,hydrostatic,hs,rg,pt,peln, adiabatic, & +!$OMP cp,delz,nwat,rainwat,liq_wat,ice_wat,snowwat, & +!$OMP graupel,q_con,r_vir,sphum,w,pk,pkz,last_step,consv, & +!$OMP do_adiabatic_init,zsum1,zsum0,te0_2d,domain, & +!$OMP ng,gridstruct,E_Flux,pdt,dtmp,reproduce_sum,q, & +!$OMP mdt,cld_amt,cappa,dtdt,out_dt,rrg,akap,do_sat_adj, & +!$OMP fast_mp_consv,kord_tm) & +!$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp) + + dtmp = 0. + if( last_step .and. (.not.do_adiabatic_init) ) then + ! NOTE : Code can enter here since do_adiabatic_init can be False + if ( consv > consv_min ) then + ! print*, "consv > consv_min = entered" +!$OMP do + do j=js,je + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 12' + do i=is,ie + gz(i) = hs(i,j) + do k=1,km + gz(i) = gz(i) + rg*pt(i,j,k)*(peln(i,k+1,j)-peln(i,k,j)) + enddo + enddo + do i=is,ie + te_2d(i,j) = pe(i,km+1,j)*hs(i,j) - pe(i,1,j)*gz(i) + enddo + do k=1,km + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cp*pt(i,j,k) + & + 0.25*gridstruct%rsin2(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j))) + enddo + enddo + else +! TE using 3D winds (pt is virtual temperature): + do i=is,ie + te_2d(i,j) = 0. + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cvm(i)*pt(i,j,k) + & + 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)))) + enddo +#else + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cv_air*pt(i,j,k) + & + 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)))) + enddo +#endif + enddo ! k-loop + endif ! end non-hydro + + do i=is,ie + te_2d(i,j) = te0_2d(i,j) - te_2d(i,j) + zsum1(i,j) = pkz(i,j,1)*delp(i,j,1) + enddo + do k=2,km + do i=is,ie + zsum1(i,j) = zsum1(i,j) + pkz(i,j,k)*delp(i,j,k) + enddo + enddo + if ( hydrostatic ) then + do i=is,ie + zsum0(i,j) = ptop*(pk(i,j,1)-pk(i,j,km+1)) + zsum1(i,j) + enddo + endif + + enddo ! j-loop + +!$OMP single + tesum = mpp_global_sum(domain, te_2d*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + E_Flux = DBLE(consv)*tesum / DBLE(grav*pdt*4.*pi*radius**2) ! unit: W/m**2 + ! Note pdt is "phys" time step + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 13' + zsum = mpp_global_sum(domain, zsum0*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = tesum / DBLE(cp*zsum) + else + zsum = mpp_global_sum(domain, zsum1*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = tesum / DBLE(cv_air*zsum) + endif +!$OMP end single + + elseif ( consv < -consv_min ) then + print*,'CODE NOT TESTED HERE 14' +!$OMP do + do j=js,je + do i=is,ie + zsum1(i,j) = pkz(i,j,1)*delp(i,j,1) + enddo + do k=2,km + do i=is,ie + zsum1(i,j) = zsum1(i,j) + pkz(i,j,k)*delp(i,j,k) + enddo + enddo + if ( hydrostatic ) then + do i=is,ie + zsum0(i,j) = ptop*(pk(i,j,1)-pk(i,j,km+1)) + zsum1(i,j) + enddo + endif + enddo + + E_Flux = consv +!$OMP single + if ( hydrostatic ) then + zsum = mpp_global_sum(domain, zsum0*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = E_Flux*(grav*pdt*4.*pi*radius**2) / (cp*zsum) + else + zsum = mpp_global_sum(domain, zsum1*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = E_Flux*(grav*pdt*4.*pi*radius**2) / (cv_air*zsum) + endif +!$OMP end single + endif ! end consv check + endif ! end last_step check + +! Note: pt at this stage is T_v + if ( do_sat_adj .and. nwat>=6 ) then + print*,'CODE NOT TESTED HERE 15' + call timing_on('sat_adj2') +!$OMP do + do k=kmp,km + do j=js,je + do i=is,ie + dpln(i,j) = peln(i,k+1,j) - peln(i,k,j) + enddo + enddo + call fv_sat_adj(abs(mdt), r_vir, is, ie, js, je, ng, hydrostatic, fast_mp_consv, & + te(isd,jsd,k), q(isd,jsd,k,sphum), q(isd,jsd,k,liq_wat), & + q(isd,jsd,k,ice_wat), q(isd,jsd,k,rainwat), & + q(isd,jsd,k,snowwat), q(isd,jsd,k,graupel), & + hs ,dpln, delz(isd:,jsd:,k), pt(isd,jsd,k), delp(isd,jsd,k), & + cappa(isd:,jsd:,k), gridstruct%area_64, dtdt(is:,js:,k), out_dt, last_step, q(isd,jsd,k,cld_amt)) + if ( .not. hydrostatic ) then + do j=js,je + do i=is,ie +#ifdef MOIST_CAPPA + pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) +#else + pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) +#endif + enddo + enddo + endif + enddo ! OpenMP k-loop + + if ( fast_mp_consv ) then +!$OMP do + do j=js,je + do i=is,ie + do k=kmp,km + te0_2d(i,j) = te0_2d(i,j) + te(i,j,k) + enddo + enddo + enddo + endif + call timing_off('sat_adj2') + endif ! do_sat_adj + + + if ( last_step .and. (.not. adiabatic) ) then + ! Output temperature if last_step + if ( .not. hydrostatic ) then + ! print*,'CODE EXECUTED' +!$OMP do + do k=1,km + do j=js,je +#ifdef USE_COND + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / ((1.+r_vir*q(i,j,k,sphum))*(1.-gz(i))) + enddo +#else + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / (1.+r_vir*q(i,j,k,sphum)) + enddo +#endif + enddo ! j-loop + enddo ! k-loop +#ifdef USE_COND +! Fill condensate output +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + q_con(i,j,k) = 0.0 + if (liq_wat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,liq_wat) + if (ice_wat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,ice_wat) + if (rainwat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,rainwat) + if (snowwat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,snowwat) + if (graupel > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,graupel) + enddo + enddo ! j-loop + enddo ! k-loop +#endif + else + print*,'CODE NOT TESTED HERE 16' +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / (1.+r_vir*q(i,j,k,sphum)) + enddo + enddo ! j-loop + enddo ! k-loop + endif + + elseif ( last_step .and. adiabatic ) then + print*,'CODE NOT TESTED HERE 17' +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) + enddo + enddo ! j-loop + enddo ! k-loop + + else + print*,'CODE NOT TESTED HERE 18' + ! Top of the loop expects PT to be Theta_V +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k)/pkz(i,j,k) + enddo + enddo + enddo + + endif +!$OMP end parallel + + end subroutine Lagrangian_to_Eulerian + +!>@brief The subroutine 'compute_total_energy' performs the FV3-consistent computation of the global total energy. +!>@details It includes the potential, internal (latent and sensible heat), kinetic terms. + subroutine compute_total_energy(is, ie, js, je, isd, ied, jsd, jed, km, & + u, v, w, delz, pt, delp, q, qc, pe, peln, hs, & + rsin2_l, cosa_s_l, & + r_vir, cp, rg, hlv, te_2d, ua, va, teq, & + moist_phys, nwat, sphum, liq_wat, rainwat, ice_wat, snowwat, graupel, hydrostatic, id_te) +!------------------------------------------------------ +! Compute vertically integrated total energy per column +!------------------------------------------------------ +! !INPUT PARAMETERS: + integer, intent(in):: km, is, ie, js, je, isd, ied, jsd, jed, id_te + integer, intent(in):: sphum, liq_wat, ice_wat, rainwat, snowwat, graupel, nwat + real, intent(inout), dimension(isd:ied,jsd:jed,km):: ua, va + real, intent(in), dimension(isd:ied,jsd:jed,km):: pt, delp + real, intent(in), dimension(isd:ied,jsd:jed,km,*):: q + real, intent(in), dimension(isd:ied,jsd:jed,km):: qc + real, intent(inout):: u(isd:ied, jsd:jed+1,km) + real, intent(inout):: v(isd:ied+1,jsd:jed, km) + real, intent(in):: w(isd:,jsd:,1:) !< vertical velocity (m/s) + real, intent(in):: delz(isd:,jsd:,1:) + real, intent(in):: hs(isd:ied,jsd:jed) !< surface geopotential + real, intent(in):: pe(is-1:ie+1,km+1,js-1:je+1) !< pressure at layer edges + real, intent(in):: peln(is:ie,km+1,js:je) !< log(pe) + real, intent(in):: cp, rg, r_vir, hlv + real, intent(in) :: rsin2_l(isd:ied, jsd:jed) + real, intent(in) :: cosa_s_l(isd:ied, jsd:jed) + logical, intent(in):: moist_phys, hydrostatic +!! Output: + real(kind=8), intent(out):: te_2d(is:ie,js:je) !< vertically integrated TE + real, intent(out):: teq(is:ie,js:je) !< Moist TE +!! Local + real, dimension(is:ie,km):: tv + real phiz(is:ie,km+1) + real cvm(is:ie), qd(is:ie) + integer i, j, k + +!---------------------- +! Output lat-lon winds: +!---------------------- +! call cubed_to_latlon(u, v, ua, va, dx, dy, rdxa, rdya, km, flagstruct%c2l_ord) + +!$OMP parallel do default(none) shared(is,ie,js,je,isd,ied,jsd,jed,km,hydrostatic,hs,pt,qc,rg,peln,te_2d, & +!$OMP pe,delp,cp,rsin2_l,u,v,cosa_s_l,delz,moist_phys,w, & +!$OMP q,nwat,liq_wat,rainwat,ice_wat,snowwat,graupel,sphum) & +!$OMP private(phiz, tv, cvm, qd) + do j=js,je + + if ( hydrostatic ) then + + do i=is,ie + phiz(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + tv(i,k) = pt(i,j,k)*(1.+qc(i,j,k)) + phiz(i,k) = phiz(i,k+1) + rg*tv(i,k)*(peln(i,k+1,j)-peln(i,k,j)) + enddo + enddo + + do i=is,ie + te_2d(i,j) = pe(i,km+1,j)*phiz(i,km+1) - pe(i,1,j)*phiz(i,1) + enddo + + do k=1,km + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cp*tv(i,k) + & + 0.25*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j))) + enddo + enddo + + else +!----------------- +! Non-hydrostatic: +!----------------- + do i=is,ie + phiz(i,km+1) = hs(i,j) + do k=km,1,-1 + phiz(i,k) = phiz(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do i=is,ie + te_2d(i,j) = 0. + enddo + if ( moist_phys ) then + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, qd, cvm) +#endif + do i=is,ie +#ifdef MOIST_CAPPA + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cvm(i)*pt(i,j,k)*(1.+qc(i,j,k))*(1.-qd(i)) + & +#else + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cv_air*pt(i,j,k)*(1.+qc(i,j,k)) + & +#endif + 0.5*(phiz(i,k)+phiz(i,k+1)+w(i,j,k)**2+0.5*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2-(u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j)))) + enddo + enddo + else + do k=1,km + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cv_air*pt(i,j,k) + & + 0.5*(phiz(i,k)+phiz(i,k+1)+w(i,j,k)**2+0.5*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2-(u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j)))) + enddo + enddo + endif + endif + enddo + +!------------------------------------- +! Diganostics computation for moist TE +!------------------------------------- + if( id_te>0 ) then +!$OMP parallel do default(none) shared(is,ie,js,je,teq,te_2d,moist_phys,km,hlv,sphum,q,delp) + do j=js,je + do i=is,ie + teq(i,j) = te_2d(i,j) + enddo + if ( moist_phys ) then + do k=1,km + do i=is,ie + teq(i,j) = teq(i,j) + hlv*q(i,j,k,sphum)*delp(i,j,k) + enddo + enddo + endif + enddo + endif + + end subroutine compute_total_energy + + + subroutine pkez(km, ifirst, ilast, jfirst, jlast, j, & + pe, pk, akap, peln, pkz, ptop) + +! INPUT PARAMETERS: + integer, intent(in):: km, j + integer, intent(in):: ifirst, ilast !< Latitude strip + integer, intent(in):: jfirst, jlast !< Latitude strip + real, intent(in):: akap + real, intent(in):: pe(ifirst-1:ilast+1,km+1,jfirst-1:jlast+1) + real, intent(in):: pk(ifirst:ilast,jfirst:jlast,km+1) + real, intent(IN):: ptop +! OUTPUT + real, intent(out):: pkz(ifirst:ilast,jfirst:jlast,km) + real, intent(inout):: peln(ifirst:ilast, km+1, jfirst:jlast) !< log (pe) +! Local + real pk2(ifirst:ilast, km+1) + real pek + real lnp + real ak1 + integer i, k + + ak1 = (akap + 1.) / akap + + pek = pk(ifirst,j,1) + do i=ifirst, ilast + pk2(i,1) = pek + enddo + + do k=2,km+1 + do i=ifirst, ilast +! peln(i,k,j) = log(pe(i,k,j)) + pk2(i,k) = pk(i,j,k) + enddo + enddo + +!---- GFDL modification + if( ptop < ptop_min ) then + do i=ifirst, ilast + peln(i,1,j) = peln(i,2,j) - ak1 + enddo + else + lnp = log( ptop ) + do i=ifirst, ilast + peln(i,1,j) = lnp + enddo + endif +!---- GFDL modification + + do k=1,km + do i=ifirst, ilast + pkz(i,j,k) = (pk2(i,k+1) - pk2(i,k) ) / & + (akap*(peln(i,k+1,j) - peln(i,k,j)) ) + enddo + enddo + + end subroutine pkez + + + + subroutine remap_z(km, pe1, q1, kn, pe2, q2, i1, i2, iv, kord) + +! INPUT PARAMETERS: + integer, intent(in) :: i1 !< Starting longitude + integer, intent(in) :: i2 !< Finishing longitude + integer, intent(in) :: kord !< Method order + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + integer, intent(in) :: iv + + real, intent(in) :: pe1(i1:i2,km+1) !< height at layer edges from model top to bottom surface + real, intent(in) :: pe2(i1:i2,kn+1) !< height at layer edges from model top to bottom surface + real, intent(in) :: q1(i1:i2,km) !< Field input + +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(i1:i2,kn) !< Field output + +! LOCAL VARIABLES: + real qs(i1:i2) + real dp1( i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, delp, esl + integer i, k, l, m, k0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) ! negative + q4(1,i,k) = q1(i,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then + call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + +! Mapping + do 3000 i=i1,i2 + k0 = 1 + do 555 k=1,kn + do 100 l=k0,km +! locate the top edge: pe2(i,k) + if(pe2(i,k) <= pe1(i,l) .and. pe2(i,k) >= pe1(i,l+1)) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) >= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) < pe1(i,m+1) ) then +! Whole layer.. + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + delp = pe2(i,k+1)-pe1(i,m) + esl = delp / dp1(i,m) + qsum = qsum + delp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif +100 continue +123 q2(i,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) +555 continue +3000 continue + + end subroutine remap_z + + subroutine map_scalar( km, pe1, q1, qs, & + kn, pe2, q2, i1, i2, & + j, ibeg, iend, jbeg, jend, iv, kord, q_min) +! iv=1 + integer, intent(in) :: i1 !< Starting longitude + integer, intent(in) :: i2 !< Finishing longitude + integer, intent(in) :: iv !< Mode: 0 == constituents 1 == temp 2 == remap temp with cs scheme + integer, intent(in) :: kord !< Method order + integer, intent(in) :: j !< Current latitude + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + real, intent(in) :: qs(i1:i2) !< bottom BC + real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) !< Field output + real, intent(in):: q_min + +! DESCRIPTION: +! IV = 0: constituents +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate +! LOCAL VARIABLES: + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + integer i, k, l, m, k0, jj + integer LM1,LP0,LP1 + +!$ser verbatim real q4_1_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_2_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_3_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_4_temp(i1:i2, i1:i2,km) +!$ser verbatim real dp1_temp(i1:i2, i1:i2,km) +!$ser verbatim real pe1_temp(i1:i2, i1:i2,km+1) +!$ser verbatim real pe2_temp(i1:i2, i1:i2,kn+1) + +!$ser verbatim real LM1_INDEX(i1:i2, i1:i2, km) +!$ser verbatim real LP0_INDEX(i1:i2, i1:i2, km) + +!$ser verbatim q4_1_temp = 0.0 +!$ser verbatim q4_2_temp = 0.0 +!$ser verbatim q4_3_temp = 0.0 +!$ser verbatim q4_4_temp = 0.0 +!$ser verbatim dp1_temp = 0.0 +!$ser verbatim pe1_temp = 0.0 +!$ser verbatim pe2_temp = 0.0 +!$ser verbatim LM1_INDEX = 0 +!$ser verbatim LP0_INDEX = 0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,j,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord > 7 ) then + call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + +! NOTE : q1 and q2 fields being passed into map_scalar are identical variables +! even though q1 is declared at INTENT(IN) and q2 is declared as INTENT(IN/OUT). +!$ser verbatim if(j == 1) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim do jj = i1, i2 +!$ser verbatim dp1_temp(i,jj,k) = dp1(i,k) +!$ser verbatim pe1_temp(i,jj,k) = pe1(i,k) +!$ser verbatim pe2_temp(i,jj,k) = pe2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do jj = i1, i2 +!$ser verbatim do i = i1, i2 +!$ser verbatim pe1_temp(i,jj,km+1) = pe1(i,km+1) +!$ser verbatim pe2_temp(i,jj,km+1) = pe2(i,km+1) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Lagrangian_Contribution-In +!$ser data q1=q1 pe1_=pe1_temp pe2_=pe2_temp q4_1=q4_1_temp q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp q4_4=q4_4_temp dp1_=dp1_temp +!$ser verbatim endif + +! Interpolate field onto target Pressures +! --------------------------------------- + do i=i1,i2 + k0 = 1 + do 555 k=1,kn + LM1 = 1 + LP0 = 1 + do while( LP0.le.km ) + if (pe1(i,LP0).lt.pe2(i,k)) then + LP0 = LP0+1 + else + exit + endif + enddo + LM1 = max(LP0-1,1) + LP0 = min(LP0, km) +! Extrapolate Linearly above first model level +! ---------------------------------------------------- + if( LM1.eq.1 .and. LP0.eq.1 ) then + q2(i,j,k) = q1(i,j,1) + ( q1(i,j,2)-q1(i,j,1) )*( pe2(i,k)-pe1(i,1) ) & + /( pe1(i,2)-pe1(i,1) ) +! Extrapolate Linearly below last model level +! --------------------------------------------------- + else if( LM1.eq.km .and. LP0.eq.km ) then + q2(i,j,k) = q1(i,j,km) + ( q1(i,j,km)-q1(i,j,km-1) )*( pe2(i,k )-pe1(i,km ) ) & + /( pe1(i,km)-pe1(i,km-1) ) +! Interpolate Linearly between levels 1 => 2 and km-1 => km +! ----------------------------------------------------------------- + else if( LM1.eq.1 .or. LP0.eq.km ) then + q2(i,j,k) = q1(i,j,LP0) + ( q1(i,j,LM1)-q1(i,j,LP0) )*( pe2(i,k )-pe1(i,LP0) ) & + /( pe1(i,LM1)-pe1(i,LP0) ) + else + do l=k0,km +! locate the top edge: pe2(i,k) + if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if( pe2(i,k+1) <= pe1(i,l+1) ) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,j,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if( pe2(i,k+1) > pe1(i,m+1) ) then +! Whole layer + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif + enddo +123 q2(i,j,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) + + endif +555 continue + enddo + +!$ser verbatim if(j == 1) then +!$ser savepoint Lagrangian_Contribution-Out +!$ser data q1=q2 +!$ser verbatim endif + + end subroutine map_scalar + + + subroutine map1_ppm( km, pe1, q1, qs, & + kn, pe2, q2, i1, i2, & + j, ibeg, iend, jbeg, jend, iv, kord) + integer, intent(in) :: i1 !< Starting longitude + integer, intent(in) :: i2 !< Finishing longitude + integer, intent(in) :: iv !< Mode: 0 == constituents 1 == ??? 2 == remap temp with cs scheme + integer, intent(in) :: kord !< Method order + integer, intent(in) :: j !< Current latitude + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + real, intent(in) :: qs(i1:i2) !< bottom BC + real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) !< Field output + +! DESCRIPTION: +! IV = 0: constituents +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate + +! LOCAL VARIABLES: + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + integer i, k, l, m, k0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,j,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then + call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + + do i=i1,i2 + k0 = 1 + do 555 k=1,kn + do l=k0,km +! locate the top edge: pe2(i,k) + if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if( pe2(i,k+1) <= pe1(i,l+1) ) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,j,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if( pe2(i,k+1) > pe1(i,m+1) ) then +! Whole layer + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif + enddo +123 q2(i,j,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) +555 continue + enddo + + end subroutine map1_ppm + + + subroutine mapn_tracer(nq, km, pe1, pe2, q1, dp2, kord, j, & + i1, i2, isd, ied, jsd, jed, q_min, fill) +! INPUT PARAMETERS: + integer, intent(in):: km !< vertical dimension + integer, intent(in):: j, nq, i1, i2 + integer, intent(in):: isd, ied, jsd, jed + integer, intent(in):: kord(nq) + real, intent(in):: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in):: pe2(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in):: dp2(i1:i2,km) + real, intent(in):: q_min + logical, intent(in):: fill + real, intent(inout):: q1(isd:ied,jsd:jed,km,nq) ! Field input +! LOCAL VARIABLES: + real:: q4(4,i1:i2,km,nq) + real:: q2(i1:i2,km,nq) !< Field output + real:: qsum(nq) + real:: dp1(i1:i2,km) + real:: qs(i1:i2) + real:: pl, pr, dp, esl, fac1, fac2 + integer:: i, k, l, m, k0, iq + + !$ser verbatim integer:: kord_iq, iv, im, js2d + !$ser verbatim js2d=jsd+3 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + enddo + enddo + + do iq=1,nq + do k=1,km + do i=i1,i2 + q4(1,i,k,iq) = q1(i,j,k,iq) + enddo + enddo + call scalar_profile( qs, q4(1,i1,1,iq), dp1, km, i1, i2, 0, kord(iq), q_min ) + enddo + +! Mapping + do 4000 i=i1,i2 + k0 = 1 + do 555 k=1,km + do 100 l=k0,km +! locate the top edge: pe2(i,k) + if(pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1)) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) <= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + fac1 = pr + pl + fac2 = r3*(pr*fac1 + pl*pl) + fac1 = 0.5*fac1 + do iq=1,nq + q2(i,k,iq) = q4(2,i,l,iq) + (q4(4,i,l,iq)+q4(3,i,l,iq)-q4(2,i,l,iq))*fac1 & + - q4(4,i,l,iq)*fac2 + enddo + k0 = l + goto 555 + else +! Fractional area... + dp = pe1(i,l+1) - pe2(i,k) + fac1 = 1. + pl + fac2 = r3*(1.+pl*fac1) + fac1 = 0.5*fac1 + do iq=1,nq + qsum(iq) = dp*(q4(2,i,l,iq) + (q4(4,i,l,iq)+ & + q4(3,i,l,iq) - q4(2,i,l,iq))*fac1 - q4(4,i,l,iq)*fac2) + enddo + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) > pe1(i,m+1) ) then + ! Whole layer.. + do iq=1,nq + qsum(iq) = qsum(iq) + dp1(i,m)*q4(1,i,m,iq) + enddo + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + fac1 = 0.5*esl + fac2 = 1.-r23*esl + do iq=1,nq + qsum(iq) = qsum(iq) + dp*( q4(2,i,m,iq) + fac1*( & + q4(3,i,m,iq)-q4(2,i,m,iq)+q4(4,i,m,iq)*fac2 ) ) + enddo + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif +100 continue +123 continue + do iq=1,nq + q2(i,k,iq) = qsum(iq) / dp2(i,k) + enddo +555 continue +4000 continue + + !$ser verbatim if(j == js2d ) then + !$ser verbatim im = i2-i1+1 + !$ser verbatim iv = 9 + !$ser savepoint Fillz-In + ! Note : Currently, serializing nq=nq and q2tracers=q2(:,:,1:nq) will not run the translate test + ! To successfully run the translate test, serialize nq=iv and q2tracers=q2(:,:,1:iv) + !$ser data im=im km=km nq=iv dp2=dp2 q2tracers=q2(:,:,1:iv) + !$ser verbatim endif + + if (fill) call fillz(i2-i1+1, km, nq, q2, dp2) + + !$ser verbatim if(j == js2d ) then + !$ser savepoint Fillz-Out + !$ser data q2tracers=q2(:,:,1:iv) + !$ser verbatim endif + + do iq=1,nq +! if (fill) call fillz(i2-i1+1, km, 1, q2(i1,1,iq), dp2) + do k=1,km + do i=i1,i2 + q1(i,j,k,iq) = q2(i,k,iq) + enddo + enddo + enddo + + end subroutine mapn_tracer + + + subroutine map1_q2(km, pe1, q1, & + kn, pe2, q2, dp2, & + i1, i2, iv, kord, j, & + ibeg, iend, jbeg, jend, q_min ) + + +! INPUT PARAMETERS: + integer, intent(in) :: j + integer, intent(in) :: i1, i2 + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: iv !< Mode: 0 == constituents 1 == ??? + integer, intent(in) :: kord + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + + real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input + real, intent(in) :: dp2(i1:i2,kn) + real, intent(in) :: q_min +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(i1:i2,kn) !< Field output +! LOCAL VARIABLES: + real qs(i1:i2) + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + + integer i, k, l, m, k0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,j,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then + call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + +! Mapping + do 5000 i=i1,i2 + k0 = 1 + do 555 k=1,kn + do 100 l=k0,km +! locate the top edge: pe2(i,k) + if(pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1)) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) <= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) > pe1(i,m+1) ) then + ! Whole layer.. + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif +100 continue +123 q2(i,k) = qsum / dp2(i,k) +555 continue +5000 continue + + end subroutine map1_q2 + + + + subroutine remap_2d(km, pe1, q1, & + kn, pe2, q2, & + i1, i2, iv, kord) + integer, intent(in):: i1, i2 + integer, intent(in):: iv !< Mode: 0 == constituents 1 ==others + integer, intent(in):: kord + integer, intent(in):: km !< Original vertical dimension + integer, intent(in):: kn !< Target vertical dimension + real, intent(in):: pe1(i1:i2,km+1) !< Pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in):: pe2(i1:i2,kn+1) !< Pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(i1:i2,km) !< Field input + real, intent(out):: q2(i1:i2,kn) !< Field output +! LOCAL VARIABLES: + real qs(i1:i2) + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + integer i, k, l, m, k0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then + call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + + do i=i1,i2 + k0 = 1 + do 555 k=1,kn +#ifdef OLD_TOP_EDGE + if( pe2(i,k+1) <= pe1(i,1) ) then +! Entire grid above old ptop + q2(i,k) = q4(2,i,1) + elseif( pe2(i,k) < pe1(i,1) .and. pe2(i,k+1)>pe1(i,1) ) then +! Partially above old ptop: + q2(i,k) = q1(i,1) +#else + if( pe2(i,k) <= pe1(i,1) ) then +! above old ptop: + q2(i,k) = q1(i,1) +#endif + else + do l=k0,km +! locate the top edge: pe2(i,k) + if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) <= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) > pe1(i,m+1) ) then + ! Whole layer.. + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif + enddo +123 q2(i,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) + endif +555 continue + enddo + + end subroutine remap_2d + + +!>@brief Optimized vertical profile reconstruction: +!> Latest: Apr 2008 S.-J. Lin, NOAA/GFDL + subroutine scalar_profile(qs, a4, delp, km, i1, i2, iv, kord, qmin) +! Optimized vertical profile reconstruction: +! Latest: Apr 2008 S.-J. Lin, NOAA/GFDL + integer, intent(in):: i1, i2 + integer, intent(in):: km !< vertical dimension + integer, intent(in):: iv !< iv =-1: winds iv = 0: positive definite scalars iv = 1: others + integer, intent(in):: kord + real, intent(in) :: qs(i1:i2) + real, intent(in) :: delp(i1:i2,km) !< Layer pressure thickness + real, intent(inout):: a4(4,i1:i2,km) !< Interpolated values + real, intent(in):: qmin +!----------------------------------------------------------------------- + logical, dimension(i1:i2,km):: extm, ext5, ext6 + real gam(i1:i2,km) + real q(i1:i2,km+1) + real d4(i1:i2) + real bet, a_bot, grat + real pmp_1, lac_1, pmp_2, lac_2, x0, x1 + integer i, k, im + + if ( iv .eq. -2 ) then + do i=i1,i2 + gam(i,2) = 0.5 + q(i,1) = 1.5*a4(1,i,1) + enddo + do k=2,km-1 + do i=i1, i2 + grat = delp(i,k-1) / delp(i,k) + bet = 2. + grat + grat - gam(i,k) + q(i,k) = (3.*(a4(1,i,k-1)+a4(1,i,k)) - q(i,k-1))/bet + gam(i,k+1) = grat / bet + enddo + enddo + do i=i1,i2 + grat = delp(i,km-1) / delp(i,km) + q(i,km) = (3.*(a4(1,i,km-1)+a4(1,i,km)) - grat*qs(i) - q(i,km-1)) / & + (2. + grat + grat - gam(i,km)) + q(i,km+1) = qs(i) + enddo + do k=km-1,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k+1)*q(i,k+1) + enddo + enddo + else + do i=i1,i2 + grat = delp(i,2) / delp(i,1) ! grid ratio + bet = grat*(grat+0.5) + q(i,1) = ( (grat+grat)*(grat+1.)*a4(1,i,1) + a4(1,i,2) ) / bet + gam(i,1) = ( 1. + grat*(grat+1.5) ) / bet + enddo + + do k=2,km + do i=i1,i2 + d4(i) = delp(i,k-1) / delp(i,k) + bet = 2. + d4(i) + d4(i) - gam(i,k-1) + q(i,k) = ( 3.*(a4(1,i,k-1)+d4(i)*a4(1,i,k)) - q(i,k-1) )/bet + gam(i,k) = d4(i) / bet + enddo + enddo + + do i=i1,i2 + a_bot = 1. + d4(i)*(d4(i)+1.5) + q(i,km+1) = (2.*d4(i)*(d4(i)+1.)*a4(1,i,km)+a4(1,i,km-1)-a_bot*q(i,km)) & + / ( d4(i)*(d4(i)+0.5) - a_bot*gam(i,km) ) + enddo + + do k=km,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k)*q(i,k+1) + enddo + enddo + endif + +!----- Perfectly linear scheme -------------------------------- + if ( abs(kord) > 16 ) then + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + enddo + return + endif +!----- Perfectly linear scheme -------------------------------- + +!------------------ +! Apply constraints +!------------------ + im = i2 - i1 + 1 + +! Apply *large-scale* constraints + do i=i1,i2 + q(i,2) = min( q(i,2), max(a4(1,i,1), a4(1,i,2)) ) + q(i,2) = max( q(i,2), min(a4(1,i,1), a4(1,i,2)) ) + enddo + + do k=2,km + do i=i1,i2 + gam(i,k) = a4(1,i,k) - a4(1,i,k-1) + enddo + enddo + +! Interior: + do k=3,km-1 + do i=i1,i2 + if ( gam(i,k-1)*gam(i,k+1)>0. ) then +! Apply large-scale constraint to ALL fields if not local max/min + q(i,k) = min( q(i,k), max(a4(1,i,k-1),a4(1,i,k)) ) + q(i,k) = max( q(i,k), min(a4(1,i,k-1),a4(1,i,k)) ) + else + if ( gam(i,k-1) > 0. ) then +! There exists a local max + q(i,k) = max(q(i,k), min(a4(1,i,k-1),a4(1,i,k))) + else +! There exists a local min + q(i,k) = min(q(i,k), max(a4(1,i,k-1),a4(1,i,k))) + if ( iv==0 ) q(i,k) = max(0., q(i,k)) + endif + endif + enddo + enddo + +! Bottom: + do i=i1,i2 + q(i,km) = min( q(i,km), max(a4(1,i,km-1), a4(1,i,km)) ) + q(i,km) = max( q(i,km), min(a4(1,i,km-1), a4(1,i,km)) ) + enddo + + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + enddo + enddo + + do k=1,km + if ( k==1 .or. k==km ) then + do i=i1,i2 + extm(i,k) = (a4(2,i,k)-a4(1,i,k)) * (a4(3,i,k)-a4(1,i,k)) > 0. + enddo + else + do i=i1,i2 + extm(i,k) = gam(i,k)*gam(i,k+1) < 0. + enddo + endif + if ( abs(kord) > 9 ) then + do i=i1,i2 + x0 = 2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k)) + x1 = abs(a4(2,i,k)-a4(3,i,k)) + a4(4,i,k) = 3.*x0 + ext5(i,k) = abs(x0) > x1 + ext6(i,k) = abs(a4(4,i,k)) > x1 + enddo + endif + enddo + +!--------------------------- +! Apply subgrid constraints: +!--------------------------- +! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) +! Top 2 and bottom 2 layers always use monotonic mapping + + if ( iv==0 ) then + do i=i1,i2 + a4(2,i,1) = max(0., a4(2,i,1)) + enddo + elseif ( iv==-1 ) then + do i=i1,i2 + if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. + enddo + elseif ( iv==2 ) then + do i=i1,i2 + a4(2,i,1) = a4(1,i,1) + a4(3,i,1) = a4(1,i,1) + a4(4,i,1) = 0. + enddo + endif + + if ( iv/=2 ) then + do i=i1,i2 + a4(4,i,1) = 3.*(2.*a4(1,i,1) - (a4(2,i,1)+a4(3,i,1))) + enddo + call cs_limiters(im, extm(i1,1), a4(1,i1,1), 1) + endif + +! k=2 + do i=i1,i2 + a4(4,i,2) = 3.*(2.*a4(1,i,2) - (a4(2,i,2)+a4(3,i,2))) + enddo + call cs_limiters(im, extm(i1,2), a4(1,i1,2), 2) + +!------------------------------------- +! Huynh's 2nd constraint for interior: +!------------------------------------- + do k=3,km-2 + if ( abs(kord)<9 ) then + do i=i1,i2 +! Left edges + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) +! Right edges + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==9 ) then + do i=i1,i2 + if ( extm(i,k) .and. extm(i,k-1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else if ( extm(i,k) .and. extm(i,k+1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else if ( extm(i,k) .and. a4(1,i,k) abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + endif + endif + enddo + elseif ( abs(kord)==10 ) then + do i=i1,i2 + if( ext5(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + elseif( ext6(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==12 ) then + do i=i1,i2 + if( extm(i,k) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else ! not a local extremum + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) +! Check within the smooth region if subgrid profile is non-monotonic + if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) + endif + endif + enddo + elseif ( abs(kord)==13 ) then + do i=i1,i2 + if( ext6(i,k) ) then + if ( ext6(i,k-1) .and. ext6(i,k+1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==14 ) then + + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==15 ) then ! Revised abs(kord)=9 scheme + do i=i1,i2 + if ( ext5(i,k) .and. ext5(i,k-1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + else if ( ext5(i,k) .and. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + else if ( ext5(i,k) .and. a4(1,i,k)@brief The subroutine 'cs_profile' performs the optimized vertical profile reconstruction: +!>@date April 2008 +!>@author S. J. Lin, NOAA/GFDL + subroutine cs_profile(qs, a4, delp, km, i1, i2, iv, kord) +! Optimized vertical profile reconstruction: +! Latest: Apr 2008 S.-J. Lin, NOAA/GFDL + integer, intent(in):: i1, i2 + integer, intent(in):: km !< vertical dimension + integer, intent(in):: iv !< iv =-1: winds + !< iv = 0: positive definite scalars + !< iv = 1: others + integer, intent(in):: kord + real, intent(in) :: qs(i1:i2) + real, intent(in) :: delp(i1:i2,km) !< layer pressure thickness + real, intent(inout):: a4(4,i1:i2,km) !< Interpolated values +!----------------------------------------------------------------------- + logical, dimension(i1:i2,km):: extm, ext5, ext6 + real gam(i1:i2,km) + real q(i1:i2,km+1) + real d4(i1:i2) + real bet, a_bot, grat + real pmp_1, lac_1, pmp_2, lac_2, x0, x1 + integer i, k, im + + if ( iv .eq. -2 ) then + do i=i1,i2 + gam(i,2) = 0.5 + q(i,1) = 1.5*a4(1,i,1) + enddo + do k=2,km-1 + do i=i1, i2 + grat = delp(i,k-1) / delp(i,k) + bet = 2. + grat + grat - gam(i,k) + q(i,k) = (3.*(a4(1,i,k-1)+a4(1,i,k)) - q(i,k-1))/bet + gam(i,k+1) = grat / bet + enddo + enddo + do i=i1,i2 + grat = delp(i,km-1) / delp(i,km) + q(i,km) = (3.*(a4(1,i,km-1)+a4(1,i,km)) - grat*qs(i) - q(i,km-1)) / & + (2. + grat + grat - gam(i,km)) + q(i,km+1) = qs(i) + enddo + do k=km-1,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k+1)*q(i,k+1) + enddo + enddo + else + do i=i1,i2 + grat = delp(i,2) / delp(i,1) ! grid ratio + bet = grat*(grat+0.5) + q(i,1) = ( (grat+grat)*(grat+1.)*a4(1,i,1) + a4(1,i,2) ) / bet + gam(i,1) = ( 1. + grat*(grat+1.5) ) / bet + enddo + + do k=2,km + do i=i1,i2 + d4(i) = delp(i,k-1) / delp(i,k) + bet = 2. + d4(i) + d4(i) - gam(i,k-1) + q(i,k) = ( 3.*(a4(1,i,k-1)+d4(i)*a4(1,i,k)) - q(i,k-1) )/bet + gam(i,k) = d4(i) / bet + enddo + enddo + + do i=i1,i2 + a_bot = 1. + d4(i)*(d4(i)+1.5) + q(i,km+1) = (2.*d4(i)*(d4(i)+1.)*a4(1,i,km)+a4(1,i,km-1)-a_bot*q(i,km)) & + / ( d4(i)*(d4(i)+0.5) - a_bot*gam(i,km) ) + enddo + + do k=km,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k)*q(i,k+1) + enddo + enddo + endif +!----- Perfectly linear scheme -------------------------------- + if ( abs(kord) > 16 ) then + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + enddo + return + endif +!----- Perfectly linear scheme -------------------------------- + +!------------------ +! Apply constraints +!------------------ + im = i2 - i1 + 1 + +! Apply *large-scale* constraints + do i=i1,i2 + q(i,2) = min( q(i,2), max(a4(1,i,1), a4(1,i,2)) ) + q(i,2) = max( q(i,2), min(a4(1,i,1), a4(1,i,2)) ) + enddo + + do k=2,km + do i=i1,i2 + gam(i,k) = a4(1,i,k) - a4(1,i,k-1) + enddo + enddo + +! Interior: + do k=3,km-1 + do i=i1,i2 + if ( gam(i,k-1)*gam(i,k+1)>0. ) then +! Apply large-scale constraint to ALL fields if not local max/min + q(i,k) = min( q(i,k), max(a4(1,i,k-1),a4(1,i,k)) ) + q(i,k) = max( q(i,k), min(a4(1,i,k-1),a4(1,i,k)) ) + else + if ( gam(i,k-1) > 0. ) then +! There exists a local max + q(i,k) = max(q(i,k), min(a4(1,i,k-1),a4(1,i,k))) + else +! There exists a local min + q(i,k) = min(q(i,k), max(a4(1,i,k-1),a4(1,i,k))) + if ( iv==0 ) q(i,k) = max(0., q(i,k)) + endif + endif + enddo + enddo + +! Bottom: + do i=i1,i2 + q(i,km) = min( q(i,km), max(a4(1,i,km-1), a4(1,i,km)) ) + q(i,km) = max( q(i,km), min(a4(1,i,km-1), a4(1,i,km)) ) + enddo + + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + enddo + enddo + + do k=1,km + if ( k==1 .or. k==km ) then + do i=i1,i2 + extm(i,k) = (a4(2,i,k)-a4(1,i,k)) * (a4(3,i,k)-a4(1,i,k)) > 0. + enddo + else + do i=i1,i2 + extm(i,k) = gam(i,k)*gam(i,k+1) < 0. + enddo + endif + if ( abs(kord) > 9 ) then + do i=i1,i2 + x0 = 2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k)) + x1 = abs(a4(2,i,k)-a4(3,i,k)) + a4(4,i,k) = 3.*x0 + ext5(i,k) = abs(x0) > x1 + ext6(i,k) = abs(a4(4,i,k)) > x1 + enddo + endif + enddo + +!--------------------------- +! Apply subgrid constraints: +!--------------------------- +! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) +! Top 2 and bottom 2 layers always use monotonic mapping + + if ( iv==0 ) then + do i=i1,i2 + a4(2,i,1) = max(0., a4(2,i,1)) + enddo + elseif ( iv==-1 ) then + do i=i1,i2 + if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. + enddo + elseif ( iv==2 ) then + do i=i1,i2 + a4(2,i,1) = a4(1,i,1) + a4(3,i,1) = a4(1,i,1) + a4(4,i,1) = 0. + enddo + endif + + if ( iv/=2 ) then + do i=i1,i2 + a4(4,i,1) = 3.*(2.*a4(1,i,1) - (a4(2,i,1)+a4(3,i,1))) + enddo + call cs_limiters(im, extm(i1,1), a4(1,i1,1), 1) + endif + +! k=2 + do i=i1,i2 + a4(4,i,2) = 3.*(2.*a4(1,i,2) - (a4(2,i,2)+a4(3,i,2))) + enddo + call cs_limiters(im, extm(i1,2), a4(1,i1,2), 2) + +!------------------------------------- +! Huynh's 2nd constraint for interior: +!------------------------------------- + do k=3,km-2 + if ( abs(kord)<9 ) then + do i=i1,i2 +! Left edges + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) +! Right edges + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==9 ) then + do i=i1,i2 + if ( extm(i,k) .and. extm(i,k-1) ) then ! c90_mp122 +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else if ( extm(i,k) .and. extm(i,k+1) ) then ! c90_mp122 +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) +! Check within the smooth region if subgrid profile is non-monotonic + if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) + endif + endif + enddo + elseif ( abs(kord)==10 ) then + do i=i1,i2 + if( ext5(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + elseif( ext6(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==12 ) then + do i=i1,i2 + if( extm(i,k) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else ! not a local extremum + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) +! Check within the smooth region if subgrid profile is non-monotonic + if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) + endif + endif + enddo + elseif ( abs(kord)==13 ) then + do i=i1,i2 + if( ext6(i,k) ) then + if ( ext6(i,k-1) .and. ext6(i,k+1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==14 ) then + + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==15 ) then ! revised kord=9 scehem + do i=i1,i2 + if ( ext5(i,k) ) then ! c90_mp122 + if ( ext5(i,k-1) .or. ext5(i,k+1) ) then ! c90_mp122 +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + endif + elseif( ext6(i,k) ) then +! Check within the smooth region if subgrid profile is non-monotonic + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==16 ) then + do i=i1,i2 + if( ext5(i,k) ) then + if ( ext5(i,k-1) .or. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then + ! Left edges + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + ! Right edges + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + else ! kord = 11 + do i=i1,i2 + if ( ext5(i,k) .and. (ext5(i,k-1) .or. ext5(i,k+1)) ) then +! Noisy region: + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + endif + enddo + endif + +! Additional constraint to ensure positivity + if ( iv==0 ) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 0) + + enddo ! k-loop + +!---------------------------------- +! Bottom layer subgrid constraints: +!---------------------------------- + if ( iv==0 ) then + do i=i1,i2 + a4(3,i,km) = max(0., a4(3,i,km)) + enddo + elseif ( iv .eq. -1 ) then + do i=i1,i2 + if ( a4(3,i,km)*a4(1,i,km) <= 0. ) a4(3,i,km) = 0. + enddo + endif + + do k=km-1,km + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + if(k==(km-1)) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 2) + if(k== km ) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 1) + enddo + + end subroutine cs_profile + + + subroutine cs_limiters(im, extm, a4, iv) + integer, intent(in) :: im + integer, intent(in) :: iv + logical, intent(in) :: extm(im) + real , intent(inout) :: a4(4,im) !< PPM array +! LOCAL VARIABLES: + real da1, da2, a6da + integer i + + if ( iv==0 ) then +! Positive definite constraint + do i=1,im + if( a4(1,i)<=0.) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + if( abs(a4(3,i)-a4(2,i)) < -a4(4,i) ) then + if( (a4(1,i)+0.25*(a4(3,i)-a4(2,i))**2/a4(4,i)+a4(4,i)*r12) < 0. ) then +! local minimum is negative + if( a4(1,i) a4(2,i) ) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + else + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + endif + endif + enddo + elseif ( iv==1 ) then + do i=1,im + if( (a4(1,i)-a4(2,i))*(a4(1,i)-a4(3,i))>=0. ) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + da1 = a4(3,i) - a4(2,i) + da2 = da1**2 + a6da = a4(4,i)*da1 + if(a6da < -da2) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + elseif(a6da > da2) then + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + enddo + else +! Standard PPM constraint + do i=1,im + if( extm(i) ) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + da1 = a4(3,i) - a4(2,i) + da2 = da1**2 + a6da = a4(4,i)*da1 + if(a6da < -da2) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + elseif(a6da > da2) then + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + enddo + endif + end subroutine cs_limiters + + + + subroutine ppm_profile(a4, delp, km, i1, i2, iv, kord) + +! INPUT PARAMETERS: + integer, intent(in):: iv !< iv =-1: winds iv = 0: positive definite scalars iv = 1: others iv = 2: temp (if remap_t) and w (iv=-2) + integer, intent(in):: i1 !< Starting longitude + integer, intent(in):: i2 !< Finishing longitude + integer, intent(in):: km !< Vertical dimension + integer, intent(in):: kord !< Order (or more accurately method no.): + ! + real , intent(in):: delp(i1:i2,km) !< Layer pressure thickness + +! !INPUT/OUTPUT PARAMETERS: + real , intent(inout):: a4(4,i1:i2,km) !< Interpolated values + +! DESCRIPTION: +! +! Perform the piecewise parabolic reconstruction +! +! !REVISION HISTORY: +! S.-J. Lin revised at GFDL 2007 +!----------------------------------------------------------------------- +! local arrays: + real dc(i1:i2,km) + real h2(i1:i2,km) + real delq(i1:i2,km) + real df2(i1:i2,km) + real d4(i1:i2,km) + +! local scalars: + integer i, k, km1, lmt, it + real fac + real a1, a2, c1, c2, c3, d1, d2 + real qm, dq, lac, qmp, pmp + + km1 = km - 1 + it = i2 - i1 + 1 + + do k=2,km + do i=i1,i2 + delq(i,k-1) = a4(1,i,k) - a4(1,i,k-1) + d4(i,k ) = delp(i,k-1) + delp(i,k) + enddo + enddo + + do k=2,km1 + do i=i1,i2 + c1 = (delp(i,k-1)+0.5*delp(i,k))/d4(i,k+1) + c2 = (delp(i,k+1)+0.5*delp(i,k))/d4(i,k) + df2(i,k) = delp(i,k)*(c1*delq(i,k) + c2*delq(i,k-1)) / & + (d4(i,k)+delp(i,k+1)) + dc(i,k) = sign( min(abs(df2(i,k)), & + max(a4(1,i,k-1),a4(1,i,k),a4(1,i,k+1))-a4(1,i,k), & + a4(1,i,k)-min(a4(1,i,k-1),a4(1,i,k),a4(1,i,k+1))), df2(i,k) ) + enddo + enddo + +!----------------------------------------------------------- +! 4th order interpolation of the provisional cell edge value +!----------------------------------------------------------- + + do k=3,km1 + do i=i1,i2 + c1 = delq(i,k-1)*delp(i,k-1) / d4(i,k) + a1 = d4(i,k-1) / (d4(i,k) + delp(i,k-1)) + a2 = d4(i,k+1) / (d4(i,k) + delp(i,k)) + a4(2,i,k) = a4(1,i,k-1) + c1 + 2./(d4(i,k-1)+d4(i,k+1)) * & + ( delp(i,k)*(c1*(a1 - a2)+a2*dc(i,k-1)) - & + delp(i,k-1)*a1*dc(i,k ) ) + enddo + enddo + +! if(km>8 .and. kord>4) call steepz(i1, i2, km, a4, df2, dc, delq, delp, d4) + +! Area preserving cubic with 2nd deriv. = 0 at the boundaries +! Top + do i=i1,i2 + d1 = delp(i,1) + d2 = delp(i,2) + qm = (d2*a4(1,i,1)+d1*a4(1,i,2)) / (d1+d2) + dq = 2.*(a4(1,i,2)-a4(1,i,1)) / (d1+d2) + c1 = 4.*(a4(2,i,3)-qm-d2*dq) / ( d2*(2.*d2*d2+d1*(d2+3.*d1)) ) + c3 = dq - 0.5*c1*(d2*(5.*d1+d2)-3.*d1*d1) + a4(2,i,2) = qm - 0.25*c1*d1*d2*(d2+3.*d1) +! Top edge: +!------------------------------------------------------- + a4(2,i,1) = d1*(2.*c1*d1**2-c3) + a4(2,i,2) +!------------------------------------------------------- +! a4(2,i,1) = (12./7.)*a4(1,i,1)-(13./14.)*a4(1,i,2)+(3./14.)*a4(1,i,3) +!------------------------------------------------------- +! No over- and undershoot condition + a4(2,i,2) = max( a4(2,i,2), min(a4(1,i,1), a4(1,i,2)) ) + a4(2,i,2) = min( a4(2,i,2), max(a4(1,i,1), a4(1,i,2)) ) + dc(i,1) = 0.5*(a4(2,i,2) - a4(1,i,1)) + enddo + +! Enforce monotonicity within the top layer + + if( iv==0 ) then + do i=i1,i2 + a4(2,i,1) = max(0., a4(2,i,1)) + a4(2,i,2) = max(0., a4(2,i,2)) + enddo + elseif( iv==-1 ) then + do i=i1,i2 + if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. + enddo + elseif( abs(iv)==2 ) then + do i=i1,i2 + a4(2,i,1) = a4(1,i,1) + a4(3,i,1) = a4(1,i,1) + enddo + endif + +! Bottom +! Area preserving cubic with 2nd deriv. = 0 at the surface + do i=i1,i2 + d1 = delp(i,km) + d2 = delp(i,km1) + qm = (d2*a4(1,i,km)+d1*a4(1,i,km1)) / (d1+d2) + dq = 2.*(a4(1,i,km1)-a4(1,i,km)) / (d1+d2) + c1 = (a4(2,i,km1)-qm-d2*dq) / (d2*(2.*d2*d2+d1*(d2+3.*d1))) + c3 = dq - 2.0*c1*(d2*(5.*d1+d2)-3.*d1*d1) + a4(2,i,km) = qm - c1*d1*d2*(d2+3.*d1) +! Bottom edge: +!----------------------------------------------------- + a4(3,i,km) = d1*(8.*c1*d1**2-c3) + a4(2,i,km) +! dc(i,km) = 0.5*(a4(3,i,km) - a4(1,i,km)) +!----------------------------------------------------- +! a4(3,i,km) = (12./7.)*a4(1,i,km)-(13./14.)*a4(1,i,km-1)+(3./14.)*a4(1,i,km-2) +! No over- and under-shoot condition + a4(2,i,km) = max( a4(2,i,km), min(a4(1,i,km), a4(1,i,km1)) ) + a4(2,i,km) = min( a4(2,i,km), max(a4(1,i,km), a4(1,i,km1)) ) + dc(i,km) = 0.5*(a4(1,i,km) - a4(2,i,km)) + enddo + + +! Enforce constraint on the "slope" at the surface + +#ifdef BOT_MONO + do i=i1,i2 + a4(4,i,km) = 0 + if( a4(3,i,km) * a4(1,i,km) <= 0. ) a4(3,i,km) = 0. + d1 = a4(1,i,km) - a4(2,i,km) + d2 = a4(3,i,km) - a4(1,i,km) + if ( d1*d2 < 0. ) then + a4(2,i,km) = a4(1,i,km) + a4(3,i,km) = a4(1,i,km) + else + dq = sign(min(abs(d1),abs(d2),0.5*abs(delq(i,km-1))), d1) + a4(2,i,km) = a4(1,i,km) - dq + a4(3,i,km) = a4(1,i,km) + dq + endif + enddo +#else + if( iv==0 ) then + do i=i1,i2 + a4(2,i,km) = max(0.,a4(2,i,km)) + a4(3,i,km) = max(0.,a4(3,i,km)) + enddo + elseif( iv<0 ) then + do i=i1,i2 + if( a4(1,i,km)*a4(3,i,km) <= 0. ) a4(3,i,km) = 0. + enddo + endif +#endif + + do k=1,km1 + do i=i1,i2 + a4(3,i,k) = a4(2,i,k+1) + enddo + enddo + +!----------------------------------------------------------- +! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) +!----------------------------------------------------------- +! Top 2 and bottom 2 layers always use monotonic mapping + do k=1,2 + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 0) + enddo + + if(kord >= 7) then +!----------------------- +! Huynh's 2nd constraint +!----------------------- + do k=2,km1 + do i=i1,i2 +! Method#1 +! h2(i,k) = delq(i,k) - delq(i,k-1) +! Method#2 - better + h2(i,k) = 2.*(dc(i,k+1)/delp(i,k+1) - dc(i,k-1)/delp(i,k-1)) & + / ( delp(i,k)+0.5*(delp(i,k-1)+delp(i,k+1)) ) & + * delp(i,k)**2 +! Method#3 +!!! h2(i,k) = dc(i,k+1) - dc(i,k-1) + enddo + enddo + + fac = 1.5 ! original quasi-monotone + + do k=3,km-2 + do i=i1,i2 +! Right edges +! qmp = a4(1,i,k) + 2.0*delq(i,k-1) +! lac = a4(1,i,k) + fac*h2(i,k-1) + 0.5*delq(i,k-1) +! + pmp = 2.*dc(i,k) + qmp = a4(1,i,k) + pmp + lac = a4(1,i,k) + fac*h2(i,k-1) + dc(i,k) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), qmp, lac)), & + max(a4(1,i,k), qmp, lac) ) +! Left edges +! qmp = a4(1,i,k) - 2.0*delq(i,k) +! lac = a4(1,i,k) + fac*h2(i,k+1) - 0.5*delq(i,k) +! + qmp = a4(1,i,k) - pmp + lac = a4(1,i,k) + fac*h2(i,k+1) - dc(i,k) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), qmp, lac)), & + max(a4(1,i,k), qmp, lac)) +!------------- +! Recompute A6 +!------------- + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo +! Additional constraint to ensure positivity when kord=7 + if (iv == 0 .and. kord >= 6 ) & + call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 2) + enddo + + else + + lmt = kord - 3 + lmt = max(0, lmt) + if (iv == 0) lmt = min(2, lmt) + + do k=3,km-2 + if( kord /= 4) then + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + endif + if(kord/=6) call ppm_limiters(dc(i1,k), a4(1,i1,k), it, lmt) + enddo + endif + + do k=km1,km + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 0) + enddo + + end subroutine ppm_profile + + + subroutine ppm_limiters(dm, a4, itot, lmt) + +! INPUT PARAMETERS: + real , intent(in):: dm(*) !< Linear slope + integer, intent(in) :: itot !< Total Longitudes + integer, intent(in) :: lmt !< 0: Standard PPM constraint 1: Improved full monotonicity constraint + !< (Lin) 2: Positive definite constraint + !< 3: do nothing (return immediately) +! INPUT/OUTPUT PARAMETERS: + real , intent(inout) :: a4(4,*) !< PPM array AA <-- a4(1,i) AL <-- a4(2,i) AR <-- a4(3,i) A6 <-- a4(4,i) +! LOCAL VARIABLES: + real qmp + real da1, da2, a6da + real fmin + integer i + +! Developer: S.-J. Lin + + if ( lmt == 3 ) return + + if(lmt == 0) then +! Standard PPM constraint + do i=1,itot + if(dm(i) == 0.) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + da1 = a4(3,i) - a4(2,i) + da2 = da1**2 + a6da = a4(4,i)*da1 + if(a6da < -da2) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + elseif(a6da > da2) then + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + enddo + + elseif (lmt == 1) then + +! Improved full monotonicity constraint (Lin 2004) +! Note: no need to provide first guess of A6 <-- a4(4,i) + do i=1, itot + qmp = 2.*dm(i) + a4(2,i) = a4(1,i)-sign(min(abs(qmp),abs(a4(2,i)-a4(1,i))), qmp) + a4(3,i) = a4(1,i)+sign(min(abs(qmp),abs(a4(3,i)-a4(1,i))), qmp) + a4(4,i) = 3.*( 2.*a4(1,i) - (a4(2,i)+a4(3,i)) ) + enddo + + elseif (lmt == 2) then + +! Positive definite constraint + do i=1,itot + if( abs(a4(3,i)-a4(2,i)) < -a4(4,i) ) then + fmin = a4(1,i)+0.25*(a4(3,i)-a4(2,i))**2/a4(4,i)+a4(4,i)*r12 + if( fmin < 0. ) then + if(a4(1,i) a4(2,i)) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + else + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + endif + enddo + + endif + + end subroutine ppm_limiters + + + + subroutine steepz(i1, i2, km, a4, df2, dm, dq, dp, d4) + integer, intent(in) :: km, i1, i2 + real , intent(in) :: dp(i1:i2,km) !< Grid size + real , intent(in) :: dq(i1:i2,km) !< Backward diff of q + real , intent(in) :: d4(i1:i2,km) !< Backward sum: dp(k)+ dp(k-1) + real , intent(in) :: df2(i1:i2,km) !< First guess mismatch + real , intent(in) :: dm(i1:i2,km) !< Monotonic mismatch +! INPUT/OUTPUT PARAMETERS: + real , intent(inout) :: a4(4,i1:i2,km) !@brief The subroutine 'rst_remap' remaps all variables required for a restart. +!>@details npz_restart /= npz (i.e., when the number of vertical levels is +!! changed at restart) + subroutine rst_remap(km, kn, is,ie,js,je, isd,ied,jsd,jed, nq, ntp, & + delp_r, u_r, v_r, w_r, delz_r, pt_r, q_r, qdiag_r, & + delp, u, v, w, delz, pt, q, qdiag, & + ak_r, bk_r, ptop, ak, bk, hydrostatic, make_nh, & + domain, square_domain) +!------------------------------------ +! Assuming hybrid sigma-P coordinate: +!------------------------------------ +! INPUT PARAMETERS: + integer, intent(in):: km !< Restart z-dimension + integer, intent(in):: kn !< Run time dimension + integer, intent(in):: nq, ntp !< Number of tracers (including H2O) + integer, intent(in):: is,ie,isd,ied !< Starting & ending X-Dir index + integer, intent(in):: js,je,jsd,jed !< Starting & ending Y-Dir index + logical, intent(in):: hydrostatic, make_nh, square_domain + real, intent(IN) :: ptop + real, intent(in) :: ak_r(km+1) + real, intent(in) :: bk_r(km+1) + real, intent(in) :: ak(kn+1) + real, intent(in) :: bk(kn+1) + real, intent(in):: delp_r(is:ie,js:je,km) !< Pressure thickness + real, intent(in):: u_r(is:ie, js:je+1,km) !< u-wind (m/s) + real, intent(in):: v_r(is:ie+1,js:je ,km) !< v-wind (m/s) + real, intent(inout):: pt_r(is:ie,js:je,km) + real, intent(in):: w_r(is:ie,js:je,km) + real, intent(in):: q_r(is:ie,js:je,km,1:ntp) + real, intent(in):: qdiag_r(is:ie,js:je,km,ntp+1:nq) + real, intent(inout)::delz_r(is:ie,js:je,km) + type(domain2d), intent(INOUT) :: domain +! Output: + real, intent(out):: delp(isd:ied,jsd:jed,kn) !< Pressure thickness + real, intent(out):: u(isd:ied ,jsd:jed+1,kn) !< u-wind (m/s) + real, intent(out):: v(isd:ied+1,jsd:jed ,kn) !< v-wind (m/s) + real, intent(out):: w(isd: ,jsd: ,1:) !< Vertical velocity (m/s) + real, intent(out):: pt(isd:ied ,jsd:jed ,kn) !< Temperature + real, intent(out):: q(isd:ied,jsd:jed,kn,1:ntp) + real, intent(out):: qdiag(isd:ied,jsd:jed,kn,ntp+1:nq) + real, intent(out):: delz(isd:,jsd:,1:) !< Delta-height (m) +!----------------------------------------------------------------------- + real r_vir, rgrav + real ps(isd:ied,jsd:jed) !< Surface pressure + real pe1(is:ie,km+1) + real pe2(is:ie,kn+1) + real pv1(is:ie+1,km+1) + real pv2(is:ie+1,kn+1) + + integer i,j,k , iq + integer, parameter:: kord=4 + +#ifdef HYDRO_DELZ_REMAP + if (is_master() .and. .not. hydrostatic) then + print*, '' + print*, ' REMAPPING IC: INITIALIZING DELZ WITH HYDROSTATIC STATE ' + print*, '' + endif +#endif + +#ifdef HYDRO_DELZ_EXTRAP + if (is_master() .and. .not. hydrostatic) then + print*, '' + print*, ' REMAPPING IC: INITIALIZING DELZ WITH HYDROSTATIC STATE ABOVE INPUT MODEL TOP ' + print*, '' + endif +#endif + +#ifdef ZERO_W_EXTRAP + if (is_master() .and. .not. hydrostatic) then + print*, '' + print*, ' REMAPPING IC: INITIALIZING W TO ZERO ABOVE INPUT MODEL TOP ' + print*, '' + endif +#endif + + r_vir = rvgas/rdgas - 1. + rgrav = 1./grav + +!$OMP parallel do default(none) shared(is,ie,js,je,ps,ak_r) + do j=js,je + do i=is,ie + ps(i,j) = ak_r(1) + enddo + enddo + +! this OpenMP do-loop setup cannot work in it's current form.... +!$OMP parallel do default(none) shared(is,ie,js,je,km,ps,delp_r) + do j=js,je + do k=1,km + do i=is,ie + ps(i,j) = ps(i,j) + delp_r(i,j,k) + enddo + enddo + enddo + +! only one cell is needed + if ( square_domain ) then + call mpp_update_domains(ps, domain, whalo=1, ehalo=1, shalo=1, nhalo=1, complete=.true.) + else + call mpp_update_domains(ps, domain, complete=.true.) + endif + +! Compute virtual Temp +!$OMP parallel do default(none) shared(is,ie,js,je,km,pt_r,r_vir,q_r) + do k=1,km + do j=js,je + do i=is,ie + pt_r(i,j,k) = pt_r(i,j,k) * (1.+r_vir*q_r(i,j,k,1)) + enddo + enddo + enddo + +!$OMP parallel do default(none) shared(is,ie,js,je,km,ak_r,bk_r,ps,kn,ak,bk,u_r,u,delp, & +!$OMP ntp,nq,hydrostatic,make_nh,w_r,w,delz_r,delp_r,delz, & +!$OMP pt_r,pt,v_r,v,q,q_r,qdiag,qdiag_r) & +!$OMP private(pe1, pe2, pv1, pv2) + do 6000 j=js,je+1 +!------ +! map u +!------ + do k=1,km+1 + do i=is,ie + pe1(i,k) = ak_r(k) + 0.5*bk_r(k)*(ps(i,j-1)+ps(i,j)) + enddo + enddo + + do k=1,kn+1 + do i=is,ie + pe2(i,k) = ak(k) + 0.5*bk(k)*(ps(i,j-1)+ps(i,j)) + enddo + enddo + + call remap_2d(km, pe1, u_r(is:ie,j:j,1:km), & + kn, pe2, u(is:ie,j:j,1:kn), & + is, ie, -1, kord) + + if ( j /= (je+1) ) then + +!--------------- +! Hybrid sigma-p +!--------------- + do k=1,km+1 + do i=is,ie + pe1(i,k) = ak_r(k) + bk_r(k)*ps(i,j) + enddo + enddo + + do k=1,kn+1 + do i=is,ie + pe2(i,k) = ak(k) + bk(k)*ps(i,j) + enddo + enddo + +!------------- +! Compute delp +!------------- + do k=1,kn + do i=is,ie + delp(i,j,k) = pe2(i,k+1) - pe2(i,k) + enddo + enddo + +!---------------- +! Map constituents +!---------------- + if( nq /= 0 ) then + do iq=1,ntp + call remap_2d(km, pe1, q_r(is:ie,j:j,1:km,iq:iq), & + kn, pe2, q(is:ie,j:j,1:kn,iq:iq), & + is, ie, 0, kord) + enddo + do iq=ntp+1,nq + call remap_2d(km, pe1, qdiag_r(is:ie,j:j,1:km,iq:iq), & + kn, pe2, qdiag(is:ie,j:j,1:kn,iq:iq), & + is, ie, 0, kord) + enddo + endif + + if ( .not. hydrostatic .and. .not. make_nh) then +! Remap vertical wind: + call remap_2d(km, pe1, w_r(is:ie,j:j,1:km), & + kn, pe2, w(is:ie,j:j,1:kn), & + is, ie, -1, kord) + +#ifdef ZERO_W_EXTRAP + do k=1,kn + do i=is,ie + if (pe2(i,k) < pe1(i,1)) then + w(i,j,k) = 0. + endif + enddo + enddo +#endif + +#ifndef HYDRO_DELZ_REMAP +! Remap delz for hybrid sigma-p coordinate + do k=1,km + do i=is,ie + delz_r(i,j,k) = -delz_r(i,j,k)/delp_r(i,j,k) ! ="specific volume"/grav + enddo + enddo + call remap_2d(km, pe1, delz_r(is:ie,j:j,1:km), & + kn, pe2, delz(is:ie,j:j,1:kn), & + is, ie, 1, kord) + do k=1,kn + do i=is,ie + delz(i,j,k) = -delz(i,j,k)*delp(i,j,k) + enddo + enddo +#endif + endif + +! Geopotential conserving remap of virtual temperature: + do k=1,km+1 + do i=is,ie + pe1(i,k) = log(pe1(i,k)) + enddo + enddo + do k=1,kn+1 + do i=is,ie + pe2(i,k) = log(pe2(i,k)) + enddo + enddo + + call remap_2d(km, pe1, pt_r(is:ie,j:j,1:km), & + kn, pe2, pt(is:ie,j:j,1:kn), & + is, ie, 1, kord) + +#ifdef HYDRO_DELZ_REMAP + !initialize delz from the hydrostatic state + do k=1,kn + do i=is,ie + delz(i,j,k) = (rdgas*rgrav)*pt(i,j,k)*(pe2(i,k)-pe2(i,k+1)) + enddo + enddo +#endif +#ifdef HYDRO_DELZ_EXTRAP + !initialize delz from the hydrostatic state + do k=1,kn + do i=is,ie + if (pe2(i,k) < pe1(i,1)) then + delz(i,j,k) = (rdgas*rgrav)*pt(i,j,k)*(pe2(i,k)-pe2(i,k+1)) + endif + enddo + enddo +#endif +!------ +! map v +!------ + do k=1,km+1 + do i=is,ie+1 + pv1(i,k) = ak_r(k) + 0.5*bk_r(k)*(ps(i-1,j)+ps(i,j)) + enddo + enddo + do k=1,kn+1 + do i=is,ie+1 + pv2(i,k) = ak(k) + 0.5*bk(k)*(ps(i-1,j)+ps(i,j)) + enddo + enddo + + call remap_2d(km, pv1, v_r(is:ie+1,j:j,1:km), & + kn, pv2, v(is:ie+1,j:j,1:kn), & + is, ie+1, -1, kord) + + endif !(j < je+1) +6000 continue + +!$OMP parallel do default(none) shared(is,ie,js,je,kn,pt,r_vir,q) + do k=1,kn + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k) / (1.+r_vir*q(i,j,k,1)) + enddo + enddo + enddo + + end subroutine rst_remap + +!>@brief The subroutine 'mappm' is a general-purpose routine for remapping +!! one set of vertical levels to another. + subroutine mappm(km, pe1, q1, kn, pe2, q2, i1, i2, iv, kord, ptop) + +! IV = 0: constituents +! IV = 1: potential temp +! IV =-1: winds + +! Mass flux preserving mapping: q1(im,km) -> q2(im,kn) + +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate + + integer, intent(in):: i1, i2, km, kn, kord, iv + real, intent(in ):: pe1(i1:i2,km+1), pe2(i1:i2,kn+1) !< pe1: pressure at layer edges from model top to bottom + !! surface in the ORIGINAL vertical coordinate + !< pe2: pressure at layer edges from model top to bottom + !! surface in the NEW vertical coordinate +! Mass flux preserving mapping: q1(im,km) -> q2(im,kn) + real, intent(in ):: q1(i1:i2,km) + real, intent(out):: q2(i1:i2,kn) + real, intent(IN) :: ptop +! local + real qs(i1:i2) + real dp1(i1:i2,km) + real a4(4,i1:i2,km) + integer i, k, l + integer k0, k1 + real pl, pr, tt, delp, qsum, dpsum, esl + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + a4(1,i,k) = q1(i,k) + enddo + enddo + + if ( kord >7 ) then + call cs_profile( qs, a4, dp1, km, i1, i2, iv, kord ) + else + call ppm_profile( a4, dp1, km, i1, i2, iv, kord ) + endif + +!------------------------------------ +! Lowest layer: constant distribution +!------------------------------------ +#ifdef NGGPS_SUBMITTED + do i=i1,i2 + a4(2,i,km) = q1(i,km) + a4(3,i,km) = q1(i,km) + a4(4,i,km) = 0. + enddo +#endif + + do 5555 i=i1,i2 + k0 = 1 + do 555 k=1,kn + + if(pe2(i,k) .le. pe1(i,1)) then +! above old ptop + q2(i,k) = q1(i,1) + elseif(pe2(i,k) .ge. pe1(i,km+1)) then +! Entire grid below old ps +#ifdef NGGPS_SUBMITTED + q2(i,k) = a4(3,i,km) ! this is not good. +#else + q2(i,k) = q1(i,km) +#endif + else + + do 45 L=k0,km +! locate the top edge at pe2(i,k) + if( pe2(i,k) .ge. pe1(i,L) .and. & + pe2(i,k) .le. pe1(i,L+1) ) then + k0 = L + PL = (pe2(i,k)-pe1(i,L)) / dp1(i,L) + if(pe2(i,k+1) .le. pe1(i,L+1)) then + +! entire new grid is within the original grid + PR = (pe2(i,k+1)-pe1(i,L)) / dp1(i,L) + TT = r3*(PR*(PR+PL)+PL**2) + q2(i,k) = a4(2,i,L) + 0.5*(a4(4,i,L)+a4(3,i,L) & + - a4(2,i,L))*(PR+PL) - a4(4,i,L)*TT + goto 555 + else +! Fractional area... + delp = pe1(i,L+1) - pe2(i,k) + TT = r3*(1.+PL*(1.+PL)) + qsum = delp*(a4(2,i,L)+0.5*(a4(4,i,L)+ & + a4(3,i,L)-a4(2,i,L))*(1.+PL)-a4(4,i,L)*TT) + dpsum = delp + k1 = L + 1 + goto 111 + endif + endif +45 continue + +111 continue + do 55 L=k1,km + if( pe2(i,k+1) .gt. pe1(i,L+1) ) then + +! Whole layer.. + + qsum = qsum + dp1(i,L)*q1(i,L) + dpsum = dpsum + dp1(i,L) + else + delp = pe2(i,k+1)-pe1(i,L) + esl = delp / dp1(i,L) + qsum = qsum + delp * (a4(2,i,L)+0.5*esl* & + (a4(3,i,L)-a4(2,i,L)+a4(4,i,L)*(1.-r23*esl)) ) + dpsum = dpsum + delp + k0 = L + goto 123 + endif +55 continue + delp = pe2(i,k+1) - pe1(i,km+1) + if(delp > 0.) then +! Extended below old ps +#ifdef NGGPS_SUBMITTED + qsum = qsum + delp * a4(3,i,km) ! not good. +#else + qsum = qsum + delp * q1(i,km) +#endif + dpsum = dpsum + delp + endif +123 q2(i,k) = qsum / dpsum + endif +555 continue +5555 continue + + end subroutine mappm + +!>@brief The subroutine 'moist_cv' computes the FV3-consistent moist heat capacity under constant volume, +!! including the heating capacity of water vapor and condensates. +!>@details See \cite emanuel1994atmospheric for information on variable heat capacities. + subroutine moist_cv(is,ie, isd,ied, jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, qd, cvm, t1) + integer, intent(in):: is, ie, isd,ied, jsd,jed, km, nwat, j, k + integer, intent(in):: sphum, liq_wat, rainwat, ice_wat, snowwat, graupel + real, intent(in), dimension(isd:ied,jsd:jed,km,nwat):: q + real, intent(out), dimension(is:ie):: cvm, qd + real, intent(in), optional:: t1(is:ie) +! + real, parameter:: t_i0 = 15. + real, dimension(is:ie):: qv, ql, qs + integer:: i + + select case (nwat) + + case(1) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + cvm(i) = (1.-qv(i))*cv_air + qv(i)*cv_vap + enddo + case(2) + if ( present(t1) ) then ! Special case for GFS physics + do i=is,ie + qd(i) = max(0., q(i,j,k,liq_wat)) + if ( t1(i) > tice ) then + qs(i) = 0. + elseif ( t1(i) < tice-t_i0 ) then + qs(i) = qd(i) + else + qs(i) = qd(i)*(tice-t1(i))/t_i0 + endif + ql(i) = qd(i) - qs(i) + qv(i) = max(0.,q(i,j,k,sphum)) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + else + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + qs(i) = max(q(i,j,k,liq_wat),0.0) + qd(i) = qs(i) + cvm(i) = (1.-qv(i))*cv_air + qv(i)*cv_vap + enddo + endif + case (3) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + qs(i) = max(q(i,j,k,ice_wat),0.0) + qd(i) = ql(i) + qs(i) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + case(4) ! K_warm_rain with fake ice + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + qd(i)*c_liq + enddo + case(5) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) + qs(i) = max(q(i,j,k,ice_wat),0.0) + max(q(i,j,k,snowwat),0.0) + qd(i) = ql(i) + qs(i) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + case(6:7) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) + qs(i) = max(q(i,j,k,ice_wat),0.0) + max(q(i,j,k,snowwat),0.0) + max(q(i,j,k,graupel),0.0) + qd(i) = ql(i) + qs(i) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + case default + !call mpp_error (NOTE, 'fv_mapz::moist_cv - using default cv_air') + do i=is,ie + qd(i) = 0. + cvm(i) = cv_air + enddo + end select + + end subroutine moist_cv + +!>@brief The subroutine 'moist_cp' computes the FV3-consistent moist heat capacity under constant pressure, +!! including the heating capacity of water vapor and condensates. + subroutine moist_cp(is,ie, isd,ied, jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, qd, cpm, t1) + + integer, intent(in):: is, ie, isd,ied, jsd,jed, km, nwat, j, k + integer, intent(in):: sphum, liq_wat, rainwat, ice_wat, snowwat, graupel + real, intent(in), dimension(isd:ied,jsd:jed,km,nwat):: q + real, intent(out), dimension(is:ie):: cpm, qd + real, intent(in), optional:: t1(is:ie) +! + real, parameter:: t_i0 = 15. + real, dimension(is:ie):: qv, ql, qs + integer:: i + + select case (nwat) + + case(2) + if ( present(t1) ) then ! Special case for GFS physics + do i=is,ie + qd(i) = max(0., q(i,j,k,liq_wat)) + if ( t1(i) > tice ) then + qs(i) = 0. + elseif ( t1(i) < tice-t_i0 ) then + qs(i) = qd(i) + else + qs(i) = qd(i)*(tice-t1(i))/t_i0 + endif + ql(i) = qd(i) - qs(i) + qv(i) = max(0.,q(i,j,k,sphum)) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + else + do i=is,ie + qv(i) = max(0.,q(i,j,k,sphum)) + qs(i) = max(0.,q(i,j,k,liq_wat)) + qd(i) = qs(i) + cpm(i) = (1.-qv(i))*cp_air + qv(i)*cp_vapor + enddo + endif + + case(3) + do i=is,ie + qv(i) = q(i,j,k,sphum) + ql(i) = q(i,j,k,liq_wat) + qs(i) = q(i,j,k,ice_wat) + qd(i) = ql(i) + qs(i) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + case(4) ! K_warm_rain scheme with fake ice + do i=is,ie + qv(i) = q(i,j,k,sphum) + qd(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + qd(i)*c_liq + enddo + case(5) + do i=is,ie + qv(i) = q(i,j,k,sphum) + ql(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) + qs(i) = q(i,j,k,ice_wat) + q(i,j,k,snowwat) + qd(i) = ql(i) + qs(i) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + case(6:7) + do i=is,ie + qv(i) = q(i,j,k,sphum) + ql(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) + qs(i) = q(i,j,k,ice_wat) + q(i,j,k,snowwat) + q(i,j,k,graupel) + qd(i) = ql(i) + qs(i) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + case default + ! call mpp_error (NOTE, 'fv_mapz::moist_cp - using default cp_air') + do i=is,ie + qd(i) = 0. + cpm(i) = cp_air + enddo + end select + + end subroutine moist_cp + +!----------------------------------------------------------------------- +!BOP +! !ROUTINE: map1_gmao --- GMAO Interpolation for vertical re-mapping +! +! !INTERFACE: + subroutine map1_gmao( km, pe1, q1, & + kn, pe2, q2, i1, i2, & + j, ibeg, iend, jbeg, jend, akap, gmao_remap, P_MAP, conserv) + implicit none + +! !INPUT PARAMETERS: + integer, intent(in) :: i1 ! Starting longitude + integer, intent(in) :: i2 ! Finishing longitude + real, intent(in) :: akap + integer, intent(in) :: P_MAP ! Thermodynamic variable to remap + ! 1:TE-logP 2:PT-PK 3:T-logP + integer, intent(in) :: gmao_remap ! 3:cubic 2:quadratic 1:linear + logical, intent(in) :: conserv + integer, intent(in) :: j ! Current latitude + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: km ! Original vertical dimension + integer, intent(in) :: kn ! Target vertical dimension + + real, intent(in) :: pe1(i1:i2,km+1) ! pressure at layer edges + ! (from model top to bottom surface) + ! in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) ! pressure at layer edges + ! (from model top to bottom surface) + ! in the new vertical coordinate + + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) ! Field input +! !INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) ! Field output + +! !DESCRIPTION: +! +! Perform Cubic Interpolation a given latitude +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate +! +! !REVISION HISTORY: +! 2005.11.14 Takacs Initial Code +! 2016.07.20 Putman Modified to make genaric for any thermodynamic variable +! +!EOP +!----------------------------------------------------------------------- +!BOC +! +! !LOCAL VARIABLES: + real qx(i1:i2,km) + real logpl1(i1:i2,km) + real logpl2(i1:i2,kn) + real dlogp1(i1:i2,km) + real vsum1(i1:i2) + real vsum2(i1:i2) + real am2,am1,ap0,ap1,P,PLP1,PLP0,PLM1,PLM2,DLP0,DLM1,DLM2 + + integer i, k, LM2,LM1,LP0,LP1 + +! Initialization +! -------------- + + select case (P_MAP) + case(1) + ! Remapping in Log(P) + do k=1,km + qx(:,k) = q1(i1:i2,j,k) + logpl1(:,k) = log( r2*(pe1(:,k)+pe1(:,k+1)) ) + enddo + do k=1,kn + logpl2(:,k) = log( r2*(pe2(:,k)+pe2(:,k+1)) ) + enddo + do k=1,km-1 + dlogp1(:,k) = logpl1(:,k+1)-logpl1(:,k) + enddo + case(2) + ! Remapping in P^KAPPA + do k=1,km + qx(:,k) = q1(i1:i2,j,k) + logpl1(:,k) = exp( akap*log( r2*(pe1(:,k)+pe1(:,k+1))) ) + enddo + do k=1,kn + logpl2(:,k) = exp( akap*log( r2*(pe2(:,k)+pe2(:,k+1))) ) + enddo + do k=1,km-1 + dlogp1(:,k) = logpl1(:,k+1)-logpl1(:,k) + enddo + end select + + if (conserv) then +! Compute vertical integral of Input field +! ---------------------------------------- + vsum1(:) = r0 + do i=i1,i2 + do k=1,km + vsum1(i) = vsum1(i) + qx(i,k)*( pe1(i,k+1)-pe1(i,k) ) + enddo + vsum1(i) = vsum1(i) / ( pe1(i,km+1)-pe1(i,1) ) + enddo + endif + +! Interpolate field onto target Pressures +! --------------------------------------- + do i=i1,i2 + do k=1,kn + LM1 = 1 + LP0 = 1 + do while( LP0.le.km ) + if (logpl1(i,LP0).lt.logpl2(i,k)) then + LP0 = LP0+1 + else + exit + endif + enddo + LM1 = max(LP0-1,1) + LP0 = min(LP0, km) + +! Extrapolate Linearly above first model level +! ---------------------------------------------------- + if( LM1.eq.1 .and. LP0.eq.1 ) then + q2(i,j,k) = qx(i,1) + ( qx(i,2)-qx(i,1) )*( logpl2(i,k)-logpl1(i,1) ) & + /( logpl1(i,2)-logpl1(i,1) ) + +! Extrapolate Linearly below last model level +! --------------------------------------------------- + else if( LM1.eq.km .and. LP0.eq.km ) then + q2(i,j,k) = qx(i,km) + ( qx(i,km)-qx(i,km-1) )*( logpl2(i,k )-logpl1(i,km ) ) & + /( logpl1(i,km)-logpl1(i,km-1) ) + +! Interpolate Linearly between levels 1 => 2 and km-1 => km +! ----------------------------------------------------------------- + else if( LM1.eq.1 .or. LP0.eq.km ) then + q2(i,j,k) = qx(i,LP0) + ( qx(i,LM1)-qx(i,LP0) )*( logpl2(i,k )-logpl1(i,LP0) ) & + /( logpl1(i,LM1)-logpl1(i,LP0) ) +! Interpolate Cubicly between other model levels +! ------------------------------------------------------ + else + LP1 = LP0+1 + LM2 = LM1-1 + P = logpl2(i,k) + PLP1 = logpl1(i,LP1) + PLP0 = logpl1(i,LP0) + PLM1 = logpl1(i,LM1) + PLM2 = logpl1(i,LM2) + DLP0 = dlogp1(i,LP0) + DLM1 = dlogp1(i,LM1) + DLM2 = dlogp1(i,LM2) + + ! Cubic Coefficients + ! ------------------ + if( gmao_remap .eq. 3 ) then + ap1 = (P-PLP0)*(P-PLM1)*(P-PLM2)/( DLP0*(DLP0+DLM1)*(DLP0+DLM1+DLM2) ) + ap0 = (PLP1-P)*(P-PLM1)*(P-PLM2)/( DLP0* DLM1 *( DLM1+DLM2) ) + am1 = (PLP1-P)*(PLP0-P)*(P-PLM2)/( DLM1* DLM2 *(DLP0+DLM1 ) ) + am2 = (PLP1-P)*(PLP0-P)*(PLM1-P)/( DLM2*(DLM1+DLM2)*(DLP0+DLM1+DLM2) ) + q2(i,j,k) = ap1*qx(i,LP1) + ap0*qx(i,LP0) + am1*qx(i,LM1) + am2*qx(i,LM2) + endif + + ! Quadratic Coefficients + ! ---------------------- + if( gmao_remap .eq. 2 ) then + ap1 = (P-PLP0)*(P-PLM1)/( (PLP1-PLP0)*(PLP1-PLM1) ) + ap0 = (PLP1-P)*(P-PLM1)/( (PLP1-PLP0)*(PLP0-PLM1) ) + am1 = (PLP1-P)*(PLP0-P)/( (PLP1-PLM1)*(PLP0-PLM1) ) + q2(i,j,k) = ap1*qx(i,LP1) + ap0*qx(i,LP0) + am1*qx(i,LM1) + endif + + ! Linear Coefficients + ! ------------------- + if( gmao_remap .eq. 1 ) then + q2(i,j,k) = qx(i,LP0) + ( qx(i,LM1)-qx(i,LP0) )*( logpl2(i,k )-logpl1(i,LP0) ) & + /( logpl1(i,LM1)-logpl1(i,LP0) ) + endif + + endif + + enddo + enddo + + if (conserv) then +! Compute vertical integral of Output field +! ----------------------------------------- + vsum2(:) = r0 + do i=i1,i2 + do k=1,kn + vsum2(i) = vsum2(i) + q2(i,j,k)*( pe2(i,k+1)-pe2(i,k) ) + enddo + vsum2(i) = vsum2(i) / ( pe2(i,kn+1)-pe2(i,1) ) + enddo +! Adjust Final field to conserve +! ------------------------------ + do i=i1,i2 + do k=1,kn + q2(i,j,k) = q2(i,j,k) + vsum1(i)-vsum2(i) + enddo + enddo + endif + + return +!EOC + end subroutine map1_gmao +!----------------------------------------------------------------------- + + +end module fv_mapz_mod From 7835b6f9f24ce4f3b73cf1558f6971bfa01cf1ea Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 23 Oct 2024 15:47:13 -0700 Subject: [PATCH 013/252] Added translate test for scalar_profile subroutine under map_scalar. Using C24 dataset, the translate test verifies for all except the q4_4 data. The q4_4 data error seems to be due to some roundoff from a constant as there are about 10 to 20 unique differences over set of thousands --- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_scalar_profile.py | 117 ++++++++++++++++++ 2 files changed, 119 insertions(+), 1 deletion(-) create mode 100644 tests/savepoint/translate/translate_scalar_profile.py diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 5361fa00..8676b617 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -88,4 +88,5 @@ from .translate_w_fix_consrv_moment import TranslateW_fix_consrv_moment from .translate_lagrangian_contribution import TranslateLagrangian_Contribution -from .translate_remapping_GEOS import TranslateRemapping_GEOS \ No newline at end of file +from .translate_remapping_GEOS import TranslateRemapping_GEOS +from .translate_scalar_profile import TranslateScalar_Profile \ No newline at end of file diff --git a/tests/savepoint/translate/translate_scalar_profile.py b/tests/savepoint/translate/translate_scalar_profile.py new file mode 100644 index 00000000..d9a5af85 --- /dev/null +++ b/tests/savepoint/translate/translate_scalar_profile.py @@ -0,0 +1,117 @@ +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + +from ndsl import StencilFactory +from ndsl.dsl.typing import FloatField, Float +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.remap_profile import RemapProfile + + +class TranslateScalar_Profile(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "qs_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_1": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_2": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_3": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_4": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "dp1_":{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + + + } + self.in_vars["parameters"] = [ + "q_min", + ] + + self.out_vars = { + "q4_1": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_2": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_3": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_4": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + } + + # Value from GEOS + self.kord = 9 + + # mode / iv set to 1 from GEOS + self.mode = 1 + + self._compute_func = RemapProfile( + self.stencil_factory, + self.quantity_factory, + self.kord, + self.mode, + dims=[X_DIM, Y_DIM, Z_DIM] + ) + + def compute_from_storage(self, inputs): + self._compute_func( + inputs["qs_"], + inputs["q4_1"], + inputs["q4_2"], + inputs["q4_3"], + inputs["q4_4"], + inputs["dp1_"], + Float(inputs["q_min"]), + ) + return inputs From 832eed0876d24beaaf1fff760f039b47955be626 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 23 Oct 2024 15:48:16 -0700 Subject: [PATCH 014/252] Added updated serialized fv_mapz.F90 file --- fv_mapz.F90.SER | 43 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index e9e8f588..e547be40 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -1663,6 +1663,8 @@ contains !$ser verbatim real LM1_INDEX(i1:i2, i1:i2, km) !$ser verbatim real LP0_INDEX(i1:i2, i1:i2, km) + !$ser verbatim real :: qs_2d(i1:i2, i1:i2) + !$ser verbatim q4_1_temp = 0.0 !$ser verbatim q4_2_temp = 0.0 !$ser verbatim q4_3_temp = 0.0 @@ -1672,7 +1674,7 @@ contains !$ser verbatim pe2_temp = 0.0 !$ser verbatim LM1_INDEX = 0 !$ser verbatim LP0_INDEX = 0 - +!$ser verbatim qs_2d = 0.0 do k=1,km do i=i1,i2 dp1(i,k) = pe1(i,k+1) - pe1(i,k) @@ -1682,7 +1684,46 @@ contains ! Compute vertical subgrid distribution if ( kord > 7 ) then +!$ser verbatim if (j == 1) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do i = i1,i2 +!$ser verbatim qs_2d(i,j) = qs(i) +!$ser verbatim enddo +!$ser savepoint Scalar_Profile-In +!$ser data qs_=qs_2d +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser data dp1_=dp1_temp +!$ser data q_min=q_min +!$ser verbatim endif call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) +!$ser verbatim if (j == 1) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Scalar_Profile-Out +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser data dp1_=dp1_temp +!$ser verbatim endif else call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) endif From f325f8030e1268c928f3f0b62b3a6326733128ba Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 25 Oct 2024 09:29:56 -0700 Subject: [PATCH 015/252] Separated the lagrangian contribution code with interpolation into its own routine. --- pyFV3/stencils/map_single.py | 63 +++++++++- tests/savepoint/translate/__init__.py | 2 +- .../translate/translate_MapN_Tracer_2d.py | 117 ++++++++++++++++++ ...anslate_lagrangian_contribution_interp.py} | 34 ++--- 4 files changed, 199 insertions(+), 17 deletions(-) create mode 100644 tests/savepoint/translate/translate_MapN_Tracer_2d.py rename tests/savepoint/translate/{translate_lagrangian_contribution.py => translate_lagrangian_contribution_interp.py} (81%) diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index b6e055c8..3cbc1778 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -15,8 +15,69 @@ def set_dp(dp1: FloatField, pe1: FloatField, lev: IntFieldIJ): with computation(FORWARD), interval(0, 1): lev = 0 - def lagrangian_contributions( + q: FloatField, + pe1: FloatField, + pe2: FloatField, + q4_1: FloatField, + q4_2: FloatField, + q4_3: FloatField, + q4_4: FloatField, + dp1: FloatField, + lev: IntFieldIJ, +): + """ + Args: + q (out): + pe1 (in): + pe2 (in): + q4_1 (in): + q4_2 (in): + q4_3 (in): + q4_4 (in): + dp1 (in): + lev (inout): + """ + # TODO: Can we make lev a 2D temporary? + with computation(FORWARD), interval(...): + pl = (pe2 - pe1[0, 0, lev]) / dp1[0, 0, lev] + if pe2[0, 0, 1] <= pe1[0, 0, lev + 1]: + pr = (pe2[0, 0, 1] - pe1[0, 0, lev]) / dp1[0, 0, lev] + q = ( + q4_2[0, 0, lev] + + 0.5 + * (q4_4[0, 0, lev] + q4_3[0, 0, lev] - q4_2[0, 0, lev]) + * (pr + pl) + - q4_4[0, 0, lev] * 1.0 / 3.0 * (pr * (pr + pl) + pl * pl) + ) + else: + qsum = (pe1[0, 0, lev + 1] - pe2) * ( + q4_2[0, 0, lev] + + 0.5 + * (q4_4[0, 0, lev] + q4_3[0, 0, lev] - q4_2[0, 0, lev]) + * (1.0 + pl) + - q4_4[0, 0, lev] * 1.0 / 3.0 * (1.0 + pl * (1.0 + pl)) + ) + lev = lev + 1 + while pe1[0, 0, lev + 1] < pe2[0, 0, 1]: + qsum += dp1[0, 0, lev] * q4_1[0, 0, lev] + lev = lev + 1 + dp = pe2[0, 0, 1] - pe1[0, 0, lev] + esl = dp / dp1[0, 0, lev] + qsum += dp * ( + q4_2[0, 0, lev] + + 0.5 + * esl + * ( + q4_3[0, 0, lev] + - q4_2[0, 0, lev] + + q4_4[0, 0, lev] * (1.0 - (2.0 / 3.0) * esl) + ) + ) + q = qsum / (pe2[0, 0, 1] - pe2) + lev = lev - 1 + +def lagrangian_contributions_interp( km: int, not_exit_loop: BoolFieldIJ, INDEX_LM1: IntField, diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 8676b617..846d7ddb 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -87,6 +87,6 @@ from .translate_ytp_v import TranslateYTP_V from .translate_w_fix_consrv_moment import TranslateW_fix_consrv_moment -from .translate_lagrangian_contribution import TranslateLagrangian_Contribution +from .translate_lagrangian_contribution_interp import TranslateLagrangian_Contribution_Interp from .translate_remapping_GEOS import TranslateRemapping_GEOS from .translate_scalar_profile import TranslateScalar_Profile \ No newline at end of file diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py new file mode 100644 index 00000000..d9a5af85 --- /dev/null +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -0,0 +1,117 @@ +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + +from ndsl import StencilFactory +from ndsl.dsl.typing import FloatField, Float +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.remap_profile import RemapProfile + + +class TranslateScalar_Profile(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "qs_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_1": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_2": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_3": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_4": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "dp1_":{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + + + } + self.in_vars["parameters"] = [ + "q_min", + ] + + self.out_vars = { + "q4_1": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_2": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_3": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_4": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + } + + # Value from GEOS + self.kord = 9 + + # mode / iv set to 1 from GEOS + self.mode = 1 + + self._compute_func = RemapProfile( + self.stencil_factory, + self.quantity_factory, + self.kord, + self.mode, + dims=[X_DIM, Y_DIM, Z_DIM] + ) + + def compute_from_storage(self, inputs): + self._compute_func( + inputs["qs_"], + inputs["q4_1"], + inputs["q4_2"], + inputs["q4_3"], + inputs["q4_4"], + inputs["dp1_"], + Float(inputs["q_min"]), + ) + return inputs diff --git a/tests/savepoint/translate/translate_lagrangian_contribution.py b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py similarity index 81% rename from tests/savepoint/translate/translate_lagrangian_contribution.py rename to tests/savepoint/translate/translate_lagrangian_contribution_interp.py index 4013507b..41382a12 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py @@ -1,23 +1,27 @@ from ndsl import StencilFactory, QuantityFactory -from ndsl.dsl.typing import FloatField, FloatFieldIJ, Float, BoolFieldIJ, IntField, IntFieldIJ, Float +from ndsl.dsl.typing import FloatField, FloatFieldIJ, BoolFieldIJ, IntField, IntFieldIJ, Float, Int, Bool from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid from gt4py.cartesian.gtscript import PARALLEL, FORWARD, BACKWARD, computation, interval -from pyFV3.stencils.map_single import lagrangian_contributions +from pyFV3.stencils.map_single import lagrangian_contributions_interp -class test_Lagragian_Contribution: +class test_Lagragian_Contribution_Interp: def __init__( self, stencil_factory: StencilFactory, grid: Grid, ): - print("In test_Lagragian_Contribution") + print("In test_Lagragian_Contribution_interp") - self._lagrangian_contributions = stencil_factory.from_origin_domain( - func=lagrangian_contributions, - origin=(3,3,0), - domain=(24,1,72), + grid_indexing = stencil_factory.grid_indexing + + self._lagrangian_contributions_interp = stencil_factory.from_origin_domain( + func=lagrangian_contributions_interp, + # origin=(3,3,0), + origin=grid_indexing.origin_compute(), + # domain=(24,1,72), + domain=(grid.nic, 1, grid.npz) ) def __call__( @@ -36,7 +40,7 @@ def __call__( dp1: FloatField, lev: IntFieldIJ, ): - self._lagrangian_contributions( + self._lagrangian_contributions_interp( km, not_exit_loop, INDEX_LM1, @@ -52,12 +56,12 @@ def __call__( lev, ) -class TranslateLagrangian_Contribution(TranslateFortranData2Py): +class TranslateLagrangian_Contribution_Interp(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist, stencil_factory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid - self.compute_func = test_Lagragian_Contribution(self.stencil_factory, self.grid) # type: ignore + self.compute_func = test_Lagragian_Contribution_Interp(self.stencil_factory, self.grid) # type: ignore self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { @@ -127,25 +131,25 @@ def compute_from_storage(self, inputs): self._not_exit_loop = self.quantity_factory.zeros( [X_DIM, Y_DIM], units="", - dtype=bool, + dtype=Bool, ) self._INDEX_LM1 = self.quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="", - dtype=int, + dtype=Int, ) self._INDEX_LP0 = self.quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="", - dtype=int, + dtype=Int, ) self._lev = self.quantity_factory.zeros( [X_DIM, Y_DIM], units="", - dtype=int, + dtype=Int, ) self.compute_func( From b10b18fa64fa75d83a05ad0e453a3a84bebe03ea Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 29 Oct 2024 15:07:58 -0700 Subject: [PATCH 016/252] Added map_scalar translate test. Note that translate test for map_scalar only has solutions for j=1, and remaining j's will have nans for solution for q1 (or pt) --- pyFV3/stencils/map_single.py | 66 +++++++--- tests/savepoint/translate/__init__.py | 4 +- .../translate/translate_MapN_Tracer_2d.py | 115 +++++------------- .../translate/translate_map_scalar.py | 72 +++++++++++ 4 files changed, 152 insertions(+), 105 deletions(-) create mode 100644 tests/savepoint/translate/translate_map_scalar.py diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index 3cbc1778..3219494b 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -265,11 +265,24 @@ def make_quantity(): compute_dims=dims, ) - self._INDEX_LM1 = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=int) - self._INDEX_LP0 = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=int) + self._lagrangian_contributions_interp = stencil_factory.from_dims_halo( + lagrangian_contributions_interp, + compute_dims=dims, + ) + + self._INDEX_LM1 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="", + dtype=int, + ) + + self._INDEX_LP0 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="", + dtype=int, + ) self._km = grid_indexing.domain[2] self._not_exit_loop = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=bool) - # self._q_temp = make_quantity() @property def i_extent(self): @@ -286,6 +299,7 @@ def __call__( pe2: FloatField, qs: Optional["FloatFieldIJ"] = None, qmin: Float = 0.0, + interp: bool = False, ): """ Compute x-flux using the PPM method. @@ -321,20 +335,32 @@ def __call__( self._dp1, qmin, ) - self._lagrangian_contributions( - self._km, - self._not_exit_loop, - self._INDEX_LM1, - self._INDEX_LP0, - # self._q_temp, - q1, - pe1, - pe2, - self._q4_1, - self._q4_2, - self._q4_3, - self._q4_4, - self._dp1, - self._lev, - ) - return q1 + + if(interp == False): + self._lagrangian_contributions( + q1, + pe1, + pe2, + self._q4_1, + self._q4_2, + self._q4_3, + self._q4_4, + self._dp1, + self._lev, + ) + else: + self._lagrangian_contributions_interp( + self._km, + self._not_exit_loop, + self._INDEX_LM1, + self._INDEX_LP0, + q1, + pe1, + pe2, + self._q4_1, + self._q4_2, + self._q4_3, + self._q4_4, + self._dp1, + self._lev, + ) diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 846d7ddb..afc28135 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -89,4 +89,6 @@ from .translate_w_fix_consrv_moment import TranslateW_fix_consrv_moment from .translate_lagrangian_contribution_interp import TranslateLagrangian_Contribution_Interp from .translate_remapping_GEOS import TranslateRemapping_GEOS -from .translate_scalar_profile import TranslateScalar_Profile \ No newline at end of file +from .translate_scalar_profile import TranslateScalar_Profile +from .translate_MapN_Tracer_2d import TranslateMapN_Tracer_2d +from .translate_map_scalar import TranslateMap_Scalar \ No newline at end of file diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index d9a5af85..9d39aacd 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -4,10 +4,10 @@ from ndsl.dsl.typing import FloatField, Float from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils.remap_profile import RemapProfile +from pyFV3.stencils.mapn_tracer import MapNTracer -class TranslateScalar_Profile(TranslateFortranData2Py): +class TranslateMapN_Tracer_2d(TranslateFortranData2Py): def __init__(self, grid, namelist, stencil_factory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory @@ -15,78 +15,18 @@ def __init__(self, grid, namelist, stencil_factory): self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { - "qs_": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "q4_1": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "q4_2": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "q4_3": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "q4_4": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "dp1_":{ - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, + "qtracers": { }, + } - self.in_vars["parameters"] = [ - "q_min", - ] + # self.in_vars["parameters"] = [ + # "q_min", + # ] self.out_vars = { - "q4_1": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "q4_2": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "q4_3": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "q4_4": {"istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, + "qtracers": { }, } @@ -96,22 +36,29 @@ def __init__(self, grid, namelist, stencil_factory): # mode / iv set to 1 from GEOS self.mode = 1 - self._compute_func = RemapProfile( - self.stencil_factory, - self.quantity_factory, - self.kord, - self.mode, - dims=[X_DIM, Y_DIM, Z_DIM] - ) + nq = 9 + + fill = True + + # self._mapn_tracer = MapNTracer( + # self.stencil_factory, + # self.quantity_factory, + # abs(self.kord), + # nq, + # fill=fill, + # tracers=tracers, + # ) def compute_from_storage(self, inputs): - self._compute_func( - inputs["qs_"], - inputs["q4_1"], - inputs["q4_2"], - inputs["q4_3"], - inputs["q4_4"], - inputs["dp1_"], - Float(inputs["q_min"]), - ) + # self._compute_func( + # inputs["qs_"], + # inputs["q4_1"], + # inputs["q4_2"], + # inputs["q4_3"], + # inputs["q4_4"], + # inputs["dp1_"], + # Float(inputs["q_min"]), + # ) + + print("qtracers shape: ", inputs["qtracers"].data.shape) return inputs diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py new file mode 100644 index 00000000..3c5eb86e --- /dev/null +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -0,0 +1,72 @@ +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + +from ndsl import StencilFactory +from ndsl.dsl.typing import FloatField, Float +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.map_single import MapSingle + + +class TranslateMap_Scalar(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "q1": { + "kend": grid.npz-1, + }, + "pe1_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + + } + self.in_vars["parameters"] = [ + "q_min", + ] + + self.out_vars = { + "q1": { + "kend": grid.npz-1, + }, + + } + + # Value from GEOS + self._kord_tm = 9 + + # mode / iv set to 1 from GEOS + self.mode = 1 + + self.dims=[X_DIM, Y_DIM, Z_DIM] + + self._compute_func = MapSingle( + self.stencil_factory, + self.quantity_factory, + self._kord_tm, + self.mode, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + def compute_from_storage(self, inputs): + self._compute_func( + inputs["q1"], + inputs["pe1_"], + inputs["pe2_"], + qmin=Float(inputs["q_min"]), + interp=True, + ) + return inputs From 571ddfb846d97747f28c5bd2383594bfd3b3bd57 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 30 Oct 2024 03:40:22 -0700 Subject: [PATCH 017/252] Clarified type parameters in translate tests --- .../translate_lagrangian_contribution_interp.py | 9 +++------ tests/savepoint/translate/translate_map_scalar.py | 11 ++++------- .../savepoint/translate/translate_scalar_profile.py | 13 ++++--------- 3 files changed, 11 insertions(+), 22 deletions(-) diff --git a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py index 41382a12..011ac497 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py @@ -1,9 +1,8 @@ -from ndsl import StencilFactory, QuantityFactory -from ndsl.dsl.typing import FloatField, FloatFieldIJ, BoolFieldIJ, IntField, IntFieldIJ, Float, Int, Bool +from ndsl import StencilFactory, Namelist +from ndsl.dsl.typing import FloatField, BoolFieldIJ, IntField, IntFieldIJ, Int, Bool from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from gt4py.cartesian.gtscript import PARALLEL, FORWARD, BACKWARD, computation, interval from pyFV3.stencils.map_single import lagrangian_contributions_interp class test_Lagragian_Contribution_Interp: @@ -18,9 +17,7 @@ def __init__( self._lagrangian_contributions_interp = stencil_factory.from_origin_domain( func=lagrangian_contributions_interp, - # origin=(3,3,0), origin=grid_indexing.origin_compute(), - # domain=(24,1,72), domain=(grid.nic, 1, grid.npz) ) @@ -57,7 +54,7 @@ def __call__( ) class TranslateLagrangian_Contribution_Interp(TranslateFortranData2Py): - def __init__(self, grid: Grid, namelist, stencil_factory): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py index 3c5eb86e..25f82a0f 100644 --- a/tests/savepoint/translate/translate_map_scalar.py +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -1,14 +1,11 @@ -from gt4py.cartesian.gtscript import PARALLEL, computation, interval - -from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField, Float +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils.map_single import MapSingle - class TranslateMap_Scalar(TranslateFortranData2Py): - def __init__(self, grid, namelist, stencil_factory): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid @@ -66,7 +63,7 @@ def compute_from_storage(self, inputs): inputs["q1"], inputs["pe1_"], inputs["pe2_"], - qmin=Float(inputs["q_min"]), + qmin=inputs["q_min"], interp=True, ) return inputs diff --git a/tests/savepoint/translate/translate_scalar_profile.py b/tests/savepoint/translate/translate_scalar_profile.py index d9a5af85..b0779d34 100644 --- a/tests/savepoint/translate/translate_scalar_profile.py +++ b/tests/savepoint/translate/translate_scalar_profile.py @@ -1,14 +1,11 @@ -from gt4py.cartesian.gtscript import PARALLEL, computation, interval - -from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField, Float +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils.remap_profile import RemapProfile - class TranslateScalar_Profile(TranslateFortranData2Py): - def __init__(self, grid, namelist, stencil_factory): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid @@ -56,8 +53,6 @@ def __init__(self, grid, namelist, stencil_factory): "jend": grid.je, "kend": grid.npz-1, }, - - } self.in_vars["parameters"] = [ "q_min", @@ -112,6 +107,6 @@ def compute_from_storage(self, inputs): inputs["q4_3"], inputs["q4_4"], inputs["dp1_"], - Float(inputs["q_min"]), + inputs["q_min"], ) return inputs From c50f11a6dcf74c56bbb1fcaac2b43c8057818083 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 30 Oct 2024 04:23:33 -0700 Subject: [PATCH 018/252] Added fv_mapz.F90.SER serialization file --- fv_mapz.F90.SER | 46 +++++++++++++++++++++++++++++++--------------- 1 file changed, 31 insertions(+), 15 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index e547be40..1006e0df 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -222,7 +222,7 @@ contains integer:: nt, liq_wat, ice_wat, rainwat, snowwat, cld_amt, graupel, iq, n, kmp, kp, k_next logical:: remap_t, remap_pt, remap_te - real, dimension(is:ie, js:je, km+1) :: pe_3d, pe1_3d, pe2_3d, pn2_3d, pk2_3d, peln_3d + real, dimension(is:ie, js:je, km+1) :: pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, peln_3d real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d @@ -240,6 +240,7 @@ contains !$ser verbatim pe2_3d = 0.0 !$ser verbatim peln_3d = 0.0 !$ser verbatim dp2_3d = 0.0 + !$ser verbatim pn1_3d = 0.0 !$ser verbatim pn2_3d = 0.0 !$ser verbatim pk2_3d = pk @@ -361,7 +362,7 @@ contains !$OMP graupel,q_con, sphum,cappa,r_vir,rcp,cp,k1k,delp, & !$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & -!$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,pe_3d, pe1_3d, pe2_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, peln_3d) & +!$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, peln_3d) & !$OMP private(gz,cvm,bkh,dp2, & !$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_) do 1000 j=js,je+1 @@ -612,16 +613,22 @@ contains is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.true.) else !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo !$ser savepoint MapScalar_2d-In !$ser verbatim mode=1 - !$ser data peln=peln pt=pt gz1d=gz pn2=pn2 mode=mode j_2d=js2d + !$ser data pe1_=pn1_3d pt=pt q1=te mode=mode j_2d=js2d q_min=-1.e25 !$ser verbatim endif call map_scalar(km, pn1, te, gz, & km, pn2, te, & is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), te_min) !$ser verbatim if(j == js2d) then !$ser savepoint MapScalar_2d-Out - !$ser data pt=pt + !$ser data q1=te !$ser verbatim endif endif else @@ -635,16 +642,23 @@ contains is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.false.) else !$ser verbatim if(j == js2d) then - !$ser savepoint MapScalar_2d-In + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Map_Scalar-In !$ser verbatim mode=1 - !$ser data peln=peln pt=pt gz1d=gz pn2=pn2 mode=mode j_2d=js2d kord_tm=kord_tm + !$ser verbatim qmin=184.0 + !$ser data pe1_=pn1_3d pe2_=pn2_3d q1=pt mode=mode j_2d=js2d q_min=qmin !$ser verbatim endif call map_scalar(km, pn1, pt, gz, & km, pn2, pt, & is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), t_min) !$ser verbatim if(j == js2d) then - !$ser savepoint MapScalar_2d-Out - !$ser data pt=pt + !$ser savepoint Map_Scalar-Out + !$ser data q1=pt !$ser verbatim endif endif endif @@ -1682,8 +1696,6 @@ contains enddo enddo -! Compute vertical subgrid distribution - if ( kord > 7 ) then !$ser verbatim if (j == 1) then !$ser verbatim do k = 1,kn !$ser verbatim do i = i1,i2 @@ -1706,7 +1718,14 @@ contains !$ser data dp1_=dp1_temp !$ser data q_min=q_min !$ser verbatim endif + +! Compute vertical subgrid distribution + if ( kord > 7 ) then call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + !$ser verbatim if (j == 1) then !$ser verbatim do k = 1,kn !$ser verbatim do i = i1,i2 @@ -1724,9 +1743,6 @@ contains !$ser data q4_4=q4_4_temp !$ser data dp1_=dp1_temp !$ser verbatim endif - else - call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) - endif ! NOTE : q1 and q2 fields being passed into map_scalar are identical variables ! even though q1 is declared at INTENT(IN) and q2 is declared as INTENT(IN/OUT). @@ -1750,7 +1766,7 @@ contains !$ser verbatim pe2_temp(i,jj,km+1) = pe2(i,km+1) !$ser verbatim enddo !$ser verbatim enddo -!$ser savepoint Lagrangian_Contribution-In +!$ser savepoint Lagrangian_Contribution_Interp-In !$ser data q1=q1 pe1_=pe1_temp pe2_=pe2_temp q4_1=q4_1_temp q4_2=q4_2_temp !$ser data q4_3=q4_3_temp q4_4=q4_4_temp dp1_=dp1_temp !$ser verbatim endif @@ -1828,7 +1844,7 @@ contains enddo !$ser verbatim if(j == 1) then -!$ser savepoint Lagrangian_Contribution-Out +!$ser savepoint Lagrangian_Contribution_Interp-Out !$ser data q1=q2 !$ser verbatim endif From 801b62e5f1f738082837da6844ef0571fd9bcbb1 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 30 Oct 2024 14:56:43 -0700 Subject: [PATCH 019/252] Added translate test for MapN_Tracer. Note that there are nans reported from the C24 translate test, but this is may be due to only data from j=1 being serialized. --- pyFV3/stencils/map_single.py | 4 +- .../translate/translate_MapN_Tracer_2d.py | 138 +++++++++++++----- 2 files changed, 107 insertions(+), 35 deletions(-) diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index 3219494b..7bd202b7 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -323,7 +323,7 @@ def __call__( self._q4_3, self._q4_4, self._dp1, - qmin, + Float(qmin), ) else: self._remap_profile( @@ -333,7 +333,7 @@ def __call__( self._q4_3, self._q4_4, self._dp1, - qmin, + Float(qmin), ) if(interp == False): diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index 9d39aacd..4459a98a 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -1,32 +1,96 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval -from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField, Float -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils.mapn_tracer import MapNTracer - class TranslateMapN_Tracer_2d(TranslateFortranData2Py): - def __init__(self, grid, namelist, stencil_factory): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { - "qtracers": { + "qvapor": { + "kend": grid.npz-1, + }, + "qliquid": { + "kend": grid.npz-1, + }, + "qice": { + "kend": grid.npz-1, + }, + "qrain": { + "kend": grid.npz-1, + }, + "qsnow": { + "kend": grid.npz-1, + }, + "qgraupel": { + "kend": grid.npz-1, + }, + "qcld": { + "kend": grid.npz-1, + }, + "qo3mr": { + "kend": grid.npz-1, + }, + "qsgs_tke": { + "kend": grid.npz-1, + }, + "pe1_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz }, - - + "dp2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1 + }, } - # self.in_vars["parameters"] = [ - # "q_min", - # ] self.out_vars = { - "qtracers": { + "qvapor": { + "kend": grid.npz-1, + }, + "qliquid": { + "kend": grid.npz-1, + }, + "qice": { + "kend": grid.npz-1, + }, + "qrain": { + "kend": grid.npz-1, + }, + "qsnow": { + "kend": grid.npz-1, + }, + "qgraupel": { + "kend": grid.npz-1, + }, + "qcld": { + "kend": grid.npz-1, + }, + "qo3mr": { + "kend": grid.npz-1, + }, + "qsgs_tke": { + "kend": grid.npz-1, }, } @@ -36,29 +100,37 @@ def __init__(self, grid, namelist, stencil_factory): # mode / iv set to 1 from GEOS self.mode = 1 - nq = 9 - - fill = True + self.nq = 9 - # self._mapn_tracer = MapNTracer( - # self.stencil_factory, - # self.quantity_factory, - # abs(self.kord), - # nq, - # fill=fill, - # tracers=tracers, - # ) + self.fill = True def compute_from_storage(self, inputs): - # self._compute_func( - # inputs["qs_"], - # inputs["q4_1"], - # inputs["q4_2"], - # inputs["q4_3"], - # inputs["q4_4"], - # inputs["dp1_"], - # Float(inputs["q_min"]), - # ) - print("qtracers shape: ", inputs["qtracers"].data.shape) + tracers = { "qvapor": inputs["qvapor"], + "qliquid": inputs["qliquid"], + "qice": inputs["qice"], + "qrain": inputs["qrain"], + "qsnow": inputs["qsnow"], + "qgraupel": inputs["qgraupel"], + "qcld": inputs["qcld"], + "qo3mr": inputs["qo3mr"], + "qsgs_tke": inputs["qsgs_tke"], + } + + self._compute_func = MapNTracer( + self.stencil_factory, + self.quantity_factory, + abs(self.kord), + self.nq, + fill=self.fill, + tracers=tracers, + ) + + self._compute_func( + inputs["pe1_"], + inputs["pe2_"], + inputs["dp2_"], + tracers, + ) + return inputs From 17a79d34c5968435fe9f688425a375cbb2cbd9f6 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 31 Oct 2024 10:50:53 -0700 Subject: [PATCH 020/252] Added translate test for cs_profile subroutine. Translate tests pass for C24 dataset. Also added updated fv_mapz.F90 serialization file --- fv_mapz.F90.SER | 83 ++++++++++++- pyFV3/stencils/remap_profile.py | 2 +- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_cs_profile.py | 112 ++++++++++++++++++ 4 files changed, 195 insertions(+), 5 deletions(-) create mode 100644 tests/savepoint/translate/translate_cs_profile.py diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 1006e0df..e12dc06c 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -667,15 +667,40 @@ contains ! Map constituents !---------------- if( nq > 5 ) then - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim + !$ser verbatim if (k < km+1) dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo !$ser savepoint MapN_Tracer_2d-In - !$ser data j_2d=js2d nq=nq qtracers=q(:,:,:,1:nq) pe1=pe1 pe2=pe2 dp2=dp2 q_min=qmin + !$ser data j_2d=js2d nq=nq pe1_=pe1_3d pe2_=pe2_3d dp2_=dp2_3d + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) !$ser verbatim endif call mapn_tracer(nq, km, pe1, pe2, q, dp2, kord_tr, j, & is, ie, isd, ied, jsd, jed, 0., fill) !$ser verbatim if(j == js2d) then !$ser savepoint MapN_Tracer_2d-Out - !$ser data qtracers=q(:,:,:,1:nq) + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) !$ser verbatim endif elseif ( nq > 0 ) then print*,'CODE NOT TESTED HERE 7' @@ -716,6 +741,7 @@ contains !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array if ( w_limiter ) then + print*,"w_limiter part" !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km !$ser verbatim do i = is, ie @@ -1882,6 +1908,19 @@ contains real pl, pr, qsum, dp, esl integer i, k, l, m, k0 +!$ser verbatim real q4_1_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_2_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_3_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_4_temp(i1:i2, i1:i2,km) +!$ser verbatim real dp1_temp(i1:i2, i1:i2,km) +!$ser verbatim real :: qs_2d(i1:i2, i1:i2) + +!$ser verbatim q4_1_temp = 0.0 +!$ser verbatim q4_2_temp = 0.0 +!$ser verbatim q4_3_temp = 0.0 +!$ser verbatim q4_4_temp = 0.0 +!$ser verbatim qs_2d = 0.0 + do k=1,km do i=i1,i2 dp1(i,k) = pe1(i,k+1) - pe1(i,k) @@ -1891,7 +1930,45 @@ contains ! Compute vertical subgrid distribution if ( kord >7 ) then +!$ser verbatim if (j == 1 .and. i2 == 24) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do i = i1,i2 +!$ser verbatim qs_2d(i,j) = qs(i) +!$ser verbatim enddo +!$ser savepoint CS_Profile-In +!$ser data qs_=qs_2d +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser data dp1_=dp1_temp +!$ser data kord_=kord +!$ser data iv_=iv +!$ser verbatim endif call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) +!$ser verbatim if (j == 1 .and. i2 == 24) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint CS_Profile-Out +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser verbatim endif else call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) endif diff --git a/pyFV3/stencils/remap_profile.py b/pyFV3/stencils/remap_profile.py index f37d32d3..5e5cc41e 100644 --- a/pyFV3/stencils/remap_profile.py +++ b/pyFV3/stencils/remap_profile.py @@ -691,5 +691,5 @@ def __call__( self._ext5, self._ext6, self._extm, - qmin, + Float(qmin), ) diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index afc28135..1d35d10d 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -91,4 +91,5 @@ from .translate_remapping_GEOS import TranslateRemapping_GEOS from .translate_scalar_profile import TranslateScalar_Profile from .translate_MapN_Tracer_2d import TranslateMapN_Tracer_2d -from .translate_map_scalar import TranslateMap_Scalar \ No newline at end of file +from .translate_map_scalar import TranslateMap_Scalar +from .translate_cs_profile import TranslateCS_Profile \ No newline at end of file diff --git a/tests/savepoint/translate/translate_cs_profile.py b/tests/savepoint/translate/translate_cs_profile.py new file mode 100644 index 00000000..7a207224 --- /dev/null +++ b/tests/savepoint/translate/translate_cs_profile.py @@ -0,0 +1,112 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.remap_profile import RemapProfile + +class TranslateCS_Profile(TranslateFortranData2Py): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "qs_": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_1": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_2": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_3": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_4": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "dp1_":{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + } + self.in_vars["parameters"] = [ + "iv_", + "kord_", + ] + + self.out_vars = { + "q4_1": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_2": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_3": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "q4_4": {"istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + } + + # # Value from GEOS + # self.kord = 9 + + # # mode / iv set to -2 from GEOS + # self.mode = -2 + + def compute_from_storage(self, inputs): + self._compute_func = RemapProfile( + self.stencil_factory, + self.quantity_factory, + inputs["kord_"], + inputs["iv_"], + dims=[X_DIM, Y_DIM, Z_DIM] + ) + + self._compute_func( + inputs["qs_"], + inputs["q4_1"], + inputs["q4_2"], + inputs["q4_3"], + inputs["q4_4"], + inputs["dp1_"], + ) + return inputs From dee5174a675738274ac676bce9c202c66960adcc Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 31 Oct 2024 13:26:29 -0700 Subject: [PATCH 021/252] Added translate test for map1_ppm subroutine. C24 translate test produces what I think are roundoff errors in the j=1 domain. There are nans appearing in the calcuations, but those appear for all j's not equal to 1 --- fv_mapz.F90.SER | 27 ++++++- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_cs_profile.py | 6 -- .../savepoint/translate/translate_map1_ppm.py | 72 +++++++++++++++++++ 4 files changed, 100 insertions(+), 8 deletions(-) create mode 100644 tests/savepoint/translate/translate_map1_ppm.py diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index e12dc06c..4908bd72 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -226,6 +226,8 @@ contains real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d + real, dimension(is:ie, js:je) :: ws_2d + !$ser verbatim real :: w_max_, w_min_ !$ser verbatim integer:: mode, abskord, iep1, iedp1, jedp1, js2d, o3mr, sgs_tke @@ -243,6 +245,7 @@ contains !$ser verbatim pn1_3d = 0.0 !$ser verbatim pn2_3d = 0.0 !$ser verbatim pk2_3d = pk + !$ser verbatim ws_2d = 0.0 remap_t = .false. remap_pt = .false. @@ -362,7 +365,8 @@ contains !$OMP graupel,q_con, sphum,cappa,r_vir,rcp,cp,k1k,delp, & !$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & -!$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, peln_3d) & +!$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,& +!$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, peln_3d, ws_2d) & !$OMP private(gz,cvm,bkh,dp2, & !$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_) do 1000 j=js,je+1 @@ -720,9 +724,30 @@ contains if ( .not. hydrostatic ) then ! Remap vertical wind: + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do i = is, ie + !$ser verbatim ws_2d(i,j) = ws(i,j) + !$ser verbatim enddo + !$ser savepoint Map1_PPM-In + !$ser data pe1_=pe1_3d + !$ser data pe2_=pe2_3d + !$ser data ws_=ws_2d + !$ser data w_=w + !$ser data kord_wz=kord_wz + !$ser verbatim endif call map1_ppm (km, pe1, w, ws(is,j), & km, pe2, w, & is, ie, j, isd, ied, jsd, jed, -2, kord_wz) + !$ser verbatim if(j == js2d) then + !$ser savepoint Map1_PPM-Out + !$ser data w_=w + !$ser verbatim endif ! Remap delz for hybrid sigma-p coordinate do k=1,km do i=is,ie diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 1d35d10d..78f41c47 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -92,4 +92,5 @@ from .translate_scalar_profile import TranslateScalar_Profile from .translate_MapN_Tracer_2d import TranslateMapN_Tracer_2d from .translate_map_scalar import TranslateMap_Scalar -from .translate_cs_profile import TranslateCS_Profile \ No newline at end of file +from .translate_cs_profile import TranslateCS_Profile +from .translate_map1_ppm import TranslateMap1_PPM \ No newline at end of file diff --git a/tests/savepoint/translate/translate_cs_profile.py b/tests/savepoint/translate/translate_cs_profile.py index 7a207224..30b73f4a 100644 --- a/tests/savepoint/translate/translate_cs_profile.py +++ b/tests/savepoint/translate/translate_cs_profile.py @@ -86,12 +86,6 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto }, } - # # Value from GEOS - # self.kord = 9 - - # # mode / iv set to -2 from GEOS - # self.mode = -2 - def compute_from_storage(self, inputs): self._compute_func = RemapProfile( self.stencil_factory, diff --git a/tests/savepoint/translate/translate_map1_ppm.py b/tests/savepoint/translate/translate_map1_ppm.py new file mode 100644 index 00000000..f6fa2db2 --- /dev/null +++ b/tests/savepoint/translate/translate_map1_ppm.py @@ -0,0 +1,72 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.map_single import MapSingle + +class TranslateMap1_PPM(TranslateFortranData2Py): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "w_": { + "kend": grid.npz-1, + }, + "pe1_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "ws_":{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + + } + self.in_vars["parameters"] = [ + "kord_wz", + ] + + self.out_vars = { + "w_": { + "kend": grid.npz-1, + }, + + } + + # mode / iv set to -2 from GEOS + self.mode = -2 + + self.dims=[X_DIM, Y_DIM, Z_DIM] + + def compute_from_storage(self, inputs): + self._compute_func = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_wz"], + self.mode, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._compute_func( + inputs["w_"], + inputs["pe1_"], + inputs["pe2_"], + qs=inputs["ws_"], + interp=False, + ) + return inputs From df937c7c3540c6060db354baf2dc651259432441 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 1 Nov 2024 12:41:15 -0700 Subject: [PATCH 022/252] Separated translate test for W and delz. Verified with C24 dataset --- tests/savepoint/translate/__init__.py | 3 +- ...te_map1_ppm.py => translate_map1_ppm_W.py} | 2 +- .../translate/translate_map1_ppm_delz.py | 72 +++++++++++++++++++ 3 files changed, 75 insertions(+), 2 deletions(-) rename tests/savepoint/translate/{translate_map1_ppm.py => translate_map1_ppm_W.py} (97%) create mode 100644 tests/savepoint/translate/translate_map1_ppm_delz.py diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 78f41c47..8708cbc5 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -93,4 +93,5 @@ from .translate_MapN_Tracer_2d import TranslateMapN_Tracer_2d from .translate_map_scalar import TranslateMap_Scalar from .translate_cs_profile import TranslateCS_Profile -from .translate_map1_ppm import TranslateMap1_PPM \ No newline at end of file +from .translate_map1_ppm_W import TranslateMap1_PPM_W +from .translate_map1_ppm_delz import TranslateMap1_PPM_delz \ No newline at end of file diff --git a/tests/savepoint/translate/translate_map1_ppm.py b/tests/savepoint/translate/translate_map1_ppm_W.py similarity index 97% rename from tests/savepoint/translate/translate_map1_ppm.py rename to tests/savepoint/translate/translate_map1_ppm_W.py index f6fa2db2..1b6e544b 100644 --- a/tests/savepoint/translate/translate_map1_ppm.py +++ b/tests/savepoint/translate/translate_map1_ppm_W.py @@ -4,7 +4,7 @@ from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils.map_single import MapSingle -class TranslateMap1_PPM(TranslateFortranData2Py): +class TranslateMap1_PPM_W(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py new file mode 100644 index 00000000..400ad750 --- /dev/null +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -0,0 +1,72 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.map_single import MapSingle + +class TranslateMap1_PPM_delz(TranslateFortranData2Py): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "delz_": { + "kend": grid.npz-1, + }, + "pe1_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "gz_":{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + + } + self.in_vars["parameters"] = [ + "kord_wz", + ] + + self.out_vars = { + "delz_": { + "kend": grid.npz-1, + }, + + } + + # mode / iv set to 1 from GEOS + self.mode = 1 + + self.dims=[X_DIM, Y_DIM, Z_DIM] + + def compute_from_storage(self, inputs): + self._compute_func = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_wz"], + self.mode, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._compute_func( + inputs["delz_"], + inputs["pe1_"], + inputs["pe2_"], + qs=inputs["gz_"], + interp=False, + ) + return inputs From 9447dc71f6efc46fa2f0a78b53aa7bf3953051cd Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 8 Nov 2024 08:56:59 -0800 Subject: [PATCH 023/252] Adding translate test for Pressures_mapU. Translate tests pass with C24 dataset --- pyFV3/stencils/remapping.py | 19 ++-- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_Pressures_mapU.py | 91 +++++++++++++++++++ .../translate/translate_remapping_GEOS.py | 19 ++++ 4 files changed, 123 insertions(+), 9 deletions(-) create mode 100644 tests/savepoint/translate/translate_Pressures_mapU.py diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 6843a59a..444e33ce 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -149,8 +149,8 @@ def moist_cv_pt_pressure( ) # NOTE : GEOS does not perform the delz computation at this location # # delz_adjust - # if __INLINED(not hydrostatic): - # delz = -delz / delp + if __INLINED(not hydrostatic): + delz = -delz / delp # pressure_updates with computation(FORWARD): @@ -201,11 +201,12 @@ def pn2_pk_delp( def pressures_mapu( pe: FloatField, - pe1: FloatField, + # pe1: FloatField, ak: FloatFieldK, bk: FloatFieldK, pe0: FloatField, pe3: FloatField, + ptop: Float, ): """ Args: @@ -219,18 +220,20 @@ def pressures_mapu( with computation(BACKWARD): with interval(-1, None): pe_bottom = pe - pe1_bottom = pe + # pe1_bottom = pe with interval(0, -1): pe_bottom = pe_bottom[0, 0, 1] - pe1_bottom = pe1_bottom[0, 0, 1] + # pe1_bottom = pe1_bottom[0, 0, 1] with computation(FORWARD): with interval(0, 1): - pe0 = pe + # pe0 = pe + pe0 = ptop with interval(1, None): - pe0 = 0.5 * (pe[0, -1, 0] + pe1) + # pe0 = 0.5 * (pe[0, -1, 0] + pe1) + pe0 = 0.5 * (pe[0, -1, 0] + pe) with computation(FORWARD), interval(...): bkh = 0.5 * bk - pe3 = ak + bkh * (pe_bottom[0, -1, 0] + pe1_bottom) + pe3 = ak + bkh * (pe_bottom[0, -1, 0] + pe_bottom) def pressures_mapv( diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 8708cbc5..106760af 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -94,4 +94,5 @@ from .translate_map_scalar import TranslateMap_Scalar from .translate_cs_profile import TranslateCS_Profile from .translate_map1_ppm_W import TranslateMap1_PPM_W -from .translate_map1_ppm_delz import TranslateMap1_PPM_delz \ No newline at end of file +from .translate_map1_ppm_delz import TranslateMap1_PPM_delz +from .translate_Pressures_mapU import TranslatePressures_mapU \ No newline at end of file diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py new file mode 100644 index 00000000..90380186 --- /dev/null +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -0,0 +1,91 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.remapping import pressures_mapu + +class TranslatePressures_mapU(TranslateFortranData2Py): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "pe_": { + "istart": grid.is_-1, + "iend": grid.ie+1, + "jstart": grid.js-1, + "jend": grid.je+1, + "kend": grid.npz + }, + "pe0_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe3_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "ak":{ + + }, + "bk":{ + + } + + } + self.in_vars["parameters"] = [ + "ptop", + ] + + self.out_vars = { + "pe0_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe3_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + + } + + grid_indexing = stencil_factory.grid_indexing + + self.dims=[X_DIM, Y_DIM, Z_DIM] + + self._domain_jextra = ( + grid_indexing.domain[0], + grid_indexing.domain[1] + 1, + grid_indexing.domain[2] + 1, + ) + + self._compute_func = stencil_factory.from_origin_domain( + pressures_mapu, + origin=grid_indexing.origin_compute(), + domain=self._domain_jextra, + ) + + def compute_from_storage(self, inputs): + self._compute_func( + inputs["pe_"], + inputs["ak"], + inputs["bk"], + inputs["pe0_"], + inputs["pe3_"], + inputs["ptop"], + ) + return inputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index fde25872..903fb817 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -14,6 +14,7 @@ Z_INTERFACE_DIM, ) from ndsl.dsl.typing import Float +from pyFV3.stencils.mapn_tracer import MapNTracer class TranslateRemapping_GEOS(TranslateDycoreFortranData2Py): def __init__( @@ -203,6 +204,15 @@ def __init__( domain=(grid.nic, 1, grid.npz+1), ) + self._compute_func = MapNTracer( + self.stencil_factory, + self.quantity_factory, + abs(self.kord), + self.nq, + fill=self.fill, + tracers=tracers, + ) + def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction @@ -251,4 +261,13 @@ def compute_from_storage(self, inputs): inputs["pk"], Float(inputs["akap"]), ) + + # now that we have the pressure profiles, we can start remapping + self._map_single_pt(pt, peln, self._pn2, qmin=self._t_min) + + self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) + + self._map_single_w(w, self._pe1, self._pe2, qs=wsd) + self._map_single_delz(delz, self._pe1, self._pe2) + return inputs From edfa8cdb845a3ae33f33f644be89b354c31cc5af Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 12 Nov 2024 16:51:11 -0800 Subject: [PATCH 024/252] Added test for map u portion of fv_mapz. Note that the translate test will likely not verify for u, mfy, and cy since the serialized data only has data for j=1, and the stencil perform over the entire stencil domain. However it does verify the pe0 and pe3 variables --- fv_mapz.F90.SER | 153 +++++++++++++++++- .../translate/translate_Pressures_mapU.py | 87 +++++++++- 2 files changed, 228 insertions(+), 12 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 4908bd72..188da288 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -222,11 +222,13 @@ contains integer:: nt, liq_wat, ice_wat, rainwat, snowwat, cld_amt, graupel, iq, n, kmp, kp, k_next logical:: remap_t, remap_pt, remap_te - real, dimension(is:ie, js:je, km+1) :: pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, peln_3d + real, dimension(is:ie, js:je, km+1) :: pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, peln_3d + real, dimension(is-1:ie+1,js-1:je+1,km+1) :: pe_3d + real, dimension(is:ie+1, js:je+1, km+1) :: pe0_3d, pe3_3d real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d - real, dimension(is:ie, js:je) :: ws_2d + real, dimension(is:ie, js:je) :: ws_2d, gz_2d !$ser verbatim real :: w_max_, w_min_ @@ -246,6 +248,7 @@ contains !$ser verbatim pn2_3d = 0.0 !$ser verbatim pk2_3d = pk !$ser verbatim ws_2d = 0.0 + !$ser verbatim gz_2d = 0.0 remap_t = .false. remap_pt = .false. @@ -366,7 +369,8 @@ contains !$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & !$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,& -!$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, peln_3d, ws_2d) & +!$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & +!$OMP peln_3d, ws_2d, gz_2d, pe0_3d, pe3_3d) & !$OMP private(gz,cvm,bkh,dp2, & !$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_) do 1000 j=js,je+1 @@ -734,7 +738,7 @@ contains !$ser verbatim do i = is, ie !$ser verbatim ws_2d(i,j) = ws(i,j) !$ser verbatim enddo - !$ser savepoint Map1_PPM-In + !$ser savepoint Map1_PPM_W-In !$ser data pe1_=pe1_3d !$ser data pe2_=pe2_3d !$ser data ws_=ws_2d @@ -745,7 +749,7 @@ contains km, pe2, w, & is, ie, j, isd, ied, jsd, jed, -2, kord_wz) !$ser verbatim if(j == js2d) then - !$ser savepoint Map1_PPM-Out + !$ser savepoint Map1_PPM_W-Out !$ser data w_=w !$ser verbatim endif ! Remap delz for hybrid sigma-p coordinate @@ -754,9 +758,30 @@ contains delz(i,j,k) = -delz(i,j,k) / delp(i,j,k) ! ="specific volume"/grav enddo enddo + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do i = is, ie + !$ser verbatim gz_2d(i,j) = gz(i) + !$ser verbatim enddo + !$ser savepoint Map1_PPM_delz-In + !$ser data pe1_=pe1_3d + !$ser data pe2_=pe2_3d + !$ser data gz_=gz_2d + !$ser data delz_=delz + !$ser data kord_wz=kord_wz + !$ser verbatim endif call map1_ppm (km, pe1, delz, gz, & km, pe2, delz, & is, ie, j, isd, ied, jsd, jed, 1, abs(kord_wz)) + !$ser verbatim if(j == js2d) then + !$ser savepoint Map1_PPM_delz-Out + !$ser data delz_=delz + !$ser verbatim endif do k=1,km do i=is,ie delz(i,j,k) = -delz(i,j,k)*dp2(i,k) @@ -766,7 +791,6 @@ contains !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array if ( w_limiter ) then - print*,"w_limiter part" !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km !$ser verbatim do i = is, ie @@ -860,6 +884,41 @@ contains endif !(j < je+1) +!$ser verbatim if(j == js2d) then +!$ser verbatim pe_3d = 0 +!$ser verbatim pe0_3d = 0 +!$ser verbatim pe3_3d = 0 +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim pe_3d(i,j-1,k) = pe(i,k,j-1) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do i = is, ie +!$ser verbatim gz_2d(i,j) = gz(i) +!$ser verbatim enddo +!$ser savepoint Pressures_mapU-In +!$ser data pe_=pe_3d +!$ser data pe0_=pe0_3d(is:ie,js:je,:) +!$ser data pe3_=pe3_3d(is:ie,js:je,:) +!$ser data ak=ak +!$ser data bk=bk +!$ser data ptop=ptop +!$ser data u_=u +!$ser data kord_mt=kord_mt +!$ser verbatim if (present(mfy)) then +!$ser data mfy_=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy_=cy +!$ser verbatim endif +!$ser verbatim endif !------ ! map u !------ @@ -880,6 +939,28 @@ contains enddo enddo +!!$ser verbatim if(j == js2d) then +!!$ser verbatim pe0_3d = 0.0 +!!$ser verbatim pe3_3d = 0.0 +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie+1 +!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Pressures_mapU-In +!!$ser data pe0_=pe0_3d(is:ie,:,:) +!!$ser data pe3_=pe3_3d(is:ie,:,:) +!!$ser data u_=u +!!$ser data kord_mt=kord_mt +!!$ser verbatim if (present(mfy)) then +!!$ser data mfy_=mfy +!!$ser verbatim endif +!!$ser verbatim if (present(cy)) then +!!$ser data cy_=cy +!!$ser verbatim endif +!!$ser verbatim endif + call map1_ppm( km, pe0(is:ie,:), u, gz, & km, pe3(is:ie,:), u, & is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) @@ -896,6 +977,60 @@ contains is, ie, j, isd, ied, js, je+1, -1, kord_mt) endif +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Pressures_mapU-Out +!$ser data pe0_=pe0_3d(is:ie,js:je,:) +!$ser data pe3_=pe3_3d(is:ie,js:je,:) +!$ser data u_=u +!$ser verbatim if (present(mfy)) then +!$ser data mfy_=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy_=cy +!$ser verbatim endif +!$ser verbatim endif + +!$ser verbatim if(j == js2d) then +!$ser verbatim pe_3d = 0 +!$ser verbatim pe0_3d = 0 +!$ser verbatim pe3_3d = 0 +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim pe_3d(i-1,j,k) = pe(i-1,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do i = is, ie +!$ser verbatim gz_2d(i,j) = gz(i) +!$ser verbatim enddo +!$ser savepoint Pressures_mapV-In +!$ser data pe_=pe_3d +!$ser data pe0_=pe0_3d(is:ie,js:je,:) +!$ser data pe3_=pe3_3d(is:ie,js:je,:) +!$ser data ak=ak +!$ser data bk=bk +!$ser data ptop=ptop +!$ser data u_=u +!$ser data kord_mt=kord_mt +!$ser verbatim if (present(mfy)) then +!$ser data mfy_=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy_=cy +!$ser verbatim endif +!$ser verbatim endif !------ ! map v !------ @@ -1165,7 +1300,7 @@ contains if( last_step .and. (.not.do_adiabatic_init) ) then ! NOTE : Code can enter here since do_adiabatic_init can be False if ( consv > consv_min ) then - ! print*, "consv > consv_min = entered" + ! print*, "consv > consv_min = entered", consv, consv_min !$OMP do do j=js,je if ( hydrostatic ) then @@ -1189,6 +1324,7 @@ contains enddo else ! TE using 3D winds (pt is virtual temperature): + ! if(j==js) print*,"CODE EXECUTED 1" do i=is,ie te_2d(i,j) = 0. phis(i,km+1) = hs(i,j) @@ -1200,6 +1336,7 @@ contains enddo do k=1,km #ifdef MOIST_CAPPA + ! if(j==js) print*,"MOIST_CAPPA EXECUTED" call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & ice_wat, snowwat, graupel, q, gz, cvm) do i=is,ie @@ -1237,6 +1374,7 @@ contains enddo ! j-loop !$OMP single + ! if(j==js) print*,"MPP GLOBAL SUM CODE EXECUTED 1" tesum = mpp_global_sum(domain, te_2d*gridstruct%area_64(is:ie,js:je), & flags=BITWISE_EFP_SUM) E_Flux = DBLE(consv)*tesum / DBLE(grav*pdt*4.*pi*radius**2) ! unit: W/m**2 @@ -1352,6 +1490,7 @@ contains enddo ! j-loop enddo ! k-loop #ifdef USE_COND + ! print*, "USE_COND active" ! Fill condensate output !$OMP do do k=1,km diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index 90380186..f944959d 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -1,8 +1,9 @@ from ndsl import StencilFactory, Namelist from ndsl.stencils.testing.grid import Grid -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Y_INTERFACE_DIM from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils.remapping import pressures_mapu +from pyFV3.stencils.map_single import MapSingle class TranslatePressures_mapU(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): @@ -38,10 +39,35 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto }, "bk":{ - } + }, + "u_": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed+1, + "kend": grid.npz-1, + }, + + "mfy_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + + "cy_": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + } self.in_vars["parameters"] = [ + "kord_mt", "ptop", ] @@ -60,7 +86,29 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "jend": grid.je, "kend": grid.npz }, - + "u_": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed+1, + "kend": grid.npz-1, + }, + + "mfy_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + + "cy_": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, } grid_indexing = stencil_factory.grid_indexing @@ -73,14 +121,22 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto grid_indexing.domain[2] + 1, ) - self._compute_func = stencil_factory.from_origin_domain( + self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), domain=self._domain_jextra, ) def compute_from_storage(self, inputs): - self._compute_func( + self._map1_ppm_u = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_mt"], + -1, + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + ) + + self._pressures_mapu( inputs["pe_"], inputs["ak"], inputs["bk"], @@ -88,4 +144,25 @@ def compute_from_storage(self, inputs): inputs["pe3_"], inputs["ptop"], ) + + self._map1_ppm_u( + inputs["u_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) + + self._map1_ppm_u( + inputs["mfy_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) + + self._map1_ppm_u( + inputs["cy_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) return inputs From 4e6367ad471c19bf18341bca6a34a4b4f696f0bf Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 13 Nov 2024 10:10:52 -0800 Subject: [PATCH 025/252] Add translate test for map v. Solution verifies with C24 test case. Take a look at the comments to properly test the various part. These comments also apply to the testing of map u --- fv_mapz.F90.SER | 60 +++++-- pyFV3/stencils/remapping.py | 5 +- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_Pressures_mapV.py | 164 ++++++++++++++++++ 4 files changed, 215 insertions(+), 17 deletions(-) create mode 100644 tests/savepoint/translate/translate_Pressures_mapV.py diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 188da288..4e4a9f9a 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -996,6 +996,8 @@ contains !$ser verbatim endif !$ser verbatim endif +! Note : This serialization portion will test the update of pe0 +! and pe3 !$ser verbatim if(j == js2d) then !$ser verbatim pe_3d = 0 !$ser verbatim pe0_3d = 0 @@ -1012,24 +1014,12 @@ contains !$ser verbatim pe3_3d(i,j,k) = pe3(i,k) !$ser verbatim enddo !$ser verbatim enddo -!$ser verbatim do i = is, ie -!$ser verbatim gz_2d(i,j) = gz(i) -!$ser verbatim enddo !$ser savepoint Pressures_mapV-In !$ser data pe_=pe_3d -!$ser data pe0_=pe0_3d(is:ie,js:je,:) -!$ser data pe3_=pe3_3d(is:ie,js:je,:) +!$ser data pe0_v=pe0_3d(:,js:je,:) +!$ser data pe3_v=pe3_3d(:,js:je,:) !$ser data ak=ak !$ser data bk=bk -!$ser data ptop=ptop -!$ser data u_=u -!$ser data kord_mt=kord_mt -!$ser verbatim if (present(mfy)) then -!$ser data mfy_=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy_=cy -!$ser verbatim endif !$ser verbatim endif !------ ! map v @@ -1048,6 +1038,29 @@ contains enddo enddo +! Note : This serialization portion will test the update of v, +! mfx, and cx +!!$ser verbatim if(j == js2d) then +!!$ser verbatim pe0_3d = 0.0 +!!$ser verbatim pe3_3d = 0.0 +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie+1 +!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Pressures_mapV-In +!!$ser data pe0_v=pe0_3d(:,js:je,:) +!!$ser data pe3_v=pe3_3d(:,js:je,:) +!!$ser data v_=v +!!$ser data kord_mt=kord_mt +!!$ser verbatim if (present(mfx)) then +!!$ser data mfx_=mfx +!!$ser verbatim endif +!!$ser verbatim if (present(cx)) then +!!$ser data cx_=cx +!!$ser verbatim endif +!!$ser verbatim endif call map1_ppm (km, pe0, v, gz, & km, pe3, v, is, ie+1, & j, isd, ied+1, jsd, jed, -1, kord_mt) @@ -1065,6 +1078,25 @@ contains endif endif ! (j < je+1) +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Pressures_mapV-Out +!$ser data pe0_v=pe0_3d(:,js:je,:) +!$ser data pe3_v=pe3_3d(:,js:je,:) +!$ser data v_=v +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser verbatim endif + 1000 continue ! Update pressure variables and get new pkz, T_v, and omega diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 444e33ce..ee735a4b 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -255,8 +255,9 @@ def pressures_mapv( pe_bottom = pe_bottom[0, 0, 1] with computation(FORWARD): with interval(0, 1): - pe3 = ak - pe0 = pe + bkh = 0.5 * bk + pe3 = ak + bkh * (pe_bottom[-1, 0, 0] + pe_bottom) + # pe0 = pe with interval(1, None): bkh = 0.5 * bk pe0 = 0.5 * (pe[-1, 0, 0] + pe) diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 106760af..2cbd7368 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -95,4 +95,5 @@ from .translate_cs_profile import TranslateCS_Profile from .translate_map1_ppm_W import TranslateMap1_PPM_W from .translate_map1_ppm_delz import TranslateMap1_PPM_delz -from .translate_Pressures_mapU import TranslatePressures_mapU \ No newline at end of file +from .translate_Pressures_mapU import TranslatePressures_mapU +from .translate_Pressures_mapV import TranslatePressures_mapV \ No newline at end of file diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py new file mode 100644 index 00000000..ef3f20af --- /dev/null +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -0,0 +1,164 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid +from ndsl.constants import X_DIM, Y_DIM, Z_DIM, X_INTERFACE_DIM +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.remapping import pressures_mapv +from pyFV3.stencils.map_single import MapSingle + +class TranslatePressures_mapV(TranslateFortranData2Py): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "pe_": { + "istart": grid.is_-1, + "iend": grid.ie+1, + "jstart": grid.js-1, + "jend": grid.je+1, + "kend": grid.npz + }, + "pe0_v": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe3_v": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "ak":{ + + }, + "bk":{ + + }, + # "v_": { + # "istart": grid.isd, + # "iend": grid.ied+1, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz-1, + # }, + + # "mfx_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + + # "cx_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz-1, + # }, + + + } + self.in_vars["parameters"] = [ + # "kord_mt", + ] + + self.out_vars = { + "pe0_v": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + "pe3_v": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, + # "v_": { + # "istart": grid.isd, + # "iend": grid.ied+1, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz-1, + # }, + + # "mfx_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + + # "cx_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz-1, + # }, + } + + grid_indexing = stencil_factory.grid_indexing + + self.dims=[X_DIM, Y_DIM, Z_DIM] + + self._pressures_mapv = stencil_factory.from_origin_domain( + pressures_mapv, + origin=grid_indexing.origin_compute(), + domain=( + grid_indexing.domain[0] + 1, + grid_indexing.domain[1], + grid_indexing.domain[2] + 1, + ), + ) + + def compute_from_storage(self, inputs): + # self._map1_ppm_v = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # inputs["kord_mt"], + # -1, + # dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + # ) + + self._pressures_mapv( + inputs["pe_"], + inputs["ak"], + inputs["bk"], + inputs["pe0_v"], + inputs["pe3_v"], + ) + + # self._map1_ppm_v( + # inputs["v_"], + # inputs["pe0_v"], + # inputs["pe3_v"], + # interp=False, + # ) + + # self._map1_ppm_v( + # inputs["mfx_"], + # inputs["pe0_v"], + # inputs["pe3_v"], + # interp=False, + # ) + + # self._map1_ppm_v( + # inputs["cx_"], + # inputs["pe0_v"], + # inputs["pe3_v"], + # interp=False, + # ) + return inputs From f663587cbc4f14fb58acadd1fb2fadbed8686d0e Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 13 Nov 2024 15:46:33 -0500 Subject: [PATCH 026/252] XPPM & YPPM update for -6 option --- pyFV3/stencils/xppm.py | 69 ++++++++++++++++++++++++++++----------- pyFV3/stencils/yppm.py | 73 ++++++++++++++++++++++++++++++++---------- 2 files changed, 107 insertions(+), 35 deletions(-) diff --git a/pyFV3/stencils/xppm.py b/pyFV3/stencils/xppm.py index 77d7780e..1df8d42a 100644 --- a/pyFV3/stencils/xppm.py +++ b/pyFV3/stencils/xppm.py @@ -46,12 +46,17 @@ def fx1_fn(courant, br, b0, bl): @gtscript.function def get_advection_mask(bl, b0, br): - from __externals__ import mord + from __externals__ import mord, i_start, i_end if __INLINED(mord == 5): smt5 = bl * br < 0 else: smt5 = (3.0 * abs(b0)) < abs(bl - br) + # Fix edge issues + with horizontal(region[i_start - 1, :], region[i_start, :]): + smt5 = bl * br < 0.0 + with horizontal(region[i_end, :], region[i_end + 1, :]): + smt5 = bl * br < 0.0 if smt5[-1, 0, 0] or smt5[0, 0, 0]: advection_mask = 1.0 @@ -162,10 +167,6 @@ def compute_al(q: FloatField, dxa: FloatFieldIJ): al = ppm.p1 * (q[-1, 0, 0] + q) + ppm.p2 * (q[-2, 0, 0] + q[1, 0, 0]) - if __INLINED(iord < 0): - compile_assert(False) - al = max(al, 0.0) - if __INLINED(grid_type < 3): with horizontal(region[i_start - 1, :], region[i_end, :]): al = ppm.c1 * q[-2, 0, 0] + ppm.c2 * q[-1, 0, 0] + ppm.c3 * q @@ -182,6 +183,9 @@ def compute_al(q: FloatField, dxa: FloatFieldIJ): with horizontal(region[i_start + 1, :], region[i_end + 2, :]): al = ppm.c3 * q[-1, 0, 0] + ppm.c2 * q[0, 0, 0] + ppm.c1 * q[1, 0, 0] + if __INLINED(iord < 0): + al = max(al, 0.0) + return al @@ -273,7 +277,10 @@ def compute_blbr_ord8plus(q: FloatField, dxa: FloatFieldIJ): def compute_x_flux( - q: FloatField, courant: FloatField, dxa: FloatFieldIJ, xflux: FloatField + q: FloatField, + courant: FloatField, + dxa: FloatFieldIJ, + xflux: FloatField, ): """ Args: @@ -296,6 +303,30 @@ def compute_x_flux( class XPiecewiseParabolic: """ Fortran name is xppm + + + `iord` is `hord_dp` which is hord for `δp`, `δz`, where: + + `δp`: Total air mass (including vapor and condensates) + Equal to hydrostatic pressure depth of layer + `δz`: Geometric layer depth (nonhydrostatic) + + Value explainers: + 5: Unlimited “fifth-order” scheme with weak 2∆x filter; fastest + and least diffusive (“inviscid”) + 6: Intermediate-strength 2∆x filter. Gives best ACC and storm + structure but weaker TCs (“minimally-diffusive”) + 8: Lin 2004 monotone PPM constraint (“monotonic”) + 9: Hunyh constraint: more expensive but less diffusive than #8 + -5: #5 with a positive-definite constraint + + Undocumented values implemented in Fortran: 7, 10, 11, 12, 13. + + The code below is capable of: + - Cube-sphere grid (no doubly periodic) + - `iord` == 8 for monotonic behaviors OR + - `iord` 5, 6 + - `iord` must be positive """ def __init__( @@ -310,21 +341,20 @@ def __init__( # Arguments come from: # namelist.grid_type # grid.dxa - if grid_type == 3 or grid_type > 4: - raise NotImplementedError( - "X Piecewise Parabolic (xppm): " - f" grid type {grid_type} not implemented. <3 or 4 available." - ) - - if abs(iord) >= 8 and iord != 8: + available_grid_options = [0, 4] + if grid_type not in available_grid_options: raise NotImplementedError( - "X Piecewise Parabolic (xppm): " - f"iord {iord} != 8 not implemented when >= 8." + "Y Piecewise Parabolic (yppm) configuration: " + f"grid type {grid_type} not implemented. " + f"Options are {available_grid_options}." ) - if iord < 0: + available_iords = [-6, 5, 6, 8] + if iord not in available_iords: raise NotImplementedError( - f"X Piecewise Parabolic (xppm): iord {iord} < 0 not implemented." + "Y Piecewise Parabolic (yppm) configuration: " + f"iord {iord} not implemented. " + f"Options are {available_iords}." ) self._dxa = dxa @@ -372,7 +402,10 @@ def __call__( # were called "get_flux", while the routine which got the flux was called # fx1_fn. The final value was called xflux instead of q_out. self._compute_flux_stencil( - q_in, c, self._dxa, q_mean_advected_through_x_interface + q_in, + c, + self._dxa, + q_mean_advected_through_x_interface, ) # bl and br are "edge perturbation values" as in equation 4.1 # of the FV3 documentation diff --git a/pyFV3/stencils/yppm.py b/pyFV3/stencils/yppm.py index 129b134b..090327a6 100644 --- a/pyFV3/stencils/yppm.py +++ b/pyFV3/stencils/yppm.py @@ -46,12 +46,19 @@ def fx1_fn(courant, br, b0, bl): @gtscript.function def get_advection_mask(bl, b0, br): - from __externals__ import mord + from __externals__ import mord, j_start, j_end if __INLINED(mord == 5): smt5 = bl * br < 0 + elif __INLINED(mord == -5): + compile_assert(False) else: smt5 = (3.0 * abs(b0)) < abs(bl - br) + # Fix edge issues + with horizontal(region[:, j_start - 1], region[:, j_start]): + smt5 = bl * br < 0.0 + with horizontal(region[:, j_end], region[:, j_end + 1]): + smt5 = bl * br < 0.0 if smt5[0, -1, 0] or smt5[0, 0, 0]: advection_mask = 1.0 @@ -162,10 +169,6 @@ def compute_al(q: FloatField, dya: FloatFieldIJ): al = ppm.p1 * (q[0, -1, 0] + q) + ppm.p2 * (q[0, -2, 0] + q[0, 1, 0]) - if __INLINED(jord < 0): - compile_assert(False) - al = max(al, 0.0) - if __INLINED(grid_type < 3): with horizontal(region[:, j_start - 1], region[:, j_end]): al = ppm.c1 * q[0, -2, 0] + ppm.c2 * q[0, -1, 0] + ppm.c3 * q @@ -182,6 +185,9 @@ def compute_al(q: FloatField, dya: FloatFieldIJ): with horizontal(region[:, j_start + 1], region[:, j_end + 2]): al = ppm.c3 * q[0, -1, 0] + ppm.c2 * q[0, 0, 0] + ppm.c1 * q[0, 1, 0] + if __INLINED(jord < 0): + al = max(al, 0.0) + return al @@ -273,7 +279,10 @@ def compute_blbr_ord8plus(q: FloatField, dya: FloatFieldIJ): def compute_y_flux( - q: FloatField, courant: FloatField, dya: FloatFieldIJ, yflux: FloatField + q: FloatField, + courant: FloatField, + dya: FloatFieldIJ, + yflux: FloatField, ): """ Args: @@ -296,6 +305,29 @@ def compute_y_flux( class YPiecewiseParabolic: """ Fortran name is yppm + + `jord` is `hord_dp` which is hord for `δp`, `δz`, where: + + `δp`: Total air mass (including vapor and condensates) + Equal to hydrostatic pressure depth of layer + `δz`: Geometric layer depth (nonhydrostatic) + + Value explainers: + 5: Unlimited “fifth-order” scheme with weak 2∆x filter; fastest + and least diffusive (“inviscid”) + 6: Intermediate-strength 2∆x filter. Gives best ACC and storm + structure but weaker TCs (“minimally-diffusive”) + 8: Lin 2004 monotone PPM constraint (“monotonic”) + 9: Hunyh constraint: more expensive but less diffusive than #8 + -5: #5 with a positive-definite constraint + + Undocumented values implemented in Fortran: 7, 10, 11, 12, 13. + + The code below is capable of: + - Cube-sphere grid (no doubly periodic) + - `jord` == 8 for monotonic behaviors OR + - `jord` 5, 6 + - `jord` must be positive """ def __init__( @@ -307,25 +339,29 @@ def __init__( origin: Index3D, domain: Index3D, ): + # Dev note: this could be rewrote to split monotonic and not, or per-type of + # scheme as described above with compiler-time `jord` conditional to + # direct the code + orchestrate(obj=self, config=stencil_factory.config.dace_config) # Arguments come from: # namelist.grid_type # grid.dya - if grid_type == 3 or grid_type > 4: - raise NotImplementedError( - "Y Piecewise Parabolic (xppm): " - f" grid type {grid_type} not implemented. <3 or 4 available." - ) - if abs(jord) >= 8 and jord != 8: + available_grid_options = [0, 4] + if grid_type not in available_grid_options: raise NotImplementedError( - "Y Piecewise Parabolic (yppm): " - f"jord {jord} != 8 not implemented when >= 8." + "Y Piecewise Parabolic (yppm) configuration: " + f"grid type {grid_type} not implemented. " + f"Options are {available_grid_options}." ) - if jord < 0: + available_jords = [-6, 5, 6, 8] + if jord not in available_jords: raise NotImplementedError( - f"Y Piecewise Parabolic (yppm): jord {jord} < 0 not implemented." + "Y Piecewise Parabolic (yppm) configuration: " + f"jord {jord} not implemented. " + f"Options are {available_jords}." ) self._dya = dya @@ -373,7 +409,10 @@ def __call__( # were called "get_flux", while the routine which got the flux was called # fx1_fn. The final value was called yflux instead of q_out. self._compute_flux_stencil( - q_in, c, self._dya, q_mean_advected_through_y_interface + q_in, + c, + self._dya, + q_mean_advected_through_y_interface, ) # bl and br are "edge perturbation values" as in equation 4.1 # of the FV3 documentation From 68b36368ddce4d6d3a656f00d7c2f3a33f12f155 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 14 Nov 2024 10:42:19 -0500 Subject: [PATCH 027/252] Temporary script to run savepoints --- tests/script/geos_fp/TEMP/run_yppm_xppm.sh | 28 ++++++++++++++++++++++ 1 file changed, 28 insertions(+) create mode 100755 tests/script/geos_fp/TEMP/run_yppm_xppm.sh diff --git a/tests/script/geos_fp/TEMP/run_yppm_xppm.sh b/tests/script/geos_fp/TEMP/run_yppm_xppm.sh new file mode 100755 index 00000000..9795418c --- /dev/null +++ b/tests/script/geos_fp/TEMP/run_yppm_xppm.sh @@ -0,0 +1,28 @@ +#!/bin/bash + +THIS_DIR=$PWD +TEST_DATA_PATH="../../../../test_data/geos/TEMP_XPPM_YPMM" +mkdir -p $TEST_DATA_PATH +cd $TEST_DATA_PATH + +wget https://portal.nccs.nasa.gov/datashare/astg/smt/geos-fp/translate/11.5.2/x86_GNU/Dycore/TBC_C24_L72_Debug/YPPM-In.nc +wget https://portal.nccs.nasa.gov/datashare/astg/smt/geos-fp/translate/11.5.2/x86_GNU/Dycore/TBC_C24_L72_Debug/YPPM-Out.nc +wget https://portal.nccs.nasa.gov/datashare/astg/smt/geos-fp/translate/11.5.2/x86_GNU/Dycore/TBC_C24_L72_Debug/XPPM-In.nc +wget https://portal.nccs.nasa.gov/datashare/astg/smt/geos-fp/translate/11.5.2/x86_GNU/Dycore/TBC_C24_L72_Debug/XPPM-Out.nc +wget https://portal.nccs.nasa.gov/datashare/astg/smt/geos-fp/translate/11.5.2/x86_GNU/Dycore/TBC_C24_L72_Debug/input.nml +wget https://portal.nccs.nasa.gov/datashare/astg/smt/geos-fp/translate/11.5.2/x86_GNU/Dycore/TBC_C24_L72_Debug/Grid-Info.nc + + +cd $THIS_DIR +rm -r ./.gt_cache_* + +export PACE_FLOAT_PRECISION=32 +export PACE_CONSTANTS=GEOS +export FV3_DACEMODE=Python + +python -m pytest -v -s -x \ + --data_path=$TEST_DATA_PATH \ + --backend=numpy \ + --which_modules=XPPM,YPPM \ + --multimodal_metric \ + ../../../savepoint From bd3022d0052d1ee0445d10f92a833d3582b4eca2 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 14 Nov 2024 11:06:23 -0500 Subject: [PATCH 028/252] Lint --- pyFV3/stencils/xppm.py | 2 +- pyFV3/stencils/yppm.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/xppm.py b/pyFV3/stencils/xppm.py index 1df8d42a..38e17670 100644 --- a/pyFV3/stencils/xppm.py +++ b/pyFV3/stencils/xppm.py @@ -46,7 +46,7 @@ def fx1_fn(courant, br, b0, bl): @gtscript.function def get_advection_mask(bl, b0, br): - from __externals__ import mord, i_start, i_end + from __externals__ import i_end, i_start, mord if __INLINED(mord == 5): smt5 = bl * br < 0 diff --git a/pyFV3/stencils/yppm.py b/pyFV3/stencils/yppm.py index 090327a6..0b92f661 100644 --- a/pyFV3/stencils/yppm.py +++ b/pyFV3/stencils/yppm.py @@ -46,7 +46,7 @@ def fx1_fn(courant, br, b0, bl): @gtscript.function def get_advection_mask(bl, b0, br): - from __externals__ import mord, j_start, j_end + from __externals__ import j_end, j_start, mord if __INLINED(mord == 5): smt5 = bl * br < 0 From 70ea601b8373a48afc5fe1f63732e0e1a32cc3bd Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 29 Nov 2024 14:11:59 -0500 Subject: [PATCH 029/252] Move RFF calculation to f64 Fix the implementation check for RayleighDamping and move it closer to the code --- pyFV3/stencils/fv_dynamics.py | 6 ------ pyFV3/stencils/ray_fast.py | 25 +++++++++++++++++++------ 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 50ffba27..edfb1bfb 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -483,12 +483,6 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool): "Dynamical Core (fv_dynamics): compute total energy is not implemented" ) - if (not self.config.rf_fast) and self.config.tau != 0: - raise NotImplementedError( - "Dynamical Core (fv_dynamics): Rayleigh_Super," - " called when rf_fast=False and tau !=0, is not implemented" - ) - if self.config.adiabatic and self.config.kord_tm > 0: raise NotImplementedError( "Dynamical Core (fv_dynamics): Adiabatic with positive kord_tm" diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index dbc082d4..1ac71a8a 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -10,6 +10,7 @@ log, region, sin, + f64, ) import ndsl.constants as constants @@ -18,7 +19,7 @@ from ndsl.dsl.typing import Float, FloatField, FloatFieldK -SDAY = 86400.0 +SDAY = Float(86400.0) # NOTE: The fortran version of this computes rf in the first timestep only. Then @@ -36,7 +37,7 @@ def compute_rf_vals(pfull, bdt, rf_cutoff, tau0, ptop): @gtscript.function def compute_rff_vals(pfull, dt, rf_cutoff, tau0, ptop): rffvals = compute_rf_vals(pfull, dt, rf_cutoff, tau0, ptop) - rffvals = 1.0 / (1.0 + rffvals) + rffvals = f64(1.0) / (f64(1.0) + rffvals) return rffvals @@ -155,14 +156,26 @@ class RayleighDamping: Fortran name: ray_fast. """ - def __init__(self, stencil_factory: StencilFactory, rf_cutoff, tau, hydrostatic): + def __init__( + self, + stencil_factory: StencilFactory, + rf_cutoff: Float, + tau: Float, + hydrostatic: bool, + ): orchestrate(obj=self, config=stencil_factory.config.dace_config) grid_indexing = stencil_factory.grid_indexing - self._rf_cutoff = rf_cutoff + self._rf_cutoff = Float(rf_cutoff) origin, domain = grid_indexing.get_origin_domain( [X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM] ) + if tau == 0: + raise NotImplementedError( + "Dynamical Core (fv_dynamics): RayleighDamping," + " with tau <= 0, is not implemented" + ) + ax_offsets = grid_indexing.axis_offsets(origin, domain) local_axis_offsets = {} for axis_offset_name, axis_offset_value in ax_offsets.items(): @@ -175,7 +188,7 @@ def __init__(self, stencil_factory: StencilFactory, rf_cutoff, tau, hydrostatic) domain=domain, externals={ "hydrostatic": hydrostatic, - "rf_cutoff": rf_cutoff, + "rf_cutoff": self._rf_cutoff, "tau": tau, **local_axis_offsets, }, @@ -191,7 +204,7 @@ def __call__( dt: Float, ptop: Float, ): - rf_cutoff_nudge = self._rf_cutoff + min(100.0, 10.0 * ptop) + rf_cutoff_nudge = self._rf_cutoff + min(Float(100.0), Float(10.0) * ptop) self._ray_fast_wind_compute( u, From 4c8403abef8467c5dbdc2837da1fcd3ae78d33e3 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 29 Nov 2024 14:16:03 -0500 Subject: [PATCH 030/252] lint --- pyFV3/stencils/ray_fast.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index 1ac71a8a..708d1e9f 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -5,12 +5,12 @@ FORWARD, PARALLEL, computation, + f64, horizontal, interval, log, region, sin, - f64, ) import ndsl.constants as constants From bad27a887e83b98ca2f974096c2a9b5d429ac3d8 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 2 Dec 2024 16:27:36 -0500 Subject: [PATCH 031/252] Use `ndsl.constants.SECONDS_PER_DAY` --- pyFV3/stencils/ray_fast.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index 708d1e9f..3a65c417 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -15,13 +15,10 @@ import ndsl.constants as constants from ndsl import StencilFactory, orchestrate -from ndsl.constants import X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.constants import SECONDS_PER_DAY, X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField, FloatFieldK -SDAY = Float(86400.0) - - # NOTE: The fortran version of this computes rf in the first timestep only. Then # rf_initialized let's you know you can skip it. Here we calculate it every # time. @@ -77,7 +74,7 @@ def ray_fast_wind_compute( if pfull < rf_cutoff: # rf is rayleigh damping increment, fraction of vertical velocity # left after doing rayleigh damping (w -> w * rf) - rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SDAY, ptop) + rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SECONDS_PER_DAY, ptop) with computation(FORWARD): with interval(0, 1): if pfull < rf_cutoff_nudge: From 9cfae859b8972761fed8a0c884879b2126041451 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 2 Dec 2024 16:30:02 -0500 Subject: [PATCH 032/252] Replicate Fortran substraction to reduce f32 errors Use `np.power` to insure proper type casting of power op Verbose the half damp calculation --- pyFV3/stencils/delnflux.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/pyFV3/stencils/delnflux.py b/pyFV3/stencils/delnflux.py index 58607d6a..dfd447be 100644 --- a/pyFV3/stencils/delnflux.py +++ b/pyFV3/stencils/delnflux.py @@ -8,6 +8,7 @@ from ndsl.dsl.stencil import get_stencils_with_varied_bounds from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK from ndsl.grid import DampingCoefficients +import numpy as np def calc_damp(damp_c: Quantity, da_min: Float, nord: Quantity) -> Quantity: @@ -16,7 +17,7 @@ def calc_damp(damp_c: Quantity, da_min: Float, nord: Quantity) -> Quantity: "current implementation requires damp_c and nord to have " "identical data shape and dims" ) - data = (damp_c.data * da_min) ** (nord.data + 1) + data = np.power((damp_c.data * da_min), (nord.data + 1), dtype=Float) return Quantity( data=data, dims=damp_c.dims, @@ -97,7 +98,7 @@ def fx_calculation(q: FloatField, del6_v: FloatField): @gtscript.function def fx_calculation_neg(q: FloatField, del6_v: FloatField): - return -del6_v * (q[-1, 0, 0] - q) + return del6_v * (q - q[-1, 0, 0]) @gtscript.function @@ -107,7 +108,7 @@ def fy_calculation(q: FloatField, del6_u: FloatField): @gtscript.function def fy_calculation_neg(q: FloatField, del6_u: FloatField): - return -del6_u * (q[0, -1, 0] - q) + return del6_u * (q - q[0, -1, 0]) def d2_highorder_stencil( @@ -120,7 +121,7 @@ def d2_highorder_stencil( ): with computation(PARALLEL), interval(...): if nord > current_nord: - d2 = ((fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) * rarea + d2 = (fx - fx[1, 0, 0] + fy - fy[0, 1, 0]) * rarea def d2_damp_interval( @@ -178,8 +179,8 @@ def diffusive_damp( damp: FloatFieldK, ): with computation(PARALLEL), interval(...): - fx = fx + 0.5 * damp * (mass[-1, 0, 0] + mass) * fx2 - fy = fy + 0.5 * damp * (mass[0, -1, 0] + mass) * fy2 + fx = fx + (0.5 * damp) * (mass[-1, 0, 0] + mass) * fx2 + fy = fy + (0.5 * damp) * (mass[0, -1, 0] + mass) * fy2 def copy_corners_y_nord( From 398ebf3c2aba930ea89f14e9648c2e50a367a8c7 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 2 Dec 2024 16:35:30 -0500 Subject: [PATCH 033/252] Lint --- pyFV3/stencils/delnflux.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/delnflux.py b/pyFV3/stencils/delnflux.py index dfd447be..ed4304f1 100644 --- a/pyFV3/stencils/delnflux.py +++ b/pyFV3/stencils/delnflux.py @@ -1,6 +1,7 @@ from typing import Optional import gt4py.cartesian.gtscript as gtscript +import numpy as np from gt4py.cartesian.gtscript import PARALLEL, computation, horizontal, interval, region from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate @@ -8,7 +9,6 @@ from ndsl.dsl.stencil import get_stencils_with_varied_bounds from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK from ndsl.grid import DampingCoefficients -import numpy as np def calc_damp(damp_c: Quantity, da_min: Float, nord: Quantity) -> Quantity: From 7117df13ce1b5aaf1332d9876d5b35ed031f92a4 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 2 Dec 2024 22:54:35 -0800 Subject: [PATCH 034/252] Added in PE_pk_delp_peln stencil and translate test. Code "verifies" on rank 0 with errors that are repeated, which may mean that the errors are due to some rounding. --- pyFV3/stencils/remapping.py | 44 +++++ tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_Pressures_mapV.py | 10 +- .../translate/translate_pe_pk_delp_peln.py | 153 ++++++++++++++++++ 4 files changed, 204 insertions(+), 6 deletions(-) create mode 100644 tests/savepoint/translate/translate_pe_pk_delp_peln.py diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index ee735a4b..e245f8c0 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -292,6 +292,50 @@ def copy_from_below(a: FloatField, b: FloatField): with computation(PARALLEL), interval(1, None): b = a[0, 0, -1] +def pe_pk_delp_peln( + pe: FloatField, + pk: FloatField, + delp: FloatField, + peln: FloatField, + pe2: FloatField, + pk2: FloatField, + pn2: FloatField, + ak: FloatFieldK, + bk: FloatFieldK, + akap: Float, + ptop: Float, +): + + with computation(BACKWARD): + with interval(-1, None): + pe_bottom = pe + with interval(0, -1): + pe_bottom = pe_bottom[0, 0, 1] + + with computation(PARALLEL): + with interval(0, 1): + pe2 = ptop + pn2 = peln + pk2 = pk + with interval(1,-1): + pe2 = ak + bk * pe_bottom + pn2 = log(pe2) + pk2 = exp(akap*pn2) + with interval(-1, None): + pe2 = pe + pn2 = peln + pk2 = pk + + with computation(PARALLEL): + with interval(0,-1): + pe = pe2 + pk = pk2 + delp = pe2[0,0,1] - pe2[0,0,0] + peln = pn2 + with interval(-1,None): + pe = pe2 + pk = pk2 + peln = pn2 class LagrangianToEulerian: """ diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 2cbd7368..abb5c2c0 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -96,4 +96,5 @@ from .translate_map1_ppm_W import TranslateMap1_PPM_W from .translate_map1_ppm_delz import TranslateMap1_PPM_delz from .translate_Pressures_mapU import TranslatePressures_mapU -from .translate_Pressures_mapV import TranslatePressures_mapV \ No newline at end of file +from .translate_Pressures_mapV import TranslatePressures_mapV +from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln \ No newline at end of file diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index ef3f20af..f0904a9b 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -18,21 +18,21 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "iend": grid.ie+1, "jstart": grid.js-1, "jend": grid.je+1, - "kend": grid.npz + "kend": grid.npz+1 }, "pe0_v": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz+1 }, "pe3_v": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz+1 }, "ak":{ @@ -76,14 +76,14 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz+1, }, "pe3_v": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz+1, }, # "v_": { # "istart": grid.isd, diff --git a/tests/savepoint/translate/translate_pe_pk_delp_peln.py b/tests/savepoint/translate/translate_pe_pk_delp_peln.py new file mode 100644 index 00000000..17fea0da --- /dev/null +++ b/tests/savepoint/translate/translate_pe_pk_delp_peln.py @@ -0,0 +1,153 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils.remapping import pe_pk_delp_peln +from ndsl.stencils.testing.grid import Grid + +class TranslatePE_pk_delp_peln(TranslateFortranData2Py): + def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): + super().__init__(grid, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + self.quantity_factory = grid.quantity_factory + + self.in_vars["data_vars"] = { + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pe_": { + "istart": grid.is_-1, + "iend": grid.ie+1, + "jstart": grid.js-1, + "jend": grid.je+1, + "kend": grid.npz+1, + }, + "peln_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pn2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pk2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "delp": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz, + }, + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "ak":{ + }, + "bk":{ + }, + } + self.in_vars["parameters"] = [ + "akap", + "ptop", + ] + + self.out_vars = { + "pe2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pe_": { + "istart": grid.is_-1, + "iend": grid.ie+1, + "jstart": grid.js-1, + "jend": grid.je+1, + "kend": grid.npz+1, + }, + "peln_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pn2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pk2_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "delp": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz, + }, + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + } + + grid_indexing = stencil_factory.grid_indexing + self._domain_kextra = ( + grid_indexing.domain[0], + 1, + grid_indexing.domain[2] + 1, + ) + + + self._pe_pk_delp_peln = stencil_factory.from_origin_domain( + pe_pk_delp_peln, + origin=grid_indexing.origin_compute(), + domain=self._domain_kextra, + ) + + def compute_from_storage(self, inputs): + + self._pe_pk_delp_peln(inputs["pe_"], + inputs["pk"], + inputs["delp"], + inputs["peln_"], + inputs["pe2_"], + inputs["pk2_"], + inputs["pn2_"], + inputs["ak"], + inputs["bk"], + inputs["akap"], + inputs["ptop"], + ) + return inputs From 8d5aa2622fe1f7fd3938d614fd1ace36f994a03e Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 3 Dec 2024 08:16:03 -0800 Subject: [PATCH 035/252] Added serialized fv_mapz.F90 file --- fv_mapz.F90.SER | 49 ++++++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 48 insertions(+), 1 deletion(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 4e4a9f9a..c2b2f28a 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -1107,11 +1107,38 @@ contains !$OMP graupel,sphum,cappa,r_vir,rcp,cp,k1k,delp, & !$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, & -!$OMP hs,w,ws,kord_wz,do_omega,omga,rrg,kord_mt) & +!$OMP hs,w,ws,kord_wz,do_omega,omga,rrg,kord_mt, & +#ifdef SERIALIZE +!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, js2d, & +#endif +!$OMP pe2_3d, pe_3d, peln_3d, pn2_3d, pk2_3d) & !$OMP private(gz,cvm,kp,k_next,bkh,dp2, & !$OMP pe2,pe3,pk2,pn2,phis,tpe,dlnp,tmp) do 2000 j=js,je +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint PE_pk_delp_peln-In +!$ser data pe2_=pe2_3d +!$ser data pe_=pe_3d +!$ser data peln_=peln_3d +!$ser data pn2_=pn2_3d +!$ser data pk2_=pk2_3d +!$ser data delp=delp +!$ser data pk=pk +!$ser data ak=ak +!$ser data bk=bk +!$ser data akap=akap +!$ser data ptop=ptop +!$ser verbatim endif !---------- ! Update pe !---------- @@ -1169,6 +1196,26 @@ contains enddo enddo +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint PE_pk_delp_peln-Out +!$ser data pe2_=pe2_3d +!$ser data pe_=pe_3d +!$ser data peln_=peln_3d +!$ser data pn2_=pn2_3d +!$ser data pk2_=pk2_3d +!$ser data delp=delp +!$ser data pk=pk +!$ser verbatim endif + !--------------------- ! Compute pkz and T_v !--------------------- From dadf41197d2ed39ae8f73b11a8672e088f2c981a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 5 Dec 2024 15:02:51 -0500 Subject: [PATCH 036/252] Heat dissipation calculation when `do_skeb` Vorticity damping revert Fix unapplicable TODO Mixed precision --- pyFV3/stencils/d_sw.py | 76 +++++++++++++++++++----------------------- 1 file changed, 35 insertions(+), 41 deletions(-) diff --git a/pyFV3/stencils/d_sw.py b/pyFV3/stencils/d_sw.py index 28940049..64397b15 100644 --- a/pyFV3/stencils/d_sw.py +++ b/pyFV3/stencils/d_sw.py @@ -26,7 +26,7 @@ from pyFV3.version import IS_GEOS -dcon_threshold = 1e-5 +dcon_threshold = Float(1e-5) def flux_capacitor( @@ -92,6 +92,8 @@ def heat_diss( damp_w (in): ke_bg (in): """ + from __externals__ import do_stochastic_ke_backscatter + with computation(PARALLEL), interval(...): heat_source = 0.0 diss_est = 0.0 @@ -99,7 +101,8 @@ def heat_diss( dd8 = ke_bg * abs(dt) dw = (fx2 - fx2[1, 0, 0] + fy2 - fy2[0, 1, 0]) * rarea heat_source = dd8 - dw * (w + 0.5 * dw) - diss_est = heat_source + if __INLINED(do_stochastic_ke_backscatter): + diss_est = heat_source @gtscript.function @@ -373,9 +376,7 @@ def vort_differencing( from __externals__ import local_ie, local_is, local_je, local_js with computation(PARALLEL), interval(...): - # TODO: this should likely be dcon[k] rather than dcon[0] so that this - # can be turned on and off per-layer - if dcon[0] > dcon_threshold: + if dcon > dcon_threshold: # Creating a gtscript function for the ub/vb computation # results in an "NotImplementedError" error for Jenkins # Inlining the ub/vb computation in this stencil resolves the Jenkins error @@ -537,7 +538,6 @@ def heat_source_from_vorticity_damping( to explicitly damp and convert into heat. """ from __externals__ import ( # noqa (see below) - d_con, do_stochastic_ke_backscatter, local_ie, local_is, @@ -631,24 +631,15 @@ def set_low_kvals(col: Mapping[str, Quantity], k): # For the column namelist at a specific k-level # set the vorticity parameters if do_vort_damp is true -def vorticity_damping_option_FV3GFS(column, k, do_vort_damp): +def vorticity_damping_option(column, k, do_vort_damp): if do_vort_damp: column["nord_v"].view[k] = 0 column["damp_vt"].view[k] = 0.5 * column["d2_divg"].view[k] -def vorticity_damping_option_GEOS(column, k, do_vort_damp): - # GEOS does not set damp_vt - if do_vort_damp: - column["nord_v"].view[k] = 0 - - def lowest_kvals(column, k, do_vort_damp): set_low_kvals(column, k) - if IS_GEOS: - vorticity_damping_option_GEOS(column, k, do_vort_damp) - else: - vorticity_damping_option_FV3GFS(column, k, do_vort_damp) + vorticity_damping_option(column, k, do_vort_damp) def get_column_namelist( @@ -993,6 +984,9 @@ def make_quantity(): self._heat_diss_stencil = stencil_factory.from_dims_halo( func=heat_diss, compute_dims=[X_DIM, Y_DIM, Z_DIM], + externals={ + "do_stochastic_ke_backscatter": config.do_skeb, + }, ) self._heat_source_from_vorticity_damping_stencil = ( stencil_factory.from_dims_halo( @@ -1034,30 +1028,30 @@ def make_quantity(): def __call__( self, - delpc, - delp, - pt, - u, - v, - w, - uc, - vc, - ua, - va, - divgd, - mfx, - mfy, - cx, - cy, - crx, - cry, - xfx, - yfx, - q_con, - zh, - heat_source, - diss_est, - dt, + delpc: FloatField, + delp: FloatField, + pt: FloatField, + u: FloatField, + v: FloatField, + w: FloatField, + uc: FloatField, + vc: FloatField, + ua: FloatField, + va: FloatField, + divgd: FloatField, + mfx: FloatField, + mfy: FloatField, + cx: FloatField, + cy: FloatField, + crx: FloatField, + cry: FloatField, + xfx: FloatField, + yfx: FloatField, + q_con: FloatField, + zh: FloatField, + heat_source: FloatField, + diss_est: FloatField, + dt: Float, ): """ D-Grid shallow water routine, peforms a full-timestep advance From a600520d383ba79b20b374995800a248ad2b4a9f Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 9 Dec 2024 13:16:19 -0800 Subject: [PATCH 037/252] Added example MPI translate test for testing --- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_getMPIprop.py | 32 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) create mode 100644 tests/savepoint/translate/translate_getMPIprop.py diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index abb5c2c0..4e43310f 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -97,4 +97,5 @@ from .translate_map1_ppm_delz import TranslateMap1_PPM_delz from .translate_Pressures_mapU import TranslatePressures_mapU from .translate_Pressures_mapV import TranslatePressures_mapV -from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln \ No newline at end of file +from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln +from .translate_getMPIprop import TranslateGetMPIProp \ No newline at end of file diff --git a/tests/savepoint/translate/translate_getMPIprop.py b/tests/savepoint/translate/translate_getMPIprop.py new file mode 100644 index 00000000..ac4b29c0 --- /dev/null +++ b/tests/savepoint/translate/translate_getMPIprop.py @@ -0,0 +1,32 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing import ParallelTranslate +from ndsl.typing import Communicator +from mpi4py import MPI + +class TranslateGetMPIProp(ParallelTranslate): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + print("Base TranslateGetMPIProp is initialized") + super().__init__(grid, namelist, stencil_factory) + self._base.in_vars["data_vars"] = { + "delz" :{ + + } + } + self._base.out_vars = { + "delz" :{ + + } + } + + def compute_parallel(self, inputs, communicator: Communicator): + print("Communicator rank = ", communicator.rank) + print("Communicator size = ", communicator.size) + global_sum = communicator.comm.allreduce(communicator.rank, op=MPI.SUM) + print("global_sum of ranks = ", global_sum) + return inputs \ No newline at end of file From 177617fba588552cdd90ae8d3bfcad9648a46937 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 11 Dec 2024 10:52:26 -0500 Subject: [PATCH 038/252] Fluxes and courant numbers are _always_ 64-bit float FIx to translate test for 32-bit float dataset --- pyFV3/stencils/d_sw.py | 18 ++++++++-------- tests/savepoint/translate/translate_d_sw.py | 23 +++++++++++---------- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/pyFV3/stencils/d_sw.py b/pyFV3/stencils/d_sw.py index 64397b15..ceb67302 100644 --- a/pyFV3/stencils/d_sw.py +++ b/pyFV3/stencils/d_sw.py @@ -12,7 +12,7 @@ from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK +from ndsl.dsl.typing import Float, FloatField, FloatField64, FloatFieldIJ, FloatFieldK from ndsl.grid import DampingCoefficients, GridData from pyFV3._config import DGridShallowWaterLagrangianDynamicsConfig from pyFV3.stencils import delnflux @@ -30,10 +30,10 @@ def flux_capacitor( - cx: FloatField, - cy: FloatField, - xflux: FloatField, - yflux: FloatField, + cx: FloatField64, + cy: FloatField64, + xflux: FloatField64, + yflux: FloatField64, crx_adv: FloatField, cry_adv: FloatField, fx: FloatField, @@ -1039,10 +1039,10 @@ def __call__( ua: FloatField, va: FloatField, divgd: FloatField, - mfx: FloatField, - mfy: FloatField, - cx: FloatField, - cy: FloatField, + mfx: FloatField64, + mfy: FloatField64, + cx: FloatField64, + cy: FloatField64, crx: FloatField, cry: FloatField, xfx: FloatField, diff --git a/tests/savepoint/translate/translate_d_sw.py b/tests/savepoint/translate/translate_d_sw.py index 836f3996..350e245d 100644 --- a/tests/savepoint/translate/translate_d_sw.py +++ b/tests/savepoint/translate/translate_d_sw.py @@ -3,7 +3,7 @@ import pyFV3 import pyFV3.stencils.d_sw as d_sw from ndsl import Namelist, StencilFactory -from ndsl.dsl.typing import FloatField, FloatFieldIJ +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from pyFV3.testing import TranslateDycoreFortranData2Py @@ -44,10 +44,10 @@ def __init__( "crx": grid.x3d_compute_domain_y_dict(), "yfx": grid.y3d_compute_domain_x_dict(), "cry": grid.y3d_compute_domain_x_dict(), - "mfx": grid.x3d_compute_dict(), - "mfy": grid.y3d_compute_dict(), - "cx": grid.x3d_compute_domain_y_dict(), - "cy": grid.y3d_compute_domain_x_dict(), + "mfx": grid.x3d_compute_dict() | {"serialname": "mfxd_R8"}, + "mfy": grid.y3d_compute_dict() | {"serialname": "mfyd_R8"}, + "cx": grid.x3d_compute_domain_y_dict() | {"serialname": "cxd_R8"}, + "cy": grid.y3d_compute_domain_x_dict() | {"serialname": "cyd_R8"}, "heat_source": {}, "diss_est": {}, "q_con": {}, @@ -58,7 +58,8 @@ def __init__( "divgd": grid.default_dict_buffer_2d(), } for name, info in self.in_vars["data_vars"].items(): - info["serialname"] = name + "d" + if name not in ["mfx", "mfy", "cx", "cy"]: + info["serialname"] = name + "d" self.in_vars["parameters"] = ["dt"] self.out_vars = self.in_vars["data_vars"].copy() del self.out_vars["zh"] @@ -71,8 +72,8 @@ def ubke( rsina: FloatFieldIJ, ut: FloatField, ub: FloatField, - dt4: float, - dt5: float, + dt4: Float, + dt5: Float, ): with computation(PARALLEL), interval(...): dt = 2.0 * dt5 @@ -118,8 +119,8 @@ def vbke( rsina: FloatFieldIJ, vt: FloatField, vb: FloatField, - dt4: float, - dt5: float, + dt4: Float, + dt5: Float, ): with computation(PARALLEL), interval(...): dt = 2.0 * dt5 @@ -218,7 +219,7 @@ def compute_from_storage(self, inputs): # TODO add these to the serialized data or remove the test inputs["damp_w"] = column_namelist["damp_w"] inputs["ke_bg"] = column_namelist["ke_bg"] - inputs["dt"] = ( + inputs["dt"] = Float( self.namelist.dt_atmos / self.namelist.k_split / self.namelist.n_split ) inputs["rarea"] = self.grid.rarea From ccf0cadb9abcef5c12d01994b7eaea9d7c067885 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 11 Dec 2024 10:54:00 -0500 Subject: [PATCH 039/252] PPM factors are true 32-bit float --- pyFV3/stencils/a2b_ord4.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/pyFV3/stencils/a2b_ord4.py b/pyFV3/stencils/a2b_ord4.py index 328d3142..514d7a56 100644 --- a/pyFV3/stencils/a2b_ord4.py +++ b/pyFV3/stencils/a2b_ord4.py @@ -19,16 +19,16 @@ # comact 4-pt cubic interpolation -c1 = 2.0 / 3.0 -c2 = -1.0 / 6.0 -d1 = 0.375 -d2 = -1.0 / 24.0 +c1 = Float(2.0) / Float(3.0) +c2 = Float(-1.0) / Float(6.0) +d1 = Float(0.375) +d2 = Float(-1.0) / Float(24.0) # PPM volume mean form -b1 = 7.0 / 12.0 -b2 = -1.0 / 12.0 +b1 = Float(7.0) / Float(12.0) +b2 = Float(-1.0) / Float(12.0) # 4-pt Lagrange interpolation -a1 = 9.0 / 16.0 -a2 = -1.0 / 16.0 +a1 = Float(9.0) / Float(16.0) +a2 = Float(-1.0) / Float(16.0) @gtscript.function From dfc061462a3904f5a7c6b95ac22b2be08a649b28 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 11 Dec 2024 12:18:48 -0500 Subject: [PATCH 040/252] Make sure scalars are typed properly Swap `math` for `np` to allow for dtype to be respected --- pyFV3/stencils/riem_solver3.py | 10 +++++----- pyFV3/stencils/riem_solver_c.py | 2 +- pyFV3/stencils/sim1_solver.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pyFV3/stencils/riem_solver3.py b/pyFV3/stencils/riem_solver3.py index 1f915ae6..1e32eebb 100644 --- a/pyFV3/stencils/riem_solver3.py +++ b/pyFV3/stencils/riem_solver3.py @@ -1,6 +1,6 @@ -import math import typing +import numpy as np from gt4py.cartesian.gtscript import ( __INLINED, BACKWARD, @@ -162,7 +162,7 @@ def __init__( grid_indexing = stencil_factory.grid_indexing self._sim1_solve = Sim1Solver( stencil_factory, - config.p_fac, + Float(config.p_fac), n_halo=0, ) orchestrate( @@ -215,7 +215,7 @@ def __init__( ) self._finalize_stencil = stencil_factory.from_origin_domain( finalize, - externals={"use_logp": config.use_logp, "beta": config.beta}, + externals={"use_logp": config.use_logp, "beta": Float(config.beta)}, origin=riemorigin, domain=domain, ) @@ -284,9 +284,9 @@ def __call__( # gm2 is gamma (cp/cv) # dz2 is delz - peln1 = math.log(ptop) + peln1 = np.log(ptop, dtype=Float) # ptk = ptop ** kappa - ptk = math.exp(constants.KAPPA * peln1) + ptk = np.exp(constants.KAPPA * peln1, dtype=Float) self._precompute_stencil( delp, diff --git a/pyFV3/stencils/riem_solver_c.py b/pyFV3/stencils/riem_solver_c.py index df909165..e498b676 100644 --- a/pyFV3/stencils/riem_solver_c.py +++ b/pyFV3/stencils/riem_solver_c.py @@ -187,7 +187,7 @@ def __init__( ) self._sim1_solve = Sim1Solver( stencil_factory, - p_fac, + Float(p_fac), n_halo=1, ) self._finalize_stencil = stencil_factory.from_origin_domain( diff --git a/pyFV3/stencils/sim1_solver.py b/pyFV3/stencils/sim1_solver.py index 0fcd8f28..0bc1566f 100644 --- a/pyFV3/stencils/sim1_solver.py +++ b/pyFV3/stencils/sim1_solver.py @@ -199,8 +199,8 @@ def __call__( # TODO: email Lucas about any remaining variable naming here - t1g = 2.0 * dt * dt - rdt = 1.0 / dt + t1g = Float(2.0) * dt * dt + rdt = Float(1.0) / dt self._compute_sim1_solve( w, delta_mass, From 843a5ff3ad40efd8770d3a7fd6c4d58b3f307d2b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 11 Dec 2024 13:03:22 -0500 Subject: [PATCH 041/252] Make sure scalar in D2A2C vector have the proper type --- pyFV3/stencils/d2a2c_vect.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyFV3/stencils/d2a2c_vect.py b/pyFV3/stencils/d2a2c_vect.py index 1d622082..292b407a 100644 --- a/pyFV3/stencils/d2a2c_vect.py +++ b/pyFV3/stencils/d2a2c_vect.py @@ -9,9 +9,9 @@ from pyFV3.stencils.a2b_ord4 import a1, a2, lagrange_x_func, lagrange_y_func -c1 = -2.0 / 14.0 -c2 = 11.0 / 14.0 -c3 = 5.0 / 14.0 +c1 = Float(-2.0) / Float(14.0) +c2 = Float(11.0) / Float(14.0) +c3 = Float(5.0) / Float(14.0) OFFSET = 2 @@ -409,7 +409,7 @@ def __init__( self._sin_sg4 = grid_data.sin_sg4 self._grid_type = grid_type - self._big_number = 1e30 # 1e8 if 32 bit + self._big_number = Float(1e30) # 1e8 if 32 bit nx = grid_indexing.iec + 1 # grid.npx + 2 ny = grid_indexing.jec + 1 # grid.npy + 2 i1 = grid_indexing.isc - 1 From 184091eafbab836c9d883c32bbf6618992fcb9db Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 11 Dec 2024 13:33:57 -0500 Subject: [PATCH 042/252] Add README.md describing the branch --- GEOS_v11.4.2_BranchTracking.md | 36 ++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 GEOS_v11.4.2_BranchTracking.md diff --git a/GEOS_v11.4.2_BranchTracking.md b/GEOS_v11.4.2_BranchTracking.md new file mode 100644 index 00000000..f91bb1be --- /dev/null +++ b/GEOS_v11.4.2_BranchTracking.md @@ -0,0 +1,36 @@ +# Up-skilling to GEOS v11.4.2 + +This branch exists solely for up-skilling pyFV3 to be able run GEOS in it's v11.4.2 FP configuration. +The need for a seperate branch from `develop` rely in the following differences: + +- GEOS run a 32bit floating point precision version (with appropriate 64bit buffers for mass conservation). This means the translate test requires a new set of data _and_ will not pass on old 8.1.3 Pace data. +- GEOS requires options that are deemed "legacy" and that we may want to replace rather than port +- Project requirements demand quick iterative development, while `pyFV3` demands concertation between all stakeholders. + +The aim is to validate and benchmark GEOS v11.4.2 with this dynamics. Once done, we will aim to move _as much code as possible_ back into develop. +The methodology goes as follows + +- Merge directly into `develop` any changes that do not demand a new set of data +- Keep track of the feature branch (below) that can't be merged in `develop` for future PR +- Keep track of GEOS vs SHiELD differences for future discussions + +## Feature branches + +Legend: + +- ⚙️ _GEOS - WIP_ : Ongoing work - can be merged temporarily +- 🔶 _GEOS - Merged_: Considered done - merged in GEOS v11.4.2 branch but NOT in `develop` +- ✅ _Develop - Merged_: Work done as part of up-skilling done for GEOS merged in `develop` AND the GEOS v11.4.2 branch. + +Branches: + +- ✅ `fix/F32/UpdateDzC`@Florian: Fix for fluxes gradient +- ✅ `fix/F32/DivergenceDamping`@Florian: Fix for 32-bit scalars in DivergenceDamping +- 🔶 `fix/RayleighDamping_mixed_precision`@Florian: fix the Ray_Fast test +- 🔶 `GEOS_update/yppm_xppm`@Florian: fix the YPPM/XPPM with `hord = -6` +- 🔶 `fix/DelnFlux_f32_support`@Florian: Fix for f32 support for DelnFlux (partial pass) +- 🔶 `fix/F32/UpdateDzD`@Florian: Fix for fluxes gradient & python computation +- ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation and column calculation (partial pass) +- ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 +- ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@FlorianL: Fix for 32-bit A2B_Ord4 +- ⚙️ `fix/GEOSv11_4_2/C_SW`: Fix for C_SW for 32-bit From f08775763c0bcb809a7ca5adf10551ec7fc3a2ab Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 11 Dec 2024 13:39:13 -0500 Subject: [PATCH 043/252] Move markdown to README --- GEOS_v11.4.2_BranchTracking.md | 36 --- ORIG.README.md | 459 +++++++++++++++++++++++++++++++ README.md | 477 ++------------------------------- 3 files changed, 486 insertions(+), 486 deletions(-) delete mode 100644 GEOS_v11.4.2_BranchTracking.md create mode 100644 ORIG.README.md diff --git a/GEOS_v11.4.2_BranchTracking.md b/GEOS_v11.4.2_BranchTracking.md deleted file mode 100644 index f91bb1be..00000000 --- a/GEOS_v11.4.2_BranchTracking.md +++ /dev/null @@ -1,36 +0,0 @@ -# Up-skilling to GEOS v11.4.2 - -This branch exists solely for up-skilling pyFV3 to be able run GEOS in it's v11.4.2 FP configuration. -The need for a seperate branch from `develop` rely in the following differences: - -- GEOS run a 32bit floating point precision version (with appropriate 64bit buffers for mass conservation). This means the translate test requires a new set of data _and_ will not pass on old 8.1.3 Pace data. -- GEOS requires options that are deemed "legacy" and that we may want to replace rather than port -- Project requirements demand quick iterative development, while `pyFV3` demands concertation between all stakeholders. - -The aim is to validate and benchmark GEOS v11.4.2 with this dynamics. Once done, we will aim to move _as much code as possible_ back into develop. -The methodology goes as follows - -- Merge directly into `develop` any changes that do not demand a new set of data -- Keep track of the feature branch (below) that can't be merged in `develop` for future PR -- Keep track of GEOS vs SHiELD differences for future discussions - -## Feature branches - -Legend: - -- ⚙️ _GEOS - WIP_ : Ongoing work - can be merged temporarily -- 🔶 _GEOS - Merged_: Considered done - merged in GEOS v11.4.2 branch but NOT in `develop` -- ✅ _Develop - Merged_: Work done as part of up-skilling done for GEOS merged in `develop` AND the GEOS v11.4.2 branch. - -Branches: - -- ✅ `fix/F32/UpdateDzC`@Florian: Fix for fluxes gradient -- ✅ `fix/F32/DivergenceDamping`@Florian: Fix for 32-bit scalars in DivergenceDamping -- 🔶 `fix/RayleighDamping_mixed_precision`@Florian: fix the Ray_Fast test -- 🔶 `GEOS_update/yppm_xppm`@Florian: fix the YPPM/XPPM with `hord = -6` -- 🔶 `fix/DelnFlux_f32_support`@Florian: Fix for f32 support for DelnFlux (partial pass) -- 🔶 `fix/F32/UpdateDzD`@Florian: Fix for fluxes gradient & python computation -- ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation and column calculation (partial pass) -- ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 -- ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@FlorianL: Fix for 32-bit A2B_Ord4 -- ⚙️ `fix/GEOSv11_4_2/C_SW`: Fix for C_SW for 32-bit diff --git a/ORIG.README.md b/ORIG.README.md new file mode 100644 index 00000000..a6192aac --- /dev/null +++ b/ORIG.README.md @@ -0,0 +1,459 @@ +> DISCLAIMER: Work in progress + +# FV3core + +FV3core is a Python version, using GridTools GT4Py with CPU and GPU backend options, of the FV3 dynamical core (fv3gfs-fortran repo). +The code here includes regression test data of computation units coming from serialized output from the Fortran model generated using the `GridTools/serialbox` framework. + +As of January 10, 2021 this documentation is outdated in that it was written when we had fv3core as its own single repository. Some functionality, such as linting, has been moved to the top level but may still be described in this document as occuring inside the fv3core folder. + +**WARNING** This repo is under active development and relies on code and data that is not publicly available at this point. + +## QuickStart + +1. Ensure you have docker installed and available for building and running and has access to the VCM cloud + +Be sure to complete any required post-installation instructions (e.g. [for linux](https://docs.docker.com/engine/install/linux-postinstall/)). Also [authorize Docker to pull from gcr](https://cloud.google.com/container-registry/docs/advanced-authentication). Your user will need to have read access to the `us.gcr.io/vcm-ml` repository. + +2. You can build the image, download the data, and run the tests using: + +```shell +$ make tests savepoint_tests savepoint_tests_mpi +``` + +If you want to develop code, you should also install the linting requirements and git hooks locally + +```shell +$ pip install -c constraints.txt -r requirements/requirements_lint.txt +$ pre-commit install + +## Getting started, in more detail +If you want to build the main fv3core docker image, run + +```shell +$ make build +``` + +If you want to download test data run + +```shell +$ make get_test_data +``` + +And the c12_6ranks_standard data will download into the `test_data` directory. + +If you do not have a GCP account, there is an option to download basic test data from a public FTP server and you can skip the GCP authentication step above. To download test data from the FTP server, use `make USE_FTP=yes get_test_data` instead and this will avoid fetching from a GCP storage bucket. You will need a valid in stallation of the `lftp` command. + +MPI parallel tests (that run that way to exercise halo updates in the model) can also be run with: + +```shell +$ make savepoint_tests_mpi +``` + +The environment image that the fv3core container uses is prebuilt and lives in the GCR. The above commands will by default pull this image before building the fv3core image and running the tests. +To build the environment from scratch (including GT4py) before running tests, either run + +``` +make build_environment +``` + +or + +```shell +$ PULL=False make savepoint_tests +``` + +which will execute the target `build_environment` for you before running the tests. + +There are `push_environment` and `rebuild_environment` targets, but these should normally not be done manually. Updating the install image should only be done by Jenkins after the tests pass using a new environment. + +### Test data options + +If you want to run different test data, discover the possible options with +```shell +$ make list_test_data_options +``` +This will list the storage buckets in the cloud. Then to run one of them, set EXPERIMENT to the folder name of the data you'd like to use: + +e.g. +```shell +$EXPERIMENT=c48_6ranks_standard make tests +``` + +If you choose an experiment with a different number of ranks than 6, also set `NUM_RANKS=` + +## Testing interactively outside the container + +After `make savepoint_tests` has been run at least once (or you have data in test_data and the docker image fv3core exists because `make build` has been run), you can iterate on code changes using + +```shell +$ DEV=y make savepoint_tests +``` +or for the parallel or non-savepoint tests: + +```shell +$ DEV=y make tests savepoint_tests_mpi +``` +These will mount your current code into the fv3core container and run it rather than the code that was built when `make build` ran. + +## Running tests inside a container + +If you to prefer to work interactively inside the fv3core container, get the test data and build the docker image (see above if you do not have a GCP account and want to get test data): +```shell +$ make get_test_data +``` + +```shell +$ make build +``` +Testing can be run with this data from `/port_dev` inside the container: + +```shell +$ make dev +``` + +Then in the container: + +```shell +$ pytest -v -s --data_path=/test_data/ /port_dev/tests --which_modules= +``` +The 'stencil name' can be determined from the associated Translate class. e.g. TranslateXPPM is a test class that translate data serialized from a run of the fortran model, and 'XPPM' is the name you can use with --which_modules. + + + + +### Test options + +All of the make endpoints involved running tests can be prefixed with the `TEST_ARGS` environment variable to set test options or pytest CLI args (see below) when running inside the container. + +* `--which_modules ` - comma separated list of which modules to test (defaults to running all of them). + +* `--print_failures` - if your test fails, it will only report the first datapoint. If you want all the nonmatching regression data to print out (so you can see if there are patterns, e.g. just incorrect for the first 'i' or whatever'), this will print out for every failing test all the non-matching data. + +* `--failure_stride` - when printing failures, print every n failures only. + +* `--data_path` - path to where you have the `Generator*.dat` and `*.json` serialization regression data. Defaults to current directory. + +* `--backend` - which backend to use for the computation. Options: `[numpy, gt:cpu_ifirst, gt:cpu_first, gt:gpu, cuda]`. Defaults to `numpy`. +* `--python_regression` - Run the tests that have Python based regression data. Only applies to running parallel tests (savepoint_tests_mpi) +Pytest provides a lot of options, which you can see by `pytest --help`. Here are some +common options for our tests, which you can add to `TEST_ARGS`: + +* `-r` - is used to report test types other than failure. It can be provided `s` for skipped (e.g. tests which were not run because earlier tests of the same stencil failed), `x` for xfail or "expected to fail" tests (like tests with no translate class), or `p` for pass. For example, to report skipped and xfail tests you would use `-rsx`. + +* `--disable-warnings` - will stop all warnings from being printed at the end of the tests, for example warnings that translate classes are not yet implemented. + +* `-v` - will increase test verbosity, while `-q` will decrease it. + +* `-s` - will let stdout print directly to console instead of capturing the output and printing it when a test fails only. Note that logger lines will always be printed both during (by setting log_cli in our pytest.ini file) and after tests. + +* `-m` - will let you run only certain groups of tests. For example, `-m=parallel` will run only parallel stencils, while `-m=sequential` will run only stencils that operate on one rank at a time. + +* `--threshold_overrides_file` - will read a yaml file with error thresholds specified for specific backend and platform (docker or metal) configurations, overriding the max_error thresholds defined in the Translate classes. Format of the yaml file is described [here](tests/savepoint/translate/overrides/README.md). + +* `--dperiodic` - run tests on a doubly-periodic domain. Will look for only one tile's worth of test data and parallel tests will be run with a TileCommunicator instead of a CubedSphereCommunicator. + +**NOTE:** FV3 is current assumed to be by default in a "development mode", where stencils are checked each time they execute for code changes (which can trigger regeneration). This process is somewhat expensive, so there is an option to put FV3 in a performance mode by telling it that stencils should not automatically be rebuilt: + +```shell +$ export FV3_STENCIL_REBUILD_FLAG=False +``` + +## Porting a new stencil + +1. Find the location in the fv3gfs-fortran repo code where the save-point is to be added, e.g. using + +```shell +$ git grep +``` + +2. Create a `translate` class from the serialized save-point data to a call to the stencil or function that calls the relevant stencil(s). + +These are usually named `tests/savepoint/translate/translate_` + +Import this class in the `tests/savepoint/translate/__init__.py` file + +3. Write a Python function wrapper that the translate function (created above) calls. + +By convention, we name these `fv3core/stencils/.py` + +4. Run the test, either with one name or a comma-separated list + +```shell +$ make dev_tests TEST_ARGS="-–which_modules=" +``` + +**Please also review the [Porting conventions](#porting-conventions) section for additional explanation** +## Installation + +### Docker Image + +To build the `us.gcr.io/vcm-ml/fv3core` image with required dependencies for running the Python code, run + +```shell +$ make build +``` + +Add `PULL=False` to build from scratch without running `docker pull`: + +```shell +PULL=False make build +``` + +## Relevant repositories + +- https://github.com/GridTools/serialbox - + Serialbox generates serialized data when the Fortran model runs and has bindings to manage data from Python + +- https://github.com/VulcanClimateModeling/fv3gfs-fortran - + This is the existing Fortran model decorated with serialization statements from which the test data is generated + +- https://github.com/GridTools/gt4py - + Python package for the DSL language + +- https://github.com/VulcanClimateModeling/util + Python specific model functionality, such as halo updates. + +- https://github.com/VulcanClimateModeling/fv3gfs-wrapper + A Python based wrapper for running the Fortran version of the FV3GFS model. + +Some of these are submodules. +While tests can work without these, it may be necessary for development to have these as well. +To add these to the local repository, run + +```shell +$ git submodule update --init +``` + +The submodules include: + +- `external/util` - git@github.com:VulcanClimateModeling/util.git +- `external/daint_venv` - git@github.com:VulcanClimateModeling/daint_venv.git + +## Dockerfiles and building + +There are two main docker files: + +1. `docker/dependencies.Dockerfile` - defines dependency images such as for mpi, serialbox, and GT4py + +2. `docker/Dockerfile` - uses the dependencies to define the final fv3core images. + +The dependencies are separated out into their own images to expedite rebuilding the docker image without having to rebuild dependencies, especially on CI. + +For the commands below using `make -C docker`, you can alternatively run `make` from within the `docker` directory. + +These dependencies can be updated, pushed, and pulled with `make -C docker build_deps`, `make -C docker push_deps`, and `make -C docker pull_deps`. The tag of the dependencies is based on the tag of the current build in the Makefile, which we will expand on below. + +Building from scratch requires both a deps and build command, such as `make -C docker pull_deps fv3core_image`. + +If any example fails for "pulled dependencies", it means the dependencies have never been built. You can +build them and push them to GCR with: + +```shell +$ make -C docker build_deps push_deps +``` + +### Building examples + +fv3core image with pulled dependencies: + +```shell +$ make -C docker pull_deps fv3core_image +``` + +CUDA-enabled fv3core image with pulled dependencies: +``` +$ CUDA=y make -C docker pull_deps fv3core_image +``` + +fv3core image with locally-built dependencies: +```shell +$ make -C docker build_deps fv3core_image +``` + +### Updating Serialbox + +If you need to install an updated version of Serialbox, you must first install cmake into the development environment. To install an updated version of Serialbox from within the container run + +```shell +$ wget https://github.com/Kitware/CMake/releases/download/v3.17.3/cmake-3.17.3.tar.gz && \ + tar xzf cmake-3.17.3.tar.gz && \ + cd cmake-3.17.3 && \ + ./bootstrap && make -j4 && make install +$ git clone -b v2.6.1 --depth 1 https://github.com/GridTools/serialbox.git /tmp/serialbox +$ cd /tmp/serialbox +$ cmake -B build -S /tmp/serialbox -DSERIALBOX_USE_NETCDF=ON -DSERIALBOX_TESTING=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/serialbox +$ cmake --build build/ -j $(nproc) --target install +$ cd - +$ rm -rf build /tmp/serialbox +``` + +## Pinned dependencies + +Dependencies are pinned using `constraints.txt`. This is auto-generated by pip-compile from the `pip-tools` package, which reads `requirements.txt` and `requirements/requirements_lint.txt`, determines the latest versions of all dependencies (including recursive dependencies) compatible those files, and writes pinned versions for all dependencies. This can be updated using: + +```shell +$ make constraints.txt +``` + +This file is committed to the repository, and gives more reproducible tests if an old commit of the repository is checked out in the future. The constraints are followed when creating the `fv3core` docker images. To ensure consistency this should ideally be run from inside a docker development environment, but you can also run it on your local system with an appropriate Python 3 environment. + +## Development + +To develop fv3core, you need to install the linting requirements in `requirements/requirements_lint.txt`. To install the pinned versions, use: + +```shell +$ pip install -c constraints.txt -r requirements/requirements_lint.txt +``` + +This adds `pre-commit`, which we use to lint and enforce style on the code. The first time you install `pre-commit`, install its git hooks using: + +```shell +$ pre-commit install +pre-commit installed at .git/hooks/pre-commit +``` + +As a convenience, the `lint` target of the top-level makefile executes `pre-commit run --all-files`. +Linting, which formats files and checks for some style conventions, is required, as the same checks are the first step in the continuous integration testing that happens when creating a pull request. +Linting locally saves time and literal energy, since CI tests do not have to be launched so many times! + + Please see the 'Development Guidelines' below for more information on the structure of the code to align your new code with the current conventions, as well as the CONTRIBUTING.md document for style guidelines. + +## GT4Py version + +FV3Core does not actually use the [GridTools/gt4py](https://github.com/gridtools/gt4py) main, it instead uses a Vulcan Climate Modeling development branch. +This is publically available version at [VCM/gt4py](https://github.com/vulcanclimatemodeling/gt4py). + +Situation: There is a new stable feature in a gt4py PR, but it is not yet merged into the GridTools/gt4py main branch. +[branches.cfg](https://github.com/VulcanClimateModeling/gt4py/blob/develop/branches.cfg) lists these features. +Steps: + +1. Add any new branches to `branches.cfg` +2. Rebuild the develop branch, either: + a. `make_develop gt4py-dev path/to/branches.cfg` (you may have to resolve conflicts...) + b. Adding new commits on top of the existing develop branch (e.g. merge or cherry-pick) +3. Force push to the develop branch: `git push -f upstream develop` + +The last step will launch Jenkins tests. If these pass: + +1. Create a git tag: `git tag v-$(git rev-parse --short HEAD)` +2. Push the tag: `git push upstream --tags` +3. Make a PR to [VCM/gt4py](https://github.com/vulcanclimatemodeling/fv3core) that updates the version in `docker/Makefile` to the new tag. + +## License +FV3Core is provided under the terms of the [GPLv3](https://www.gnu.org/licenses/gpl-3.0.en.html) license. + +# Development guidelines + +## File structure / conventions +The main functionality of the FV3 dynamical core, which has been ported from the Fortran version in the fv3gfs-fortran repo, is defined using GT4py stencils and python 'compute' functions in fv3core/stencils. The core is comprised of units of calculations defined for regression testing. These were initially generally separated into distinct files in fv3core/stencils with corresponding files in tests/savepoint/translate/translate_.py defining the translation of variables from Fortran to Python. Exceptions exist in cases where topical and logical grouping allowed for code reuse. As refactors optimize the model, these units may be merged to occupy the same files and even methods/stencils, but the units should still be tested separately, unless determined to be redundant. + +The core has most of its calculations happening in GT4py stencils, but there are still several instances of operations happening in Python directly, which will need to be replaced with GT4py code for optimal performance. + +The namelist and grid are global variables defined in fv3core/_config.py The namelist is 'flattened' so that the grouping name of the option is not required to access the data (we may want to change this). + +The grid variables are mostly 2d variables and are 'global' to the model thread per mpi rank. The grid object also contains domain and layout information relevant to the current rank being operated on. + +Utility functions in `fv3core/utils/` include: + - `gt4py_utils.py`: + - default gt4py and model settings + - methods for generating gt4py storages + - methods for using numpy and cupy arrays in python functions that have not been put into GT4py + - methods for handling complex patterns that did not immediately map to gt4py, and will mostly be removed with future refactors (e.g. k_split_run) + - some general model math computations (e.g. great_circle_dist), that will eventually be put into gt4py with a future refactor + - `grid.py`: + - A Grid class definition that provides information about the grid layout, current tile informationm access to grid variables used globally, and convenience methods related to tile indexing, origins and domains commonly used + - A grid is defined for each MPI rank (minimum 6 ranks, 1 for each tile face of the cubed sphere grid represnting the whole Earth) + - Also provides functionality for generating a Quantity object used for halo updates and other utilities + - `corners`: port of corner calculations, initially direct Python calculations, being replaced with GT4py gtscript functions as the GT4py regions feature is implemented + - `mpi.py`: a wrapper for importing mpi4py when available + - `global_constants.py`: constants for use throughout the model + - `typing.py`: Clean names for common types we use in the model. This is new and + hasn't been adopted throughout the model yet, but will eventually be our + standard. A shorthand 'sd' has been used in the intial version. + +The `tests/` directory currently includes a framework for translating fields serialized (using +Serialbox from GridTools) from a Fortran run into gt4py storages that can be inputs to +fv3core unit computations, and compares the results of the ported code to serialized +data following a unit computation. + +The `docker/` directory provides Dockerfiles for building a repeatable environment in which +to run the core + +The `external/` directory is for submoduled repos that provide essential functionality + +The build system uses Makefiles following the convention of other repos within VulcanClimateModeling. + +## Model Interface + +The top level functions fv_dynamics and fv_sugridz can currenty only be run in parallel using mpi with a minimum of 6 ranks (there are a few other units that also require this, e.g. whenever there is a halo update involved in a unit). These are the interface to the rest of the model and currently have different conventions than the rest of the model. + - A 'state' object (currently a SimpleNamespace) stores pointers to the allocated data fields + - Most functions within dyn_core can be run sequentially per rank + - Currently a list of ArgSpecs must decorate an interface function, where each ArgSpec provides useful information about the argument, e.g.: `@state_inputs( ArgSpec("qvapor", "specific_humidity", "kg/kg", intent="inout")` + - The format is (fortran_name, long_name, units, intent) + - We currently provide a duplicate of most of the metadata in the specification of the unit test, but that may be removed eventually. + - Then the function itself, e.g. fv_dynamics, has arguments of 'state', 'comm' (the communicator) and all of the scalar parameters being provided. + +### Porting conventions + +Generation of regression data occurs in the fv3gfs-fortran repo (https://github.com/VulcanClimateModeling/fv3gfs-fortran) with serialization statements and a build procedure defined in `tests/serialized_test_data_generation`. The version of data this repo currently tests against is defined in `FORTRAN_SERIALIZED_DATA_VERSION` in this repo's `docker/Makefile.image_names`. Fields serialized are defined in Fortran code with serialization comment statements such as: + +``` + !$ser savepoint C_SW-In + !$ser data delpcd=delpc delpd=delp ptcd=ptc +``` + +where the name being assigned is the name the fv3core uses to identify the variable in the test code. When this name is not equal to the name of the variable, this was usually done to avoid conflicts with other parts of the code where the same name is used to reference a differently sized field. + +The majority of the logic for translating from data serialized from Fortran to something that can be used by Python, and the comparison of the results, is encompassed by the main Translate class in the tests/savepoint/translate/translate.py file. Any units not involving a halo update can be run using this framework, while those that need to be run in parallel can look to the ParallelTranslate class as the parent class in tests/savepoint/translate/parallel_translate.py. These parent classes provide generally useful operations for translating serialized data between Fortran and Python specifications, and for applying regression tests. + +A new unit test can be defined as a new child class of one of these, with a naming convention of `Translate` where `Savepoint Name` is the name used in the serialization statements in the Fortran code, without the `-In` and `-Out` part of the name. A translate class can usually be minimally specify the input and output fields. Then, in cases where the parent compute function is insuffient to handle the complexity of either the data translation or the compute function, the appropriate methods can be overridden. + +For Translate objects + - The init function establishes the assumed translation setup for the class, which can be dynamically overridden as needed. + - the parent compute function does: + - Makes gt4py storages of the max shape (grid.npx+1, grid.npy+1, grid.npz+1) aligning the data based on the start indices specified. (gt4py requires data fields have the same shape, so in this model we have buffer points so all calculations can be done easily without worrying about shape matching). + - runs the compute function (defined in self.compute_func) on the input data storages + - slices the computed Python fields to be compared to fortran regression data + - The unit test then uses a modified relative error metric to determine whether the unit passes + - The init method for a Translate class: + - The input (self.in_vars["data_vars"]) and output(self.out_vars) variables are specified in dictionaries, where the keys are the name of the variable used in the model and the values are dictionaries specifying metadata for translation of serialized data to gt4py storages. The metadata that can be specied to override defaults are: + - Indices to line up data arrays into gt4py storages (which all get created as the max possible size needed by all operations, for simplicity): "istart", "iend", "jstart", "jend", "kstart", "kend". These should be set using the 'grid' object available to the Translate object, using equivalent index names as in the declaration of variables in the Fortran code, e.g. real:: cx(bd%is:bd%ie+1,bd%jsd:bd%jed ) means we should assign. Example: + +```python + self.in_vars["data_vars"]["cx"] = {"istart": self.is\_, "iend": self.ie + 1, + "jstart": self.jsd, "jend": self.jed,} +``` + - There is only a limited set of Fortran shapes declared, so abstractions defined in the grid can also be used, + e.g.: `self.out_vars["cx"] = self.grid.x3d_compute_domain_y_dict()`. Note that the variables, e.g. `grid.is\_` and `grid.ie` specify the 'compute' domain in the x direction of the current tile, equivalent to `bd%is` and `bd%ie` in the Fortran model EXCEPT that the Python variables are local to the current MPI rank (a subset of the tile face), while the Fortran values are global to the tile face. This is because these indices are used to slice into fields, which in Python is 0-based, and in Fortran is based on however the variables are declared. But, for the purposes of aligning data for computations and comparisons, we can match them in this framework. Shapes need to be defined in a dictionary per variable including `"istart"`, `"iend"`, `"jstart"`, `"jend"`, `"kstart"`, `"kend"` that represent the shape of that variable as defined in the Fortran code. The default shape assumed if a variable is specified with an empty dictionary is `isd:ied, jsd:jed, 0:npz - 1` inclusive, and variables that aren't that shape in the Fortran code need to have the 'start' indices specified for the in_vars dictionary , and 'start' and 'end' for the out_vars. + - `"serialname"` can be used to specify a name used in the Fortran code declaration if we'd like the model to use a different name + - `"kaxis"`: which dimension is the vertical direction. For most variables this is '2' and does not need to be specified. For Fortran variables that assign the vertical dimension to a different axis, this can be set to ensure we end up with 3d storages that have the vertical dimension where it is expected by GT4py. + - `"dummy_axes"`: If set this will set of the storage to have singleton dimensions in the axes defined. This is to enable testing stencils where the full 3d data has not been collected and we want to run stencil tests on the data for a particular slice. + - `"names_4d"`: If a 4d variable is being serialized, this can be set to specify the names of each 3d field. By default this is the list of tracers. + - input variables that are scalars should be added to `self.in_vars["parameters"]` + - `self.compute_func` is the name of the model function that should be run by the compute method in the translate class + - `self.max_error` overrides the parent classes relative error threshold. This should only be changed when the reasons for non-bit reproducibility are understood. + - `self.max_shape` sets the size of the gt4py storage created for testing + - `self.ignore_near_zero_errors[] = True`: This is an option to let some fields pass with higher relative error if the absolute error is very small + - `self.skip_test`: This is an option to jump over the test case, to be used in the override file for temporary deactivation of tests. + +For `ParallelTranslate` objects: + - Inputs and outputs are defined at the class level, and these include metadata such as the "name" (e.g. understandable name for the symbol), dimensions, units and n_halo(numb er of halo lines) + - Both `compute_sequential` and `compute_parallel` methods may be defined, where a mock communicator is used in the `compute_sequential` case + - The parent assumes a state object for tracking fields and methods exist for translating from inputs to a state object and extracting the output variables from the state. It is assumed that Quantity objects are needed in the model method in order to do halo updates. + - `ParallelTranslate2Py` is a slight variation of this used for many of the parallel units that do not yet utilize a state object and relies on the specification of the same index metadata of the Translate classes + - `ParallelTranslateBaseSlicing` makes use of the state but relies on the Translate object of self._base, a Translate class object, to align the data before making quantities, computing and comparing. + +### Debugging Tests + +Pytest can be configured to give you a pdb session when a test fails. To route this properly through docker, you can run: + +```bash +TEST_ARGS="-v -s --pdb" RUN_FLAGS="--rm -it" make tests +``` + +This can be done with any pytest target, such as `make savepoint_tests` and `make savepoint_tests_mpi`. + +### GEOS API + +The `GeosDycoreWrapper` class provides an API to run the dynamical core in a Python component of a GEOS model run. A `GeosDycoreWrapper` object is initialized with a namelist, communicator, and backend, which creates the communicators, partitioners, dycore state, and dycore object required to run the Pace dycore. A wrapper object takes numpy arrays of `u, v, w, delz, pt, delp, q, ps, pe, pk, peln, pkz, phis, q_con, omga, ua, va, uc, vc, mfxd, mfyd, cxd, cyd,` and `diss_estd` and returns a dictionary containing numpy arrays of those same variables. Wrapper objects contain a `timer` attrubite that tracks the amount of time moving input data to the dycore state, running the dynamical core, and retrieving the data from the state. diff --git a/README.md b/README.md index a6192aac..6b516c62 100644 --- a/README.md +++ b/README.md @@ -1,459 +1,36 @@ -> DISCLAIMER: Work in progress +# TEMPORARY BRANCH: Up-skilling to GEOS v11.4.2 -# FV3core +This branch exists solely for up-skilling pyFV3 to be able run GEOS in it's v11.4.2 FP configuration. +The need for a seperate branch from `develop` rely in the following differences: -FV3core is a Python version, using GridTools GT4Py with CPU and GPU backend options, of the FV3 dynamical core (fv3gfs-fortran repo). -The code here includes regression test data of computation units coming from serialized output from the Fortran model generated using the `GridTools/serialbox` framework. +- GEOS run a 32bit floating point precision version (with appropriate 64bit buffers for mass conservation). This means the translate test requires a new set of data _and_ will not pass on old 8.1.3 Pace data. +- GEOS requires options that are deemed "legacy" and that we may want to replace rather than port +- Project requirements demand quick iterative development, while `pyFV3` demands concertation between all stakeholders. -As of January 10, 2021 this documentation is outdated in that it was written when we had fv3core as its own single repository. Some functionality, such as linting, has been moved to the top level but may still be described in this document as occuring inside the fv3core folder. +The aim is to validate and benchmark GEOS v11.4.2 with this dynamics. Once done, we will aim to move _as much code as possible_ back into develop. +The methodology goes as follows -**WARNING** This repo is under active development and relies on code and data that is not publicly available at this point. +- Merge directly into `develop` any changes that do not demand a new set of data +- Keep track of the feature branch (below) that can't be merged in `develop` for future PR +- Keep track of GEOS vs SHiELD differences for future discussions -## QuickStart +## Feature branches -1. Ensure you have docker installed and available for building and running and has access to the VCM cloud +Legend: -Be sure to complete any required post-installation instructions (e.g. [for linux](https://docs.docker.com/engine/install/linux-postinstall/)). Also [authorize Docker to pull from gcr](https://cloud.google.com/container-registry/docs/advanced-authentication). Your user will need to have read access to the `us.gcr.io/vcm-ml` repository. +- ⚙️ _GEOS - WIP_ : Ongoing work - can be merged temporarily +- 🔶 _GEOS - Merged_: Considered done - merged in GEOS v11.4.2 branch but NOT in `develop` +- ✅ _Develop - Merged_: Work done as part of up-skilling done for GEOS merged in `develop` AND the GEOS v11.4.2 branch. -2. You can build the image, download the data, and run the tests using: +Branches: -```shell -$ make tests savepoint_tests savepoint_tests_mpi -``` - -If you want to develop code, you should also install the linting requirements and git hooks locally - -```shell -$ pip install -c constraints.txt -r requirements/requirements_lint.txt -$ pre-commit install - -## Getting started, in more detail -If you want to build the main fv3core docker image, run - -```shell -$ make build -``` - -If you want to download test data run - -```shell -$ make get_test_data -``` - -And the c12_6ranks_standard data will download into the `test_data` directory. - -If you do not have a GCP account, there is an option to download basic test data from a public FTP server and you can skip the GCP authentication step above. To download test data from the FTP server, use `make USE_FTP=yes get_test_data` instead and this will avoid fetching from a GCP storage bucket. You will need a valid in stallation of the `lftp` command. - -MPI parallel tests (that run that way to exercise halo updates in the model) can also be run with: - -```shell -$ make savepoint_tests_mpi -``` - -The environment image that the fv3core container uses is prebuilt and lives in the GCR. The above commands will by default pull this image before building the fv3core image and running the tests. -To build the environment from scratch (including GT4py) before running tests, either run - -``` -make build_environment -``` - -or - -```shell -$ PULL=False make savepoint_tests -``` - -which will execute the target `build_environment` for you before running the tests. - -There are `push_environment` and `rebuild_environment` targets, but these should normally not be done manually. Updating the install image should only be done by Jenkins after the tests pass using a new environment. - -### Test data options - -If you want to run different test data, discover the possible options with -```shell -$ make list_test_data_options -``` -This will list the storage buckets in the cloud. Then to run one of them, set EXPERIMENT to the folder name of the data you'd like to use: - -e.g. -```shell -$EXPERIMENT=c48_6ranks_standard make tests -``` - -If you choose an experiment with a different number of ranks than 6, also set `NUM_RANKS=` - -## Testing interactively outside the container - -After `make savepoint_tests` has been run at least once (or you have data in test_data and the docker image fv3core exists because `make build` has been run), you can iterate on code changes using - -```shell -$ DEV=y make savepoint_tests -``` -or for the parallel or non-savepoint tests: - -```shell -$ DEV=y make tests savepoint_tests_mpi -``` -These will mount your current code into the fv3core container and run it rather than the code that was built when `make build` ran. - -## Running tests inside a container - -If you to prefer to work interactively inside the fv3core container, get the test data and build the docker image (see above if you do not have a GCP account and want to get test data): -```shell -$ make get_test_data -``` - -```shell -$ make build -``` -Testing can be run with this data from `/port_dev` inside the container: - -```shell -$ make dev -``` - -Then in the container: - -```shell -$ pytest -v -s --data_path=/test_data/ /port_dev/tests --which_modules= -``` -The 'stencil name' can be determined from the associated Translate class. e.g. TranslateXPPM is a test class that translate data serialized from a run of the fortran model, and 'XPPM' is the name you can use with --which_modules. - - - - -### Test options - -All of the make endpoints involved running tests can be prefixed with the `TEST_ARGS` environment variable to set test options or pytest CLI args (see below) when running inside the container. - -* `--which_modules ` - comma separated list of which modules to test (defaults to running all of them). - -* `--print_failures` - if your test fails, it will only report the first datapoint. If you want all the nonmatching regression data to print out (so you can see if there are patterns, e.g. just incorrect for the first 'i' or whatever'), this will print out for every failing test all the non-matching data. - -* `--failure_stride` - when printing failures, print every n failures only. - -* `--data_path` - path to where you have the `Generator*.dat` and `*.json` serialization regression data. Defaults to current directory. - -* `--backend` - which backend to use for the computation. Options: `[numpy, gt:cpu_ifirst, gt:cpu_first, gt:gpu, cuda]`. Defaults to `numpy`. -* `--python_regression` - Run the tests that have Python based regression data. Only applies to running parallel tests (savepoint_tests_mpi) -Pytest provides a lot of options, which you can see by `pytest --help`. Here are some -common options for our tests, which you can add to `TEST_ARGS`: - -* `-r` - is used to report test types other than failure. It can be provided `s` for skipped (e.g. tests which were not run because earlier tests of the same stencil failed), `x` for xfail or "expected to fail" tests (like tests with no translate class), or `p` for pass. For example, to report skipped and xfail tests you would use `-rsx`. - -* `--disable-warnings` - will stop all warnings from being printed at the end of the tests, for example warnings that translate classes are not yet implemented. - -* `-v` - will increase test verbosity, while `-q` will decrease it. - -* `-s` - will let stdout print directly to console instead of capturing the output and printing it when a test fails only. Note that logger lines will always be printed both during (by setting log_cli in our pytest.ini file) and after tests. - -* `-m` - will let you run only certain groups of tests. For example, `-m=parallel` will run only parallel stencils, while `-m=sequential` will run only stencils that operate on one rank at a time. - -* `--threshold_overrides_file` - will read a yaml file with error thresholds specified for specific backend and platform (docker or metal) configurations, overriding the max_error thresholds defined in the Translate classes. Format of the yaml file is described [here](tests/savepoint/translate/overrides/README.md). - -* `--dperiodic` - run tests on a doubly-periodic domain. Will look for only one tile's worth of test data and parallel tests will be run with a TileCommunicator instead of a CubedSphereCommunicator. - -**NOTE:** FV3 is current assumed to be by default in a "development mode", where stencils are checked each time they execute for code changes (which can trigger regeneration). This process is somewhat expensive, so there is an option to put FV3 in a performance mode by telling it that stencils should not automatically be rebuilt: - -```shell -$ export FV3_STENCIL_REBUILD_FLAG=False -``` - -## Porting a new stencil - -1. Find the location in the fv3gfs-fortran repo code where the save-point is to be added, e.g. using - -```shell -$ git grep -``` - -2. Create a `translate` class from the serialized save-point data to a call to the stencil or function that calls the relevant stencil(s). - -These are usually named `tests/savepoint/translate/translate_` - -Import this class in the `tests/savepoint/translate/__init__.py` file - -3. Write a Python function wrapper that the translate function (created above) calls. - -By convention, we name these `fv3core/stencils/.py` - -4. Run the test, either with one name or a comma-separated list - -```shell -$ make dev_tests TEST_ARGS="-–which_modules=" -``` - -**Please also review the [Porting conventions](#porting-conventions) section for additional explanation** -## Installation - -### Docker Image - -To build the `us.gcr.io/vcm-ml/fv3core` image with required dependencies for running the Python code, run - -```shell -$ make build -``` - -Add `PULL=False` to build from scratch without running `docker pull`: - -```shell -PULL=False make build -``` - -## Relevant repositories - -- https://github.com/GridTools/serialbox - - Serialbox generates serialized data when the Fortran model runs and has bindings to manage data from Python - -- https://github.com/VulcanClimateModeling/fv3gfs-fortran - - This is the existing Fortran model decorated with serialization statements from which the test data is generated - -- https://github.com/GridTools/gt4py - - Python package for the DSL language - -- https://github.com/VulcanClimateModeling/util - Python specific model functionality, such as halo updates. - -- https://github.com/VulcanClimateModeling/fv3gfs-wrapper - A Python based wrapper for running the Fortran version of the FV3GFS model. - -Some of these are submodules. -While tests can work without these, it may be necessary for development to have these as well. -To add these to the local repository, run - -```shell -$ git submodule update --init -``` - -The submodules include: - -- `external/util` - git@github.com:VulcanClimateModeling/util.git -- `external/daint_venv` - git@github.com:VulcanClimateModeling/daint_venv.git - -## Dockerfiles and building - -There are two main docker files: - -1. `docker/dependencies.Dockerfile` - defines dependency images such as for mpi, serialbox, and GT4py - -2. `docker/Dockerfile` - uses the dependencies to define the final fv3core images. - -The dependencies are separated out into their own images to expedite rebuilding the docker image without having to rebuild dependencies, especially on CI. - -For the commands below using `make -C docker`, you can alternatively run `make` from within the `docker` directory. - -These dependencies can be updated, pushed, and pulled with `make -C docker build_deps`, `make -C docker push_deps`, and `make -C docker pull_deps`. The tag of the dependencies is based on the tag of the current build in the Makefile, which we will expand on below. - -Building from scratch requires both a deps and build command, such as `make -C docker pull_deps fv3core_image`. - -If any example fails for "pulled dependencies", it means the dependencies have never been built. You can -build them and push them to GCR with: - -```shell -$ make -C docker build_deps push_deps -``` - -### Building examples - -fv3core image with pulled dependencies: - -```shell -$ make -C docker pull_deps fv3core_image -``` - -CUDA-enabled fv3core image with pulled dependencies: -``` -$ CUDA=y make -C docker pull_deps fv3core_image -``` - -fv3core image with locally-built dependencies: -```shell -$ make -C docker build_deps fv3core_image -``` - -### Updating Serialbox - -If you need to install an updated version of Serialbox, you must first install cmake into the development environment. To install an updated version of Serialbox from within the container run - -```shell -$ wget https://github.com/Kitware/CMake/releases/download/v3.17.3/cmake-3.17.3.tar.gz && \ - tar xzf cmake-3.17.3.tar.gz && \ - cd cmake-3.17.3 && \ - ./bootstrap && make -j4 && make install -$ git clone -b v2.6.1 --depth 1 https://github.com/GridTools/serialbox.git /tmp/serialbox -$ cd /tmp/serialbox -$ cmake -B build -S /tmp/serialbox -DSERIALBOX_USE_NETCDF=ON -DSERIALBOX_TESTING=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/serialbox -$ cmake --build build/ -j $(nproc) --target install -$ cd - -$ rm -rf build /tmp/serialbox -``` - -## Pinned dependencies - -Dependencies are pinned using `constraints.txt`. This is auto-generated by pip-compile from the `pip-tools` package, which reads `requirements.txt` and `requirements/requirements_lint.txt`, determines the latest versions of all dependencies (including recursive dependencies) compatible those files, and writes pinned versions for all dependencies. This can be updated using: - -```shell -$ make constraints.txt -``` - -This file is committed to the repository, and gives more reproducible tests if an old commit of the repository is checked out in the future. The constraints are followed when creating the `fv3core` docker images. To ensure consistency this should ideally be run from inside a docker development environment, but you can also run it on your local system with an appropriate Python 3 environment. - -## Development - -To develop fv3core, you need to install the linting requirements in `requirements/requirements_lint.txt`. To install the pinned versions, use: - -```shell -$ pip install -c constraints.txt -r requirements/requirements_lint.txt -``` - -This adds `pre-commit`, which we use to lint and enforce style on the code. The first time you install `pre-commit`, install its git hooks using: - -```shell -$ pre-commit install -pre-commit installed at .git/hooks/pre-commit -``` - -As a convenience, the `lint` target of the top-level makefile executes `pre-commit run --all-files`. -Linting, which formats files and checks for some style conventions, is required, as the same checks are the first step in the continuous integration testing that happens when creating a pull request. -Linting locally saves time and literal energy, since CI tests do not have to be launched so many times! - - Please see the 'Development Guidelines' below for more information on the structure of the code to align your new code with the current conventions, as well as the CONTRIBUTING.md document for style guidelines. - -## GT4Py version - -FV3Core does not actually use the [GridTools/gt4py](https://github.com/gridtools/gt4py) main, it instead uses a Vulcan Climate Modeling development branch. -This is publically available version at [VCM/gt4py](https://github.com/vulcanclimatemodeling/gt4py). - -Situation: There is a new stable feature in a gt4py PR, but it is not yet merged into the GridTools/gt4py main branch. -[branches.cfg](https://github.com/VulcanClimateModeling/gt4py/blob/develop/branches.cfg) lists these features. -Steps: - -1. Add any new branches to `branches.cfg` -2. Rebuild the develop branch, either: - a. `make_develop gt4py-dev path/to/branches.cfg` (you may have to resolve conflicts...) - b. Adding new commits on top of the existing develop branch (e.g. merge or cherry-pick) -3. Force push to the develop branch: `git push -f upstream develop` - -The last step will launch Jenkins tests. If these pass: - -1. Create a git tag: `git tag v-$(git rev-parse --short HEAD)` -2. Push the tag: `git push upstream --tags` -3. Make a PR to [VCM/gt4py](https://github.com/vulcanclimatemodeling/fv3core) that updates the version in `docker/Makefile` to the new tag. - -## License -FV3Core is provided under the terms of the [GPLv3](https://www.gnu.org/licenses/gpl-3.0.en.html) license. - -# Development guidelines - -## File structure / conventions -The main functionality of the FV3 dynamical core, which has been ported from the Fortran version in the fv3gfs-fortran repo, is defined using GT4py stencils and python 'compute' functions in fv3core/stencils. The core is comprised of units of calculations defined for regression testing. These were initially generally separated into distinct files in fv3core/stencils with corresponding files in tests/savepoint/translate/translate_.py defining the translation of variables from Fortran to Python. Exceptions exist in cases where topical and logical grouping allowed for code reuse. As refactors optimize the model, these units may be merged to occupy the same files and even methods/stencils, but the units should still be tested separately, unless determined to be redundant. - -The core has most of its calculations happening in GT4py stencils, but there are still several instances of operations happening in Python directly, which will need to be replaced with GT4py code for optimal performance. - -The namelist and grid are global variables defined in fv3core/_config.py The namelist is 'flattened' so that the grouping name of the option is not required to access the data (we may want to change this). - -The grid variables are mostly 2d variables and are 'global' to the model thread per mpi rank. The grid object also contains domain and layout information relevant to the current rank being operated on. - -Utility functions in `fv3core/utils/` include: - - `gt4py_utils.py`: - - default gt4py and model settings - - methods for generating gt4py storages - - methods for using numpy and cupy arrays in python functions that have not been put into GT4py - - methods for handling complex patterns that did not immediately map to gt4py, and will mostly be removed with future refactors (e.g. k_split_run) - - some general model math computations (e.g. great_circle_dist), that will eventually be put into gt4py with a future refactor - - `grid.py`: - - A Grid class definition that provides information about the grid layout, current tile informationm access to grid variables used globally, and convenience methods related to tile indexing, origins and domains commonly used - - A grid is defined for each MPI rank (minimum 6 ranks, 1 for each tile face of the cubed sphere grid represnting the whole Earth) - - Also provides functionality for generating a Quantity object used for halo updates and other utilities - - `corners`: port of corner calculations, initially direct Python calculations, being replaced with GT4py gtscript functions as the GT4py regions feature is implemented - - `mpi.py`: a wrapper for importing mpi4py when available - - `global_constants.py`: constants for use throughout the model - - `typing.py`: Clean names for common types we use in the model. This is new and - hasn't been adopted throughout the model yet, but will eventually be our - standard. A shorthand 'sd' has been used in the intial version. - -The `tests/` directory currently includes a framework for translating fields serialized (using -Serialbox from GridTools) from a Fortran run into gt4py storages that can be inputs to -fv3core unit computations, and compares the results of the ported code to serialized -data following a unit computation. - -The `docker/` directory provides Dockerfiles for building a repeatable environment in which -to run the core - -The `external/` directory is for submoduled repos that provide essential functionality - -The build system uses Makefiles following the convention of other repos within VulcanClimateModeling. - -## Model Interface - -The top level functions fv_dynamics and fv_sugridz can currenty only be run in parallel using mpi with a minimum of 6 ranks (there are a few other units that also require this, e.g. whenever there is a halo update involved in a unit). These are the interface to the rest of the model and currently have different conventions than the rest of the model. - - A 'state' object (currently a SimpleNamespace) stores pointers to the allocated data fields - - Most functions within dyn_core can be run sequentially per rank - - Currently a list of ArgSpecs must decorate an interface function, where each ArgSpec provides useful information about the argument, e.g.: `@state_inputs( ArgSpec("qvapor", "specific_humidity", "kg/kg", intent="inout")` - - The format is (fortran_name, long_name, units, intent) - - We currently provide a duplicate of most of the metadata in the specification of the unit test, but that may be removed eventually. - - Then the function itself, e.g. fv_dynamics, has arguments of 'state', 'comm' (the communicator) and all of the scalar parameters being provided. - -### Porting conventions - -Generation of regression data occurs in the fv3gfs-fortran repo (https://github.com/VulcanClimateModeling/fv3gfs-fortran) with serialization statements and a build procedure defined in `tests/serialized_test_data_generation`. The version of data this repo currently tests against is defined in `FORTRAN_SERIALIZED_DATA_VERSION` in this repo's `docker/Makefile.image_names`. Fields serialized are defined in Fortran code with serialization comment statements such as: - -``` - !$ser savepoint C_SW-In - !$ser data delpcd=delpc delpd=delp ptcd=ptc -``` - -where the name being assigned is the name the fv3core uses to identify the variable in the test code. When this name is not equal to the name of the variable, this was usually done to avoid conflicts with other parts of the code where the same name is used to reference a differently sized field. - -The majority of the logic for translating from data serialized from Fortran to something that can be used by Python, and the comparison of the results, is encompassed by the main Translate class in the tests/savepoint/translate/translate.py file. Any units not involving a halo update can be run using this framework, while those that need to be run in parallel can look to the ParallelTranslate class as the parent class in tests/savepoint/translate/parallel_translate.py. These parent classes provide generally useful operations for translating serialized data between Fortran and Python specifications, and for applying regression tests. - -A new unit test can be defined as a new child class of one of these, with a naming convention of `Translate` where `Savepoint Name` is the name used in the serialization statements in the Fortran code, without the `-In` and `-Out` part of the name. A translate class can usually be minimally specify the input and output fields. Then, in cases where the parent compute function is insuffient to handle the complexity of either the data translation or the compute function, the appropriate methods can be overridden. - -For Translate objects - - The init function establishes the assumed translation setup for the class, which can be dynamically overridden as needed. - - the parent compute function does: - - Makes gt4py storages of the max shape (grid.npx+1, grid.npy+1, grid.npz+1) aligning the data based on the start indices specified. (gt4py requires data fields have the same shape, so in this model we have buffer points so all calculations can be done easily without worrying about shape matching). - - runs the compute function (defined in self.compute_func) on the input data storages - - slices the computed Python fields to be compared to fortran regression data - - The unit test then uses a modified relative error metric to determine whether the unit passes - - The init method for a Translate class: - - The input (self.in_vars["data_vars"]) and output(self.out_vars) variables are specified in dictionaries, where the keys are the name of the variable used in the model and the values are dictionaries specifying metadata for translation of serialized data to gt4py storages. The metadata that can be specied to override defaults are: - - Indices to line up data arrays into gt4py storages (which all get created as the max possible size needed by all operations, for simplicity): "istart", "iend", "jstart", "jend", "kstart", "kend". These should be set using the 'grid' object available to the Translate object, using equivalent index names as in the declaration of variables in the Fortran code, e.g. real:: cx(bd%is:bd%ie+1,bd%jsd:bd%jed ) means we should assign. Example: - -```python - self.in_vars["data_vars"]["cx"] = {"istart": self.is\_, "iend": self.ie + 1, - "jstart": self.jsd, "jend": self.jed,} -``` - - There is only a limited set of Fortran shapes declared, so abstractions defined in the grid can also be used, - e.g.: `self.out_vars["cx"] = self.grid.x3d_compute_domain_y_dict()`. Note that the variables, e.g. `grid.is\_` and `grid.ie` specify the 'compute' domain in the x direction of the current tile, equivalent to `bd%is` and `bd%ie` in the Fortran model EXCEPT that the Python variables are local to the current MPI rank (a subset of the tile face), while the Fortran values are global to the tile face. This is because these indices are used to slice into fields, which in Python is 0-based, and in Fortran is based on however the variables are declared. But, for the purposes of aligning data for computations and comparisons, we can match them in this framework. Shapes need to be defined in a dictionary per variable including `"istart"`, `"iend"`, `"jstart"`, `"jend"`, `"kstart"`, `"kend"` that represent the shape of that variable as defined in the Fortran code. The default shape assumed if a variable is specified with an empty dictionary is `isd:ied, jsd:jed, 0:npz - 1` inclusive, and variables that aren't that shape in the Fortran code need to have the 'start' indices specified for the in_vars dictionary , and 'start' and 'end' for the out_vars. - - `"serialname"` can be used to specify a name used in the Fortran code declaration if we'd like the model to use a different name - - `"kaxis"`: which dimension is the vertical direction. For most variables this is '2' and does not need to be specified. For Fortran variables that assign the vertical dimension to a different axis, this can be set to ensure we end up with 3d storages that have the vertical dimension where it is expected by GT4py. - - `"dummy_axes"`: If set this will set of the storage to have singleton dimensions in the axes defined. This is to enable testing stencils where the full 3d data has not been collected and we want to run stencil tests on the data for a particular slice. - - `"names_4d"`: If a 4d variable is being serialized, this can be set to specify the names of each 3d field. By default this is the list of tracers. - - input variables that are scalars should be added to `self.in_vars["parameters"]` - - `self.compute_func` is the name of the model function that should be run by the compute method in the translate class - - `self.max_error` overrides the parent classes relative error threshold. This should only be changed when the reasons for non-bit reproducibility are understood. - - `self.max_shape` sets the size of the gt4py storage created for testing - - `self.ignore_near_zero_errors[] = True`: This is an option to let some fields pass with higher relative error if the absolute error is very small - - `self.skip_test`: This is an option to jump over the test case, to be used in the override file for temporary deactivation of tests. - -For `ParallelTranslate` objects: - - Inputs and outputs are defined at the class level, and these include metadata such as the "name" (e.g. understandable name for the symbol), dimensions, units and n_halo(numb er of halo lines) - - Both `compute_sequential` and `compute_parallel` methods may be defined, where a mock communicator is used in the `compute_sequential` case - - The parent assumes a state object for tracking fields and methods exist for translating from inputs to a state object and extracting the output variables from the state. It is assumed that Quantity objects are needed in the model method in order to do halo updates. - - `ParallelTranslate2Py` is a slight variation of this used for many of the parallel units that do not yet utilize a state object and relies on the specification of the same index metadata of the Translate classes - - `ParallelTranslateBaseSlicing` makes use of the state but relies on the Translate object of self._base, a Translate class object, to align the data before making quantities, computing and comparing. - -### Debugging Tests - -Pytest can be configured to give you a pdb session when a test fails. To route this properly through docker, you can run: - -```bash -TEST_ARGS="-v -s --pdb" RUN_FLAGS="--rm -it" make tests -``` - -This can be done with any pytest target, such as `make savepoint_tests` and `make savepoint_tests_mpi`. - -### GEOS API - -The `GeosDycoreWrapper` class provides an API to run the dynamical core in a Python component of a GEOS model run. A `GeosDycoreWrapper` object is initialized with a namelist, communicator, and backend, which creates the communicators, partitioners, dycore state, and dycore object required to run the Pace dycore. A wrapper object takes numpy arrays of `u, v, w, delz, pt, delp, q, ps, pe, pk, peln, pkz, phis, q_con, omga, ua, va, uc, vc, mfxd, mfyd, cxd, cyd,` and `diss_estd` and returns a dictionary containing numpy arrays of those same variables. Wrapper objects contain a `timer` attrubite that tracks the amount of time moving input data to the dycore state, running the dynamical core, and retrieving the data from the state. +- ✅ `fix/F32/UpdateDzC`@Florian: Fix for fluxes gradient +- ✅ `fix/F32/DivergenceDamping`@Florian: Fix for 32-bit scalars in DivergenceDamping +- 🔶 `fix/RayleighDamping_mixed_precision`@Florian: fix the Ray_Fast test +- 🔶 `GEOS_update/yppm_xppm`@Florian: fix the YPPM/XPPM with `hord = -6` +- 🔶 `fix/DelnFlux_f32_support`@Florian: Fix for f32 support for DelnFlux (partial pass) +- 🔶 `fix/F32/UpdateDzD`@Florian: Fix for fluxes gradient & python computation +- ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation and column calculation (partial pass) +- ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 +- ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@FlorianL: Fix for 32-bit A2B_Ord4 +- ⚙️ `fix/GEOSv11_4_2/C_SW`: Fix for C_SW for 32-bit From 6a46e053f038b36b75a4179fe119adfc9153b93c Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Dec 2024 16:19:39 -0500 Subject: [PATCH 044/252] Tracer class to handle all things tracers. --- pyFV3/tracers.py | 127 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 127 insertions(+) create mode 100644 pyFV3/tracers.py diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py new file mode 100644 index 00000000..301b1f6b --- /dev/null +++ b/pyFV3/tracers.py @@ -0,0 +1,127 @@ +from __future__ import annotations +from ndsl import Quantity, QuantityFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from typing import Dict, List +import numpy as np + +# FOR REFERENCE - previous descriptive of the tracers, lining up with Pace work +# tracer_variables = [ +# "qvapor", +# "qliquid", +# "qrain", +# "qice", +# "qsnow", +# "qgraupel", +# "qo3mr", +# "qsgs_tke", +# "qcld", +# ] + + +class Tracers: + def __init__(self, factory: QuantityFactory) -> None: + self._quantities: Dict[str, Quantity] = {} + self._quantity_factory = factory + self._dims = [X_DIM, Y_DIM, Z_DIM] + + def copy_tracer_data( + self, + name: str, + data: np.ndarray, + units="unknown", + ): + qty = self._quantity_factory.empty(dims=self._dims, units=units) + if data.shape > qty.data.shape: + raise ValueError( + f"[pyFV3] Tracer {name} size ({data.shape}" + f" is bigger than grid {qty.data.shape})" + ) + qty.data[: data.shape[0], : data.shape[1], : data.shape[2]] = data + self._quantities[name] = qty + + @property + def count(self) -> int: + return len(self._quantities) + + def values(self): + return self._quantities.values() + + def names(self): + return self._quantities.keys() + + def items(self): + return self._quantities.items() + + def as_fortran_4D(self) -> np.ndarray: + shape = self._quantity_factory.sizer.get_shape(self._dims) + var4d = np.empty( + ( + shape[0] - 1, + shape[1] - 1, + shape[2] - 1, + self.count, + ) + ) + for idx, q in enumerate(self.values()): + var4d[:, :, :, idx] = q.data[:-1, :-1, :-1] + return var4d + + def as_dict(self) -> Dict[str, Quantity]: + return self._quantities + + def __getitem__(self, key): + return self._quantities[key] + + def __setitem__(self, key, value): + self._quantities[key] = value + + def __str__(self) -> str: + return self.__repr__() + + def __repr__(self) -> str: + msg = f"[pyFV3] {self.count} Tracers:\n" + for q in self.names(): + msg += f" {q}\n" + return msg + + @classmethod + def make( + cls, + quantity_factory: QuantityFactory, + tracer_mapping: List[str], + ): + tracers = cls(quantity_factory) + for name in tracer_mapping: + qty = quantity_factory.empty(dims=tracers._dims, units="kg/m^2") + tracers._quantities[name] = qty + return tracers + + @classmethod + def make_from_fortran( + cls, + quantity_factory: QuantityFactory, + tracer_mapping: List[str], + tracer_data: np.ndarray, + ) -> Tracers: + if len(tracer_data.shape) != 4: + raise ValueError("Expected 4D field as input") + count = len(tracer_mapping) + if count > tracer_data.shape[3]: + raise ValueError( + f"Mapping size {len(tracer_mapping)} is bigger than" + f" data dimensionality {tracer_data.shape[3]}" + ) + tracers = cls(quantity_factory) + for idx in range(0, count): + tracers.copy_tracer_data( + name=tracer_mapping[idx] or f"Tracer_{idx}", + data=tracer_data[:, :, :, idx], + units="kg/m^2", + ) + return tracers + + @staticmethod + def make_mapping(tracer_data: np.ndarray): + if len(tracer_data.shape) != 4: + raise ValueError("Expected 4D field as input") + return tracer_data.shape[3] * [None] From 2542184b712dcacbafba7200122b9930ac485f1f Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Dec 2024 16:21:51 -0500 Subject: [PATCH 045/252] Move Tracer2D1L (and test) to Tracers --- pyFV3/stencils/tracer_2d_1l.py | 13 +++--- .../translate/translate_tracer2d1l.py | 41 ++++++++++++++----- 2 files changed, 38 insertions(+), 16 deletions(-) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index 27514a37..e90600e4 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -22,6 +22,7 @@ from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from ndsl.typing import Communicator from pyFV3.stencils.fvtp2d import FiniteVolumeTransport +from pyFV3.tracers import Tracers @gtscript.function @@ -192,7 +193,7 @@ def __init__( transport: FiniteVolumeTransport, grid_data, comm: Communicator, - tracers: Dict[str, Quantity], + tracers: Tracers, ): orchestrate( obj=self, @@ -201,7 +202,7 @@ def __init__( ) grid_indexing = stencil_factory.grid_indexing self.grid_indexing = grid_indexing # needed for selective validation - self._tracer_count = len(tracers) + self._tracer_count = tracers.count self.grid_data = grid_data self._x_area_flux = quantity_factory.zeros( @@ -283,14 +284,14 @@ def __init__( dtype=Float, ) self._tracers_halo_updater = WrappedHaloUpdater( - comm.get_scalar_halo_updater([tracer_halo_spec] * self._tracer_count), - tracers, - [t for t in tracers.keys()], + comm.get_scalar_halo_updater([tracer_halo_spec] * tracers.count), + tracers.as_dict(), + [t for t in tracers.names()], ) def __call__( self, - tracers: Dict[str, Quantity], + tracers: Tracers, dp1, x_mass_flux, y_mass_flux, diff --git a/tests/savepoint/translate/translate_tracer2d1l.py b/tests/savepoint/translate/translate_tracer2d1l.py index f3ad0f70..5d068180 100644 --- a/tests/savepoint/translate/translate_tracer2d1l.py +++ b/tests/savepoint/translate/translate_tracer2d1l.py @@ -1,11 +1,12 @@ import pytest import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import Namelist, StencilFactory, QuantityFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import ParallelTranslate from pyFV3.stencils import FiniteVolumeTransport, TracerAdvection from pyFV3.utils.functional_validation import get_subset_func +from pyFV3.tracers import Tracers class TranslateTracer2D1L(ParallelTranslate): @@ -34,6 +35,10 @@ def __init__( self._base.in_vars["parameters"] = ["nq"] self._base.out_vars = self._base.in_vars["data_vars"] self.stencil_factory = stencil_factory + self._quantity_factory = QuantityFactory.from_backend( + sizer=stencil_factory.grid_indexing._sizer, + backend=stencil_factory.backend, + ) self.namelist = namelist self._subset = get_subset_func( self.grid.grid_indexing, @@ -46,11 +51,23 @@ def collect_input_data(self, serializer, savepoint): return input_data def compute_parallel(self, inputs, communicator): - self._base.make_storage_data_input_vars(inputs) - all_tracers = inputs["tracers"] - inputs["tracers"] = self.get_advected_tracer_dict( - inputs["tracers"], int(inputs.pop("nq")) + tracers = Tracers.make_from_fortran( + quantity_factory=self._quantity_factory, + tracer_mapping=[ + "vapor", + "liquid", + "rain", + "ice", + "snow", + "graupel", + "qo3mr", + "qsgs_tke", + ], + tracer_data=inputs["tracers"], ) + self._base.make_storage_data_input_vars(inputs, dict_4d=False) + inputs.pop("nq") # Fortran NQ is intrinsic to Tracers (e.g Tracers.count) + all_tracers = inputs.pop("tracers") transport = FiniteVolumeTransport( stencil_factory=self.stencil_factory, quantity_factory=self.grid.quantity_factory, @@ -66,20 +83,24 @@ def compute_parallel(self, inputs, communicator): transport, self.grid.grid_data, communicator, - inputs["tracers"], + tracers, ) inputs["x_mass_flux"] = inputs.pop("mfxd") inputs["y_mass_flux"] = inputs.pop("mfyd") inputs["x_courant"] = inputs.pop("cxd") inputs["y_courant"] = inputs.pop("cyd") - self.tracer_advection(**inputs) + self.tracer_advection(tracers=tracers, **inputs) inputs["mfxd"] = inputs.pop("x_mass_flux") inputs["mfyd"] = inputs.pop("y_mass_flux") inputs["cxd"] = inputs.pop("x_courant") inputs["cyd"] = inputs.pop("y_courant") - inputs[ - "tracers" - ] = all_tracers # some aren't advected, still need to be validated + # Put back un-advected tracers + # Tracers have -1 on all cartesian because of NDSL padding + # Dev note: qcld is not advected in Pace dataset for some reason + tracers_as_4d = tracers.as_fortran_4D() + for idx in range(0, tracers_as_4d.shape[3]): + all_tracers[:-1, :-1, :-1, idx] = tracers_as_4d[:, :, :, idx] + inputs["tracers"] = all_tracers # need to convert tracers dict to [x, y, z, n_tracer] array before subsetting outputs = self._base.slice_output(inputs) outputs["tracers"] = self.subset_output("tracers", outputs["tracers"]) From 8071d297e7ce28d213437efb781d7e75a92d2f25 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Dec 2024 16:29:22 -0500 Subject: [PATCH 046/252] Move Fillz (and test) to Tracers --- pyFV3/stencils/fillz.py | 18 ++------ tests/savepoint/translate/translate_fillz.py | 48 ++++++++++++++------ 2 files changed, 38 insertions(+), 28 deletions(-) diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index 5cd5c239..2b7cdb76 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -1,12 +1,11 @@ import typing -from typing import Dict from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval -import ndsl.dsl.gt4py_utils as utils -from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntFieldIJ +from pyFV3.tracers import Tracers @typing.no_type_check @@ -117,15 +116,12 @@ def __init__( self, stencil_factory: StencilFactory, quantity_factory: QuantityFactory, - nq: int, - tracers: Dict[str, Quantity], ): orchestrate( obj=self, config=stencil_factory.config.dace_config, dace_compiletime_args=["tracers"], ) - self._nq = int(nq) self._fix_tracer_stencil = stencil_factory.from_dims_halo( fix_tracer, compute_dims=[X_DIM, Y_DIM, Z_DIM], @@ -145,23 +141,19 @@ def __init__( dtype=Float, ) - self._filtered_tracer_dict = { - name: tracers[name] for name in utils.tracer_variables[0 : self._nq] - } - def __call__( self, dp2: FloatField, - tracers: Dict[str, Quantity], + tracers: Tracers, ): """ Args: dp2 (in): pressure thickness of atmospheric layer tracers (inout): tracers to fix negative masses in """ - for tracer_name in self._filtered_tracer_dict.keys(): + for tracer in tracers.values(): self._fix_tracer_stencil( - tracers[tracer_name], + tracer, dp2, self._zfix, self._sum0, diff --git a/tests/savepoint/translate/translate_fillz.py b/tests/savepoint/translate/translate_fillz.py index c08b5323..32f559fa 100644 --- a/tests/savepoint/translate/translate_fillz.py +++ b/tests/savepoint/translate/translate_fillz.py @@ -1,11 +1,12 @@ import numpy as np -import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import Namelist, StencilFactory, QuantityFactory from ndsl.stencils.testing import pad_field_in_j from ndsl.utils import safe_assign_array -from pyFV3.stencils import fillz +from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.testing import TranslateDycoreFortranData2Py +from pyFV3.tracers import Tracers +from typing import List class TranslateFillz(TranslateDycoreFortranData2Py): @@ -33,18 +34,20 @@ def __init__( self.max_error = 1e-13 self.ignore_near_zero_errors = {"q2tracers": True} self.stencil_factory = stencil_factory + self._quantity_factory = QuantityFactory.from_backend( + sizer=stencil_factory.grid_indexing._sizer, + backend=stencil_factory.backend, + ) - def make_storage_data_input_vars(self, inputs, storage_vars=None): - if storage_vars is None: - storage_vars = self.storage_vars() + def make_storage_data_input_vars(self, inputs, tracer_mapping: List[str]): + storage_vars = self.storage_vars() info = storage_vars["dp2"] inputs["dp2"] = self.make_storage_data( np.squeeze(inputs["dp2"]), istart=info["istart"], axis=info["axis"] ) - inputs["tracers"] = {} info = storage_vars["q2tracers"] for i in range(int(inputs["nq"])): - inputs["tracers"][utils.tracer_variables[i]] = self.make_storage_data( + inputs["tracers"][tracer_mapping[i]] = self.make_storage_data( np.squeeze(inputs["q2tracers"][:, :, i]), istart=info["istart"], axis=info["axis"], @@ -52,7 +55,23 @@ def make_storage_data_input_vars(self, inputs, storage_vars=None): del inputs["q2tracers"] def compute(self, inputs): - self.make_storage_data_input_vars(inputs) + tracer_mapping = [ + "vapor", + "liquid", + "rain", + "ice", + "snow", + "graupel", + "qo3mr", + "qsgs_tke", + ] + tracers = Tracers.make( + quantity_factory=self._quantity_factory, + tracer_mapping=tracer_mapping, + ) + inputs["tracers"] = tracers + + self.make_storage_data_input_vars(inputs, tracer_mapping=tracer_mapping) for name, value in tuple(inputs.items()): if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: inputs[name] = self.make_storage_data( @@ -67,18 +86,17 @@ def compute(self, inputs): value, self.grid.njd, backend=self.stencil_factory.backend ) ) - run_fillz = fillz.FillNegativeTracerValues( + inputs.pop("nq") + fillz = FillNegativeTracerValues( self.stencil_factory, self.grid.quantity_factory, - inputs.pop("nq"), - inputs["tracers"], ) - run_fillz(**inputs) + fillz(**inputs) ds = self.grid.default_domain_dict() ds.update(self.out_vars["q2tracers"]) - tracers = np.zeros((self.grid.nic, self.grid.npz, len(inputs["tracers"]))) + tracers = np.zeros((self.grid.nic, self.grid.npz, inputs["tracers"].count)) for varname, data in inputs["tracers"].items(): - index = utils.tracer_variables.index(varname) + index = tracer_mapping.index(varname) data[self.grid.slice_dict(ds)] safe_assign_array( tracers[:, :, index], np.squeeze(data[self.grid.slice_dict(ds)]) From 1a859a16d8afd44eba2f665b008e598efe2af90a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 13 Dec 2024 12:10:16 -0500 Subject: [PATCH 047/252] Move Remapping & MapNTracers to the new Tracers Rewire `qcloud` in the larger Tracers Fix bad `kord` on `MapNTracer` --- pyFV3/stencils/mapn_tracer.py | 29 ++++----- pyFV3/stencils/remapping.py | 60 +++++++++---------- .../translate/translate_remapping.py | 28 +++++++-- 3 files changed, 65 insertions(+), 52 deletions(-) diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index 0696d145..d5ddfadb 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -6,6 +6,7 @@ from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.stencils.map_single import MapSingle +from pyFV3.tracers import Tracers class MapNTracer: @@ -18,43 +19,39 @@ def __init__( stencil_factory: StencilFactory, quantity_factory: QuantityFactory, kord: int, - nq: int, fill: bool, - tracers: Dict[str, Quantity], + tracers: Tracers, ): orchestrate( obj=self, config=stencil_factory.config.dace_config, dace_compiletime_args=["tracers"], ) - self._nq = int(nq) self._qs = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="unknown", dtype=Float, ) - kord_tracer = [kord] * self._nq - kord_tracer[5] = 9 # qcld - - self._list_of_remap_objects = [ - MapSingle( + self._map_single = {} + for name in tracers.names(): + if name == "cloud": + this_kord = 9 + else: + this_kord = kord + self._map_single[name] = MapSingle( stencil_factory, quantity_factory, - kord_tracer[i], + this_kord, 0, dims=[X_DIM, Y_DIM, Z_DIM], ) - for i in range(len(kord_tracer)) - ] if fill: self._fill_negative_tracers = True self._fillz = FillNegativeTracerValues( stencil_factory, quantity_factory, - self._nq, - tracers, ) else: self._fill_negative_tracers = False @@ -64,7 +61,7 @@ def __call__( pe1: FloatField, pe2: FloatField, dp2: FloatField, - tracers: Dict[str, Quantity], + tracers: Tracers, ): """ Remaps the tracer species onto the Eulerian grid @@ -77,8 +74,8 @@ def __call__( dp2 (in): Difference in pressure between Eulerian levels tracers (inout): tracers to be remapped """ - for i, q in enumerate(utils.tracer_variables[0 : self._nq]): - self._list_of_remap_objects[i](tracers[q], pe1, pe2, self._qs) + for name, tracer in tracers.items(): + self._map_single[name](tracer, pe1, pe2, self._qs) if self._fill_negative_tracers is True: self._fillz(dp2, tracers) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index a14c38a1..591bb231 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -31,6 +31,7 @@ from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.moist_cv import moist_pt_func, moist_pt_last_step from pyFV3.stencils.saturation_adjustment import SatAdjust3d +from pyFV3.tracers import Tracers # TODO: Should this be set here or in global_constants? @@ -292,9 +293,8 @@ def __init__( quantity_factory: QuantityFactory, config: RemappingConfig, area_64, - nq, pfull, - tracers: Dict[str, Quantity], + tracers: Tracers, checkpointer: Optional[Checkpointer] = None, ): orchestrate( @@ -314,7 +314,6 @@ def __init__( raise NotImplementedError("Hydrostatic is not implemented") self._t_min = 184.0 - self._nq = nq # do_omega = hydrostatic and last_step # TODO pull into inputs self._domain_jextra = ( grid_indexing.domain[0], @@ -410,7 +409,6 @@ def __init__( stencil_factory, quantity_factory, abs(config.kord_tr), - nq, fill=config.fill, tracers=tracers, ) @@ -518,7 +516,7 @@ def __init__( def __call__( self, - tracers: Dict[str, Quantity], + tracers: Tracers, pt: FloatField, delp: FloatField, delz: FloatField, @@ -528,7 +526,6 @@ def __call__( w: FloatField, cappa: FloatField, q_con: FloatField, - q_cld: FloatField, pkz: FloatField, pk: FloatField, pe: FloatField, @@ -562,7 +559,6 @@ def __call__( va (inout): A-grid y-velocity cappa (inout): Power to raise pressure to q_con (out): Total condensate mixing ratio - q_cld (out): Cloud fraction pkz (in): Layer mean pressure raised to the power of Kappa pk (out): Interface pressure raised to power of kappa, final acoustic value pe (in): Pressure at layer edges @@ -593,12 +589,12 @@ def __call__( # pe2 is final Eulerian edge pressures self._moist_cv_pt_pressure( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], + tracers["vapor"], + tracers["liquid"], + tracers["rain"], + tracers["snow"], + tracers["ice"], + tracers["graupel"], q_con, pt, cappa, @@ -633,12 +629,12 @@ def __call__( # it clear the outputs are not needed until then? # or, are its outputs actually used? can we delete this stencil call? self._moist_cv_pkz( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], + tracers["vapor"], + tracers["liquid"], + tracers["rain"], + tracers["snow"], + tracers["ice"], + tracers["graupel"], q_con, self._gz, self._cvm, @@ -683,13 +679,13 @@ def __call__( fast_mp_consv = consv_te > CONSV_MIN self._saturation_adjustment( dp1, - tracers["qvapor"], - tracers["qliquid"], - tracers["qice"], - tracers["qrain"], - tracers["qsnow"], - tracers["qgraupel"], - q_cld, + tracers["vapor"], + tracers["liquid"], + tracers["ice"], + tracers["rain"], + tracers["snow"], + tracers["graupel"], + tracers["cloud"], hs, peln, delp, @@ -711,12 +707,12 @@ def __call__( # to the physics, but if we're staying in dynamics we need # to keep it as the virtual potential temperature self._moist_cv_last_step_stencil( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], + tracers["vapor"], + tracers["liquid"], + tracers["rain"], + tracers["snow"], + tracers["ice"], + tracers["graupel"], self._gz, pt, pkz, diff --git a/tests/savepoint/translate/translate_remapping.py b/tests/savepoint/translate/translate_remapping.py index 43ddb27d..316ab76f 100644 --- a/tests/savepoint/translate/translate_remapping.py +++ b/tests/savepoint/translate/translate_remapping.py @@ -1,9 +1,10 @@ import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import Namelist, StencilFactory, QuantityFactory from ndsl.constants import Z_DIM from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import LagrangianToEulerian from pyFV3.testing import TranslateDycoreFortranData2Py +from pyFV3.tracers import Tracers class TranslateRemapping(TranslateDycoreFortranData2Py): @@ -97,6 +98,10 @@ def __init__( self.ignore_near_zero_errors = {"q_con": True, "tracers": True} self.stencil_factory = stencil_factory self.namelist = DynamicalCoreConfig.from_namelist(namelist) + self._quantity_factory = QuantityFactory.from_backend( + sizer=stencil_factory.grid_indexing._sizer, + backend=stencil_factory.backend, + ) def compute_from_storage(self, inputs): wsd_2d = utils.make_storage_from_shape( @@ -104,19 +109,34 @@ def compute_from_storage(self, inputs): ) wsd_2d[:, :] = inputs["wsd"][:, :, 0] inputs["wsd"] = wsd_2d - inputs["q_cld"] = inputs["tracers"]["qcld"] + tracers = Tracers.make_from_fortran( + quantity_factory=self._quantity_factory, + tracer_mapping=[ + "vapor", + "liquid", + "rain", + "ice", + "snow", + "graupel", + "qo3mr", + "qsgs_tke", + "cloud", + ], + tracer_data=inputs["tracers"], + ) inputs["last_step"] = bool(inputs["last_step"]) pfull = self.grid.quantity_factory.zeros([Z_DIM], units="Pa") pfull.data[:] = pfull.np.asarray(inputs.pop("pfull")) + inputs.pop("nq") + inputs["tracers"] = tracers l_to_e_obj = LagrangianToEulerian( self.stencil_factory, quantity_factory=self.grid.quantity_factory, config=DynamicalCoreConfig.from_namelist(self.namelist).remapping, area_64=self.grid.area_64, - nq=inputs.pop("nq"), pfull=pfull, tracers=inputs["tracers"], ) l_to_e_obj(**inputs) - inputs.pop("q_cld") + inputs["tracers"] = tracers.as_fortran_4D() return inputs From 68afa44e0411d18a56934efea30182caaa3f7e79 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 13 Dec 2024 12:16:39 -0500 Subject: [PATCH 048/252] Insure scalars are properly typed State: forward precision for mixed fields --- pyFV3/dycore_state.py | 20 ++++++++++++++---- pyFV3/stencils/dyn_core.py | 42 +++++++++++++++++++++++--------------- 2 files changed, 41 insertions(+), 21 deletions(-) diff --git a/pyFV3/dycore_state.py b/pyFV3/dycore_state.py index 9058e4d5..9fdeaf12 100644 --- a/pyFV3/dycore_state.py +++ b/pyFV3/dycore_state.py @@ -1,5 +1,5 @@ from dataclasses import asdict, dataclass, field, fields -from typing import Any, Dict, Mapping, Union +from typing import Any, Dict, Mapping, Optional, Union import xarray as xr @@ -310,17 +310,27 @@ def __post_init__(self): ) @classmethod - def init_zeros(cls, quantity_factory: QuantityFactory): + def init_zeros( + cls, + quantity_factory: QuantityFactory, + dtype_dict: Optional[Dict[str, type]] = None, + allow_mismatch_float_precision=False, + ): initial_storages = {} for _field in fields(cls): if "dims" in _field.metadata.keys(): initial_storages[_field.name] = quantity_factory.zeros( _field.metadata["dims"], _field.metadata["units"], - dtype=Float, + dtype=dtype_dict[_field.name] + if dtype_dict and _field.name in dtype_dict.keys() + else Float, + allow_mismatch_float_precision=allow_mismatch_float_precision, ).data return cls.init_from_storages( - storages=initial_storages, sizer=quantity_factory.sizer + storages=initial_storages, + sizer=quantity_factory.sizer, + allow_mismatch_float_precision=allow_mismatch_float_precision, ) @classmethod @@ -355,6 +365,7 @@ def init_from_storages( sizer: GridSizer, bdt: float = 0.0, mdt: float = 0.0, + allow_mismatch_float_precision=False, ): inputs = {} for _field in fields(cls): @@ -366,6 +377,7 @@ def init_from_storages( _field.metadata["units"], origin=sizer.get_origin(dims), extent=sizer.get_extent(dims), + allow_mismatch_float_precision=allow_mismatch_float_precision, ) inputs[_field.name] = quantity return cls(**inputs, bdt=bdt, mdt=mdt) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 530d7f26..29d3591f 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -40,7 +40,7 @@ Z_INTERFACE_DIM, ) from ndsl.dsl.dace.orchestration import dace_inhibitor -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.typing import Float, FloatField, FloatField64, FloatFieldIJ from ndsl.grid import DampingCoefficients, GridData from ndsl.typing import Checkpointer, Communicator from pyFV3._config import AcousticDynamicsConfig @@ -59,10 +59,10 @@ def zero_data( - mfxd: FloatField, - mfyd: FloatField, - cxd: FloatField, - cyd: FloatField, + mfxd: FloatField64, + mfyd: FloatField64, + cxd: FloatField64, + cyd: FloatField64, heat_source: FloatField, diss_estd: FloatField, first_timestep: bool, @@ -442,7 +442,15 @@ def __init__( if config.d_ext != 0: raise RuntimeError("Acoustics (dyn_core): d_ext != 0 is not implemented") if config.beta != 0: - raise RuntimeError("Acoustics (dyn_core): beta != 0 is not implemented") + raise RuntimeError( + "Acoustics (dyn_core): beta != 0 is not implemented" + " (split_p_grad, etc.)" + ) + if config.beta < -0.1: + raise RuntimeError( + "Acoustics (dyn_core): beta < 0.1 is not implemented" + " (one_grad_p, etc.)" + ) if config.use_logp: raise RuntimeError("Acoustics (dyn_core): use_logp=True is not implemented") self._da_min = damping_coefficients.da_min @@ -460,6 +468,7 @@ def __init__( quantity_factory=quantity_factory, grid_data=grid_data, grid_type=config.grid_type, + use_logp=config.use_logp, ) ) self._akap = Float(constants.KAPPA) @@ -643,8 +652,8 @@ def __init__( # See divergence_damping.py, _get_da_min for explanation of this function @dace_inhibitor - def _get_da_min(self) -> float: - return self._da_min + def _get_da_min(self) -> Float: + return Float(self._da_min) def _checkpoint_csw(self, state: DycoreState, tag: str): if self.call_checkpointer: @@ -712,12 +721,12 @@ def _checkpoint_dsw_out(self, state: DycoreState): # See https://github.com/GEOS-ESM/pace/issues/9 @dace_inhibitor def dt_acoustic_substep(self, timestep: Float) -> Float: - return timestep / self.config.n_split + return Float(timestep / self.config.n_split) # TODO: Same as above @dace_inhibitor def dt2(self, dt_acoustic_substep: Float) -> Float: - return 0.5 * dt_acoustic_substep + return Float(0.5) * dt_acoustic_substep def __call__( self, @@ -730,8 +739,8 @@ def __call__( # akap, ptop, n_map, comm): end_step = n_map == self.config.k_split # dt = state.mdt / self.config.n_split - dt_acoustic_substep: Float = self.dt_acoustic_substep(timestep) - dt2: Float = self.dt2(dt_acoustic_substep) + dt_acoustic_substep = self.dt_acoustic_substep(timestep) + dt2 = self.dt2(dt_acoustic_substep) n_split = self.config.n_split # NOTE: In Fortran model the halo update starts happens in fv_dynamics, not here self._halo_updaters.q_con__cappa.start() @@ -995,19 +1004,18 @@ def __call__( if self.config.grid_type < 4: self._halo_updaters.interface_uc__vc.interface() - # we are here - if self._do_del2cubed: self._halo_updaters.heat_source.update() # TODO: move dependence on da_min into init of hyperdiffusion class - da_min: Float = self._get_da_min() - cd = constants.CNST_0P20 * da_min + cd = constants.CNST_0P20 * self._get_da_min() # we want to diffuse the heat source from damping before we apply it, # so that we don't reinforce the same grid-scale patterns we're trying # to damp self._hyperdiffusion(self._heat_source, cd) if not self.config.hydrostatic: - delt_time_factor = abs(dt_acoustic_substep * self.config.delt_max) + delt_time_factor = abs( + dt_acoustic_substep * Float(self.config.delt_max) + ) # TODO: it looks like state.pkz is being used as a temporary here, # and overwritten at the start of remapping. See if we can make it # an internal temporary of this stencil. From 424e538e4226657c9b75f6406636ddbdf8ab653a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 13 Dec 2024 12:17:24 -0500 Subject: [PATCH 049/252] Translate: allocate states with proper types --- pyFV3/testing/translate_dyncore.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/pyFV3/testing/translate_dyncore.py b/pyFV3/testing/translate_dyncore.py index 8524d090..3c398f4a 100644 --- a/pyFV3/testing/translate_dyncore.py +++ b/pyFV3/testing/translate_dyncore.py @@ -140,7 +140,15 @@ def compute_parallel(self, inputs, communicator): grid_data.bk = inputs["bk"] grid_data.ptop = inputs["ptop"] self._base.make_storage_data_input_vars(inputs) - state = DycoreState.init_zeros(quantity_factory=self.grid.quantity_factory) + inputs_dtypes = {} + for k, v in inputs.items(): + if hasattr(v, "dtype"): + inputs_dtypes[k] = v.dtype + state = DycoreState.init_zeros( + quantity_factory=self.grid.quantity_factory, + dtype_dict=inputs_dtypes, + allow_mismatch_float_precision=True, + ) wsd: Quantity = self.grid.quantity_factory.zeros( dims=[X_DIM, Y_DIM], units="unknown", From 4ada7428f08a049366f4d473fed89d7f576d5de2 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 13 Dec 2024 12:18:48 -0500 Subject: [PATCH 050/252] Update README with new branch --- README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 6b516c62..04a8062a 100644 --- a/README.md +++ b/README.md @@ -32,5 +32,6 @@ Branches: - 🔶 `fix/F32/UpdateDzD`@Florian: Fix for fluxes gradient & python computation - ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation and column calculation (partial pass) - ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 -- ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@FlorianL: Fix for 32-bit A2B_Ord4 -- ⚙️ `fix/GEOSv11_4_2/C_SW`: Fix for C_SW for 32-bit +- ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit A2B_Ord4 +- ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit +- ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit From f69a84155fddaaff994ad481f71bb6f3c9307e23 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 16 Dec 2024 16:36:52 -0500 Subject: [PATCH 051/252] Tracer: Freeze units and dims Lint --- pyFV3/tracers.py | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py index 301b1f6b..43802665 100644 --- a/pyFV3/tracers.py +++ b/pyFV3/tracers.py @@ -1,9 +1,13 @@ from __future__ import annotations -from ndsl import Quantity, QuantityFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM + from typing import Dict, List + import numpy as np +from ndsl import Quantity, QuantityFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM + + # FOR REFERENCE - previous descriptive of the tracers, lining up with Pace work # tracer_variables = [ # "qvapor", @@ -19,18 +23,20 @@ class Tracers: + unit = "g/kg" + dims = [X_DIM, Y_DIM, Z_DIM] + def __init__(self, factory: QuantityFactory) -> None: self._quantities: Dict[str, Quantity] = {} self._quantity_factory = factory - self._dims = [X_DIM, Y_DIM, Z_DIM] def copy_tracer_data( self, name: str, data: np.ndarray, - units="unknown", + unit="unknown", ): - qty = self._quantity_factory.empty(dims=self._dims, units=units) + qty = self._quantity_factory.empty(dims=self.dims, units=unit) if data.shape > qty.data.shape: raise ValueError( f"[pyFV3] Tracer {name} size ({data.shape}" @@ -52,8 +58,8 @@ def names(self): def items(self): return self._quantities.items() - def as_fortran_4D(self) -> np.ndarray: - shape = self._quantity_factory.sizer.get_shape(self._dims) + def as_4D_array(self) -> np.ndarray: + shape = self._quantity_factory.sizer.get_shape(self.dims) var4d = np.empty( ( shape[0] - 1, @@ -62,13 +68,12 @@ def as_fortran_4D(self) -> np.ndarray: self.count, ) ) + # Skip the extra data point that is meant to align interface + # and non interface fields for idx, q in enumerate(self.values()): var4d[:, :, :, idx] = q.data[:-1, :-1, :-1] return var4d - def as_dict(self) -> Dict[str, Quantity]: - return self._quantities - def __getitem__(self, key): return self._quantities[key] @@ -92,12 +97,12 @@ def make( ): tracers = cls(quantity_factory) for name in tracer_mapping: - qty = quantity_factory.empty(dims=tracers._dims, units="kg/m^2") + qty = quantity_factory.empty(dims=cls.dims, units=cls.unit) tracers._quantities[name] = qty return tracers @classmethod - def make_from_fortran( + def make_from_4D_array( cls, quantity_factory: QuantityFactory, tracer_mapping: List[str], @@ -116,7 +121,7 @@ def make_from_fortran( tracers.copy_tracer_data( name=tracer_mapping[idx] or f"Tracer_{idx}", data=tracer_data[:, :, :, idx], - units="kg/m^2", + unit=cls.unit, ) return tracers From 13b39f47bb11123992d8b3f2b17601d77a65c627 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 16 Dec 2024 17:17:06 -0500 Subject: [PATCH 052/252] Exclude tracers for Remapping & Advection --- pyFV3/stencils/fillz.py | 10 ++++-- pyFV3/stencils/mapn_tracer.py | 15 +++++--- tests/savepoint/translate/translate_fillz.py | 10 +++--- .../translate/translate_remapping.py | 7 ++-- .../translate/translate_tracer2d1l.py | 34 +++++-------------- 5 files changed, 36 insertions(+), 40 deletions(-) diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index 2b7cdb76..6d86b20f 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -1,4 +1,4 @@ -import typing +from typing import List, no_type_check from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval @@ -8,7 +8,7 @@ from pyFV3.tracers import Tracers -@typing.no_type_check +@no_type_check def fix_tracer( q: FloatField, dp: FloatField, @@ -116,6 +116,7 @@ def __init__( self, stencil_factory: StencilFactory, quantity_factory: QuantityFactory, + exclude_tracers: List[str], ): orchestrate( obj=self, @@ -126,6 +127,7 @@ def __init__( fix_tracer, compute_dims=[X_DIM, Y_DIM, Z_DIM], ) + self._exclude_tracers = exclude_tracers # Setting initial value of upper_fix to zero is only needed for validation. # The values in the compute domain are set to zero in the stencil. @@ -151,7 +153,9 @@ def __call__( dp2 (in): pressure thickness of atmospheric layer tracers (inout): tracers to fix negative masses in """ - for tracer in tracers.values(): + for name, tracer in tracers.items(): + if name in self._exclude_tracers: + continue self._fix_tracer_stencil( tracer, dp2, diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index d5ddfadb..ea4f3f05 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -1,12 +1,12 @@ -from typing import Dict +from typing import List -import ndsl.dsl.gt4py_utils as utils -from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.stencils.map_single import MapSingle from pyFV3.tracers import Tracers +import dace class MapNTracer: @@ -21,12 +21,14 @@ def __init__( kord: int, fill: bool, tracers: Tracers, + exclude_tracers: List[str], ): orchestrate( obj=self, config=stencil_factory.config.dace_config, dace_compiletime_args=["tracers"], ) + self._exclude_tracers = exclude_tracers self._qs = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="unknown", @@ -52,6 +54,7 @@ def __init__( self._fillz = FillNegativeTracerValues( stencil_factory, quantity_factory, + exclude_tracers=self._exclude_tracers, ) else: self._fill_negative_tracers = False @@ -74,8 +77,10 @@ def __call__( dp2 (in): Difference in pressure between Eulerian levels tracers (inout): tracers to be remapped """ - for name, tracer in tracers.items(): - self._map_single[name](tracer, pe1, pe2, self._qs) + for name in tracers.names(): + if name in self._exclude_tracers: + continue + self._map_single[name](tracers[name], pe1, pe2, self._qs) if self._fill_negative_tracers is True: self._fillz(dp2, tracers) diff --git a/tests/savepoint/translate/translate_fillz.py b/tests/savepoint/translate/translate_fillz.py index 32f559fa..2e90ff8e 100644 --- a/tests/savepoint/translate/translate_fillz.py +++ b/tests/savepoint/translate/translate_fillz.py @@ -1,12 +1,13 @@ +from typing import List + import numpy as np -from ndsl import Namelist, StencilFactory, QuantityFactory +from ndsl import Namelist, QuantityFactory, StencilFactory from ndsl.stencils.testing import pad_field_in_j from ndsl.utils import safe_assign_array from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.testing import TranslateDycoreFortranData2Py from pyFV3.tracers import Tracers -from typing import List class TranslateFillz(TranslateDycoreFortranData2Py): @@ -62,8 +63,8 @@ def compute(self, inputs): "ice", "snow", "graupel", - "qo3mr", - "qsgs_tke", + "o3mr", + "sgs_tke", ] tracers = Tracers.make( quantity_factory=self._quantity_factory, @@ -90,6 +91,7 @@ def compute(self, inputs): fillz = FillNegativeTracerValues( self.stencil_factory, self.grid.quantity_factory, + exclude_tracers=[], ) fillz(**inputs) ds = self.grid.default_domain_dict() diff --git a/tests/savepoint/translate/translate_remapping.py b/tests/savepoint/translate/translate_remapping.py index 316ab76f..c644c6bb 100644 --- a/tests/savepoint/translate/translate_remapping.py +++ b/tests/savepoint/translate/translate_remapping.py @@ -1,5 +1,5 @@ import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory, QuantityFactory +from ndsl import Namelist, QuantityFactory, StencilFactory from ndsl.constants import Z_DIM from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import LagrangianToEulerian @@ -109,7 +109,7 @@ def compute_from_storage(self, inputs): ) wsd_2d[:, :] = inputs["wsd"][:, :, 0] inputs["wsd"] = wsd_2d - tracers = Tracers.make_from_fortran( + tracers = Tracers.make_from_4D_array( quantity_factory=self._quantity_factory, tracer_mapping=[ "vapor", @@ -136,7 +136,8 @@ def compute_from_storage(self, inputs): area_64=self.grid.area_64, pfull=pfull, tracers=inputs["tracers"], + exclude_tracers=["cloud"], ) l_to_e_obj(**inputs) - inputs["tracers"] = tracers.as_fortran_4D() + inputs["tracers"] = tracers.as_4D_array() return inputs diff --git a/tests/savepoint/translate/translate_tracer2d1l.py b/tests/savepoint/translate/translate_tracer2d1l.py index 5d068180..dcc9c104 100644 --- a/tests/savepoint/translate/translate_tracer2d1l.py +++ b/tests/savepoint/translate/translate_tracer2d1l.py @@ -1,12 +1,11 @@ import pytest -import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory, QuantityFactory +from ndsl import Namelist, QuantityFactory, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import ParallelTranslate from pyFV3.stencils import FiniteVolumeTransport, TracerAdvection -from pyFV3.utils.functional_validation import get_subset_func from pyFV3.tracers import Tracers +from pyFV3.utils.functional_validation import get_subset_func class TranslateTracer2D1L(ParallelTranslate): @@ -51,7 +50,7 @@ def collect_input_data(self, serializer, savepoint): return input_data def compute_parallel(self, inputs, communicator): - tracers = Tracers.make_from_fortran( + tracers = Tracers.make_from_4D_array( quantity_factory=self._quantity_factory, tracer_mapping=[ "vapor", @@ -60,14 +59,15 @@ def compute_parallel(self, inputs, communicator): "ice", "snow", "graupel", - "qo3mr", - "qsgs_tke", + "o3mr", + "sgs_tke", + "cloud", ], tracer_data=inputs["tracers"], ) self._base.make_storage_data_input_vars(inputs, dict_4d=False) + inputs.pop("tracers") inputs.pop("nq") # Fortran NQ is intrinsic to Tracers (e.g Tracers.count) - all_tracers = inputs.pop("tracers") transport = FiniteVolumeTransport( stencil_factory=self.stencil_factory, quantity_factory=self.grid.quantity_factory, @@ -84,6 +84,7 @@ def compute_parallel(self, inputs, communicator): self.grid.grid_data, communicator, tracers, + exclude_tracers=["cloud"], ) inputs["x_mass_flux"] = inputs.pop("mfxd") inputs["y_mass_flux"] = inputs.pop("mfyd") @@ -97,28 +98,11 @@ def compute_parallel(self, inputs, communicator): # Put back un-advected tracers # Tracers have -1 on all cartesian because of NDSL padding # Dev note: qcld is not advected in Pace dataset for some reason - tracers_as_4d = tracers.as_fortran_4D() - for idx in range(0, tracers_as_4d.shape[3]): - all_tracers[:-1, :-1, :-1, idx] = tracers_as_4d[:, :, :, idx] - inputs["tracers"] = all_tracers - # need to convert tracers dict to [x, y, z, n_tracer] array before subsetting + inputs["tracers"] = tracers.as_4D_array() outputs = self._base.slice_output(inputs) outputs["tracers"] = self.subset_output("tracers", outputs["tracers"]) return outputs - def get_advected_tracer_dict(self, all_tracers, nq): - all_tracers = {**all_tracers} # make a new dict so we don't modify the input - properties = self.inputs["tracers"] - for name in utils.tracer_variables: - self.grid.quantity_dict_update( - all_tracers, - name, - dims=properties["dims"], - units=properties["units"], - ) - tracer_names = utils.tracer_variables[:nq] - return {name: all_tracers[name + "_quantity"] for name in tracer_names} - def compute_sequential(self, a, b): pytest.skip( f"{self.__class__} only has a mpirun implementation, " From 1474a4f26c175bdb94fc28352fc0184d4ffa5144 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 16 Dec 2024 17:18:28 -0500 Subject: [PATCH 053/252] Missing file from previous commit --- pyFV3/stencils/tracer_2d_1l.py | 30 ++++++++++++++++++------------ 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index e90600e4..2521a64c 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -1,16 +1,10 @@ import math -from typing import Dict +from typing import List import gt4py.cartesian.gtscript as gtscript from gt4py.cartesian.gtscript import PARALLEL, computation, horizontal, interval, region -from ndsl import ( - Quantity, - QuantityFactory, - StencilFactory, - WrappedHaloUpdater, - orchestrate, -) +from ndsl import QuantityFactory, StencilFactory, WrappedHaloUpdater, orchestrate from ndsl.constants import ( N_HALO_DEFAULT, X_DIM, @@ -194,6 +188,7 @@ def __init__( grid_data, comm: Communicator, tracers: Tracers, + exclude_tracers: List[str], ): orchestrate( obj=self, @@ -204,6 +199,7 @@ def __init__( self.grid_indexing = grid_indexing # needed for selective validation self._tracer_count = tracers.count self.grid_data = grid_data + self._exclude_tracers = exclude_tracers self._x_area_flux = quantity_factory.zeros( [X_INTERFACE_DIM, Y_DIM, Z_DIM], @@ -283,10 +279,18 @@ def __init__( n_halo=N_HALO_DEFAULT, dtype=Float, ) + + # We can exclude tracers from advecting and therefore also + # halo exchanging + advected_tracers = {} + for name, tracer in tracers.items(): + if name in exclude_tracers: + continue + advected_tracers[name] = tracer self._tracers_halo_updater = WrappedHaloUpdater( - comm.get_scalar_halo_updater([tracer_halo_spec] * tracers.count), - tracers.as_dict(), - [t for t in tracers.names()], + comm.get_scalar_halo_updater([tracer_halo_spec] * len(advected_tracers)), + advected_tracers, + [t for t in advected_tracers.keys()], ) def __call__( @@ -393,7 +397,9 @@ def __call__( self.grid_data.rarea, dp2, ) - for q in tracers.values(): + for name, q in tracers.items(): + if name in self._exclude_tracers: + continue self.finite_volume_transport( q, x_courant, From 714f1a8e627bb2091eb123747ac0645895c20a9e Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 16 Dec 2024 17:18:52 -0500 Subject: [PATCH 054/252] Dynamics with Tracers including DycoreState --- examples/standalone/runfile/dynamics.py | 1 + pyFV3/dycore_state.py | 312 ++++++++++++------------ pyFV3/stencils/fv_dynamics.py | 72 ++++-- pyFV3/testing/translate_dyncore.py | 5 +- pyFV3/testing/translate_fvdynamics.py | 66 ++++- 5 files changed, 268 insertions(+), 188 deletions(-) diff --git a/examples/standalone/runfile/dynamics.py b/examples/standalone/runfile/dynamics.py index 88326001..103161d1 100755 --- a/examples/standalone/runfile/dynamics.py +++ b/examples/standalone/runfile/dynamics.py @@ -263,6 +263,7 @@ def setup_dycore( config=dycore_config, phis=state.phis, state=state, + exclude_tracers=[], timestep=timedelta(seconds=dycore_config.dt_atmos), ) return dycore, state, stencil_factory diff --git a/pyFV3/dycore_state.py b/pyFV3/dycore_state.py index 9058e4d5..742fd55c 100644 --- a/pyFV3/dycore_state.py +++ b/pyFV3/dycore_state.py @@ -1,10 +1,10 @@ from dataclasses import asdict, dataclass, field, fields -from typing import Any, Dict, Mapping, Union +from typing import Any, Dict, List, Mapping, Union import xarray as xr import ndsl.dsl.gt4py_utils as gt_utils -from ndsl import GridSizer, Quantity, QuantityFactory +from ndsl import Quantity, QuantityFactory from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -16,6 +16,65 @@ from ndsl.dsl.typing import Float from ndsl.restart._legacy_restart import open_restart from ndsl.typing import Communicator +from pyFV3.tracers import Tracers + + +DEFAULT_TRACER_PROPERTIES = { + "specific_humidity": { + "pyFV3_key": "vapor", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "sphum", + "units": "g/kg", + }, + "cloud_liquid_water_mixing_ratio": { + "pyFV3_key": "liquid", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "liq_wat", + "units": "g/kg", + }, + "cloud_ice_mixing_ratio": { + "pyFV3_key": "ice", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "ice_wat", + "units": "g/kg", + }, + "rain_mixing_ratio": { + "pyFV3_key": "rain", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "rainwat", + "units": "g/kg", + }, + "snow_mixing_ratio": { + "pyFV3_key": "snow", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "snowwat", + "units": "g/kg", + }, + "graupel_mixing_ratio": { + "pyFV3_key": "graupel", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "graupel", + "units": "g/kg", + }, + "ozone_mixing_ratio": { + "pyFV3_key": "o3mr", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "o3mr", + "units": "g/kg", + }, + "turbulent_kinetic_energy": { + "pyFV3_key": "sgs_tke", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "sgs_tke", + "units": "g/kg", + }, + "cloud_fraction": { + "pyFV3_key": "cloud", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "cld_amt", + "units": "g/kg", + }, +} @dataclass() @@ -148,74 +207,10 @@ class DycoreState: "intent": "inout", } ) - qvapor: Quantity = field( - metadata={ - "name": "specific_humidity", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - } - ) - qliquid: Quantity = field( - metadata={ - "name": "cloud_water_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - "intent": "inout", - } - ) - qice: Quantity = field( + tracers: Tracers = field( metadata={ - "name": "cloud_ice_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - "intent": "inout", - } - ) - qrain: Quantity = field( - metadata={ - "name": "rain_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - "intent": "inout", - } - ) - qsnow: Quantity = field( - metadata={ - "name": "snow_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - "intent": "inout", - } - ) - qgraupel: Quantity = field( - metadata={ - "name": "graupel_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - "intent": "inout", - } - ) - qo3mr: Quantity = field( - metadata={ - "name": "ozone_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - "intent": "inout", - } - ) - qsgs_tke: Quantity = field( - metadata={ - "name": "turbulent_kinetic_energy", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "m**2/s**2", - "intent": "inout", - } - ) - qcld: Quantity = field( - metadata={ - "name": "cloud_fraction", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "", + "name": "tracers", + "units": "g/kg", "intent": "inout", } ) @@ -297,6 +292,8 @@ class DycoreState: def __post_init__(self): for _field in fields(self): + if _field.name == "tracers": + continue for check_name in ["units", "dims"]: if check_name in _field.metadata: required = _field.metadata[check_name] @@ -310,7 +307,7 @@ def __post_init__(self): ) @classmethod - def init_zeros(cls, quantity_factory: QuantityFactory): + def init_zeros(cls, quantity_factory: QuantityFactory, tracer_list: List[str]): initial_storages = {} for _field in fields(cls): if "dims" in _field.metadata.keys(): @@ -319,13 +316,25 @@ def init_zeros(cls, quantity_factory: QuantityFactory): _field.metadata["units"], dtype=Float, ).data + for name in tracer_list: + initial_storages[name] = quantity_factory.zeros( + Tracers.dims, + Tracers.unit, + dtype=Float, + ).data return cls.init_from_storages( - storages=initial_storages, sizer=quantity_factory.sizer + storages=initial_storages, + quantity_factory=quantity_factory, + tracer_list=tracer_list, ) @classmethod def init_from_numpy_arrays( - cls, dict_of_numpy_arrays, sizer: GridSizer, backend: str + cls, + dict_of_numpy_arrays, + quantity_factory: QuantityFactory, + backend: str, + tracer_list: List[str], ): field_names = [_field.name for _field in fields(cls)] for variable_name in dict_of_numpy_arrays.keys(): @@ -341,10 +350,22 @@ def init_from_numpy_arrays( dict_of_numpy_arrays[_field.name], dims, _field.metadata["units"], - origin=sizer.get_origin(dims), - extent=sizer.get_extent(dims), + origin=quantity_factory.sizer.get_origin(dims), + extent=quantity_factory.sizer.get_extent(dims), gt4py_backend=backend, ) + elif issubclass(_field.type, Tracers): + if len(dict_of_numpy_arrays[_field.name]) != len(tracer_list): + raise ValueError( + "[pyFV3] DycoreState init:" + f" tracer list size ({len(tracer_list)})" + " doesn't match the inputs size" + f" ({len(dict_of_numpy_arrays[_field.name])})" + ) + dict_state[_field.name] = Tracers.make( + quantity_factory=quantity_factory, + tracer_mapping=tracer_list, + ) state = cls(**dict_state) # type: ignore return state @@ -352,7 +373,8 @@ def init_from_numpy_arrays( def init_from_storages( cls, storages: Mapping[str, Any], - sizer: GridSizer, + quantity_factory: QuantityFactory, + tracer_list: List[str], bdt: float = 0.0, mdt: float = 0.0, ): @@ -364,10 +386,19 @@ def init_from_storages( storages[_field.name], dims, _field.metadata["units"], - origin=sizer.get_origin(dims), - extent=sizer.get_extent(dims), + origin=quantity_factory.sizer.get_origin(dims), + extent=quantity_factory.sizer.get_extent(dims), ) inputs[_field.name] = quantity + elif "tracers" == _field.name: + tracers = Tracers.make( + quantity_factory=quantity_factory, + tracer_mapping=tracer_list, + ) + for name in tracer_list: + tracers[name].data[:] = storages[name][:] + inputs[_field.name] = tracers + return cls(**inputs, bdt=bdt, mdt=mdt) @classmethod @@ -381,10 +412,14 @@ def from_fortran_restart( state_dict: Mapping[str, Quantity] = open_restart( dirname=path, communicator=communicator, - tracer_properties=TRACER_PROPERTIES, + tracer_properties=DEFAULT_TRACER_PROPERTIES, + ) + new = cls.init_zeros( + quantity_factory=quantity_factory, + tracer_list=[ + str(prop["pyFV3_key"]) for prop in DEFAULT_TRACER_PROPERTIES.values() + ], ) - - new = cls.init_zeros(quantity_factory=quantity_factory) new.pt.view[:] = new.pt.np.asarray( state_dict["air_temperature"].transpose(new.pt.dims).view[:] ) @@ -405,31 +440,35 @@ def from_fortran_restart( new.v.view[:] = new.v.np.asarray( state_dict["y_wind"].transpose(new.v.dims).view[:] ) - new.qvapor.view[:] = new.qvapor.np.asarray( - state_dict["specific_humidity"].transpose(new.qvapor.dims).view[:] + new.tracers["vapor"].view[:] = new.tracers["vapor"].np.asarray( + state_dict["specific_humidity"].transpose(new.tracers["vapor"].dims).view[:] ) - new.qliquid.view[:] = new.qliquid.np.asarray( + new.tracers["liquid"].view[:] = new.tracers["liquid"].np.asarray( state_dict["cloud_liquid_water_mixing_ratio"] - .transpose(new.qliquid.dims) + .transpose(new.tracers["liquid"].dims) .view[:] ) - new.qice.view[:] = new.qice.np.asarray( - state_dict["cloud_ice_mixing_ratio"].transpose(new.qice.dims).view[:] + new.tracers["ice"].view[:] = new.tracers["ice"].np.asarray( + state_dict["cloud_ice_mixing_ratio"] + .transpose(new.tracers["ice"].dims) + .view[:] ) - new.qrain.view[:] = new.qrain.np.asarray( - state_dict["rain_mixing_ratio"].transpose(new.qrain.dims).view[:] + new.tracers["rain"].view[:] = new.tracers["rain"].np.asarray( + state_dict["rain_mixing_ratio"].transpose(new.tracers["rain"].dims).view[:] ) - new.qsnow.view[:] = new.qsnow.np.asarray( - state_dict["snow_mixing_ratio"].transpose(new.qsnow.dims).view[:] + new.tracers["snow"].view[:] = new.tracers["snow"].np.asarray( + state_dict["snow_mixing_ratio"].transpose(new.tracers["snow"].dims).view[:] ) - new.qgraupel.view[:] = new.qgraupel.np.asarray( - state_dict["graupel_mixing_ratio"].transpose(new.qgraupel.dims).view[:] + new.tracers["graupel"].view[:] = new.tracers["graupel"].np.asarray( + state_dict["graupel_mixing_ratio"] + .transpose(new.tracers["graupel"].dims) + .view[:] ) - new.qo3mr.view[:] = new.qo3mr.np.asarray( - state_dict["ozone_mixing_ratio"].transpose(new.qo3mr.dims).view[:] + new.tracers["o3mr"].view[:] = new.tracers["o3mr"].np.asarray( + state_dict["ozone_mixing_ratio"].transpose(new.tracers["o3mr"].dims).view[:] ) - new.qcld.view[:] = new.qcld.np.asarray( - state_dict["cloud_fraction"].transpose(new.qcld.dims).view[:] + new.tracers["cloud"].view[:] = new.tracers["cld"].np.asarray( + state_dict["cloud_fraction"].transpose(new.tracers["cld"].dims).view[:] ) new.delz.view[:] = new.delz.np.asarray( state_dict["vertical_thickness_of_atmospheric_layer"] @@ -439,22 +478,34 @@ def from_fortran_restart( return new + def _xr_dataarray_from_quantity(self, name: str, metadata: Dict[str, Any], data): + dims = [f"{dim_name}_{name}" for dim_name in metadata["dims"]] + return xr.DataArray( + gt_utils.asarray(data), + dims=dims, + attrs={ + "long_name": metadata["name"], + "units": metadata.get("units", "unknown"), + }, + ) + @property def xr_dataset(self): data_vars = {} for name, field_info in self.__dataclass_fields__.items(): if issubclass(field_info.type, Quantity): - dims = [ - f"{dim_name}_{name}" for dim_name in field_info.metadata["dims"] - ] - data_vars[name] = xr.DataArray( - gt_utils.asarray(getattr(self, name).data), - dims=dims, - attrs={ - "long_name": field_info.metadata["name"], - "units": field_info.metadata.get("units", "unknown"), - }, + data_vars[name] = self._xr_dataarray_from_quantity( + name=name, + metadata=field_info.metadata, + data=getattr(self, name).data, ) + if isinstance(field_info.type, Tracers): + for tracer in getattr(self, name).values(): + data_vars[name] = self._xr_dataarray_from_quantity( + name=name, + metadata=field_info.metadata, + data=tracer, + ) return xr.Dataset(data_vars=data_vars) def __getitem__(self, item): @@ -465,52 +516,3 @@ def as_dict(self, quantity_only=True) -> Dict[str, Union[Quantity, int]]: return {k: v for k, v in asdict(self).items() if isinstance(v, Quantity)} else: return {k: v for k, v in asdict(self).items()} - - -TRACER_PROPERTIES = { - "specific_humidity": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "sphum", - "units": "g/kg", - }, - "cloud_liquid_water_mixing_ratio": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "liq_wat", - "units": "g/kg", - }, - "cloud_ice_mixing_ratio": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "ice_wat", - "units": "g/kg", - }, - "rain_mixing_ratio": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "rainwat", - "units": "g/kg", - }, - "snow_mixing_ratio": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "snowwat", - "units": "g/kg", - }, - "graupel_mixing_ratio": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "graupel", - "units": "g/kg", - }, - "ozone_mixing_ratio": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "o3mr", - "units": "g/kg", - }, - "turbulent_kinetic_energy": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "sgs_tke", - "units": "g/kg", - }, - "cloud_fraction": { - "dims": [Z_DIM, Y_DIM, X_DIM], - "restart_name": "cld_amt", - "units": "g/kg", - }, -} diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 50ffba27..9a71d3c7 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -1,10 +1,9 @@ from datetime import timedelta -from typing import Mapping, Optional +from typing import List, Mapping, Optional from dace.frontend.python.interface import nounroll as dace_no_unroll from gt4py.cartesian.gtscript import PARALLEL, computation, interval -import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater from ndsl.checkpointer import NullCheckpointer @@ -98,6 +97,7 @@ def __init__( config: DynamicalCoreConfig, phis: Quantity, state: DycoreState, + exclude_tracers: List[str], timestep: timedelta, checkpointer: Optional[Checkpointer] = None, ): @@ -111,6 +111,8 @@ def __init__( the namelist in the Fortran model phis: surface geopotential height state: model state + exclude_tracer: List of named tracer to be excluded from the Advection, + and Remapping schemes timestep: model timestep checkpointer: if given, used to perform operations on model data at specific points in model execution, such as testing against @@ -198,6 +200,28 @@ def __init__( f" nwat=={config.nwat} is not implemented." " Only nwat=6 has been implemented." ) + + # Implemented dynamics options require those tracers to be present at minima + # this is a more granular list than carried by the `nwat` single integer + # but cover the same topic + required_tracers = [ + "vapor", + "liquid", + "rain", + "snow", + "ice", + "graupel", + "cloud", + ] + if not all(n in state.tracers.names() for n in required_tracers): + raise NotImplementedError( + "Dynamical core (fv_dynamics):" + " missing required tracers. Dynamics requires:\n" + f" {required_tracers}\n" + "but only the following where given:\n" + f" {state.tracers.names()}" + ) + self.comm_rank = comm.rank self.grid_data = grid_data self.grid_indexing = grid_indexing @@ -213,10 +237,6 @@ def __init__( hord=config.hord_tr, ) - self.tracers = {} - for name in utils.tracer_variables[0:NQ]: - self.tracers[name] = state.__dict__[name] - temporaries = fvdyn_temporaries(quantity_factory) self._te_2d = temporaries["te_2d"] self._te0_2d = temporaries["te0_2d"] @@ -231,7 +251,8 @@ def __init__( tracer_transport, self.grid_data, comm, - self.tracers, + state.tracers, + exclude_tracers=exclude_tracers, ) self._ak = grid_data.ak self._bk = grid_data.bk @@ -312,9 +333,9 @@ def __init__( quantity_factory=quantity_factory, config=config.remapping, area_64=grid_data.area_64, - nq=NQ, pfull=self._pfull, - tracers=self.tracers, + tracers=state.tracers, + exclude_tracers=exclude_tracers, checkpointer=checkpointer, ) @@ -352,7 +373,7 @@ def _checkpoint_fvdynamics(self, state: DycoreState, tag: str): va=state.va, uc=state.uc, vc=state.vc, - qvapor=state.qvapor, + qvapor=state.tracers["vapor"], ) def _checkpoint_remapping_in( @@ -462,12 +483,12 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool): log_on_rank_0("FV Setup") self._fv_setup_stencil( - state.qvapor, - state.qliquid, - state.qrain, - state.qsnow, - state.qice, - state.qgraupel, + state.tracers["vapor"], + state.tracers["liquid"], + state.tracers["rain"], + state.tracers["snow"], + state.tracers["ice"], + state.tracers["graupel"], state.q_con, self._cvm, state.pkz, @@ -536,7 +557,7 @@ def _compute(self, state: DycoreState, timer: Timer): with timer.clock("TracerAdvection"): self._checkpoint_tracer_advection_in(state) self.tracer_advection( - self.tracers, + state.tracers, self._dp_initial, state.mfxd, state.mfyd, @@ -573,7 +594,7 @@ def _compute(self, state: DycoreState, timer: Timer): # time # When NQ=8, we do need qcld passed explicitely self._lagrangian_to_eulerian_obj( - self.tracers, + state.tracers, state.pt, state.delp, state.delz, @@ -583,7 +604,6 @@ def _compute(self, state: DycoreState, timer: Timer): state.w, self._cappa, state.q_con, - state.qcld, state.pkz, state.pk, state.pe, @@ -625,13 +645,13 @@ def _compute(self, state: DycoreState, timer: Timer): if __debug__: log_on_rank_0("Neg Adj 3") self._adjust_tracer_mixing_ratio( - state.qvapor, - state.qliquid, - state.qrain, - state.qsnow, - state.qice, - state.qgraupel, - state.qcld, + state.tracers["vapor"], + state.tracers["liquid"], + state.tracers["rain"], + state.tracers["snow"], + state.tracers["ice"], + state.tracers["graupel"], + state.tracers["cloud"], state.pt, state.delp, ) diff --git a/pyFV3/testing/translate_dyncore.py b/pyFV3/testing/translate_dyncore.py index 8524d090..111d6844 100644 --- a/pyFV3/testing/translate_dyncore.py +++ b/pyFV3/testing/translate_dyncore.py @@ -140,7 +140,10 @@ def compute_parallel(self, inputs, communicator): grid_data.bk = inputs["bk"] grid_data.ptop = inputs["ptop"] self._base.make_storage_data_input_vars(inputs) - state = DycoreState.init_zeros(quantity_factory=self.grid.quantity_factory) + state = DycoreState.init_zeros( + quantity_factory=self.grid.quantity_factory, + tracer_list=[], # No tracers used in acoustics + ) wsd: Quantity = self.grid.quantity_factory.zeros( dims=[X_DIM, Y_DIM], units="unknown", diff --git a/pyFV3/testing/translate_fvdynamics.py b/pyFV3/testing/translate_fvdynamics.py index 5daea9d0..a4178640 100644 --- a/pyFV3/testing/translate_fvdynamics.py +++ b/pyFV3/testing/translate_fvdynamics.py @@ -5,7 +5,7 @@ import pytest import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, Quantity, StencilFactory +from ndsl import Namelist, Quantity, QuantityFactory, StencilFactory from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -33,6 +33,31 @@ def __init__( self.namelist = DynamicalCoreConfig.from_namelist(namelist) +TRACERS_IN_PYFV3 = [ + "vapor", + "liquid", + "ice", + "rain", + "snow", + "graupel", + "o3mr", + "sgs_tke", + "cloud", +] + +TRACERS_IN_FORTRAN = [ + "qvapor", + "qliquid", + "qice", + "qrain", + "qsnow", + "qgraupel", + "qo3mr", + "qsgs_tke", + "qcld", +] + + class TranslateFVDynamics(ParallelTranslateBaseSlicing): compute_grid_option = True inputs: Dict[str, Any] = { @@ -290,25 +315,50 @@ def __init__( self.max_error = 1e-5 self.ignore_near_zero_errors = {} - for qvar in utils.tracer_variables: - self.ignore_near_zero_errors[qvar] = True + self.ignore_near_zero_errors["qvapor"] = True + self.ignore_near_zero_errors["qliquid"] = True + self.ignore_near_zero_errors["qice"] = True + self.ignore_near_zero_errors["qrain"] = True + self.ignore_near_zero_errors["qsnow"] = True + self.ignore_near_zero_errors["qgraupel"] = True + self.ignore_near_zero_errors["qo3mr"] = True + self.ignore_near_zero_errors["qsgs_tke"] = True + self.ignore_near_zero_errors["qcld"] = True self.ignore_near_zero_errors["q_con"] = True self.dycore: Optional[fv_dynamics.DynamicalCore] = None self.stencil_factory = stencil_factory + self._quantity_factory = QuantityFactory.from_backend( + sizer=stencil_factory.grid_indexing._sizer, + backend=stencil_factory.backend, + ) self.namelist: DynamicalCoreConfig = DynamicalCoreConfig.from_namelist(namelist) def state_from_inputs(self, inputs): input_storages = super().state_from_inputs(inputs) + # extract tracers + input_storages["vapor"] = input_storages.pop("qvapor") + input_storages["liquid"] = input_storages.pop("qliquid") + input_storages["ice"] = input_storages.pop("qice") + input_storages["rain"] = input_storages.pop("qrain") + input_storages["snow"] = input_storages.pop("qsnow") + input_storages["graupel"] = input_storages.pop("qgraupel") + input_storages["o3mr"] = input_storages.pop("qo3mr") + input_storages["sgs_tke"] = input_storages.pop("qsgs_tke") + input_storages["cloud"] = input_storages.pop("qcld") # making sure we init DycoreState with the exact set of variables accepted_keys = [_field.name for _field in fields(DycoreState)] + accepted_keys += TRACERS_IN_PYFV3 todelete = [] for name in input_storages.keys(): if name not in accepted_keys: todelete.append(name) for name in todelete: del input_storages[name] - - state = DycoreState.init_from_storages(input_storages, sizer=self.grid.sizer) + state = DycoreState.init_from_storages( + storages=input_storages, + quantity_factory=self._quantity_factory, + tracer_list=TRACERS_IN_PYFV3, + ) return state def prepare_data(self, inputs) -> Tuple[DycoreState, GridData]: @@ -340,6 +390,7 @@ def compute_parallel(self, inputs, communicator): config=DynamicalCoreConfig.from_namelist(self.namelist), phis=state.phis, state=state, + exclude_tracers=["cloud"], timestep=timedelta(seconds=inputs["bdt"]), ) self.dycore.step_dynamics(state, NullTimer()) @@ -352,7 +403,10 @@ def outputs_from_state(self, state: dict): outputs = {} storages = {} for name, properties in self.outputs.items(): - if isinstance(state[name], Quantity): + if name in TRACERS_IN_FORTRAN: + idx = TRACERS_IN_FORTRAN.index(name) + storages[name] = state["tracers"][TRACERS_IN_PYFV3[idx]].data + elif isinstance(state[name], Quantity): storages[name] = state[name].data elif len(self.outputs[name]["dims"]) > 0: storages[name] = state[name] # assume it's a storage From bd8380a6ff3e1fb18f78a495ed7f45705a4b8a99 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 16 Dec 2024 17:19:58 -0500 Subject: [PATCH 055/252] Lint --- pyFV3/stencils/mapn_tracer.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index ea4f3f05..ef4aafc7 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -1,12 +1,13 @@ from typing import List +import dace + from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.stencils.map_single import MapSingle from pyFV3.tracers import Tracers -import dace class MapNTracer: From a41c7c77e9c3650a91e13a03433faabbd37fe940 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 17 Dec 2024 16:41:00 -0800 Subject: [PATCH 056/252] Added python implementation of mpp_global_sum within a translate test framework. --- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_getMPIprop.py | 43 ++++- .../translate/translate_mpp_global_sum.py | 181 ++++++++++++++++++ 3 files changed, 223 insertions(+), 4 deletions(-) create mode 100644 tests/savepoint/translate/translate_mpp_global_sum.py diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 4e43310f..e956b33d 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -98,4 +98,5 @@ from .translate_Pressures_mapU import TranslatePressures_mapU from .translate_Pressures_mapV import TranslatePressures_mapV from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln -from .translate_getMPIprop import TranslateGetMPIProp \ No newline at end of file +from .translate_getMPIprop import TranslateGetMPIProp +from .translate_mpp_global_sum import TranslateMpp_global_sum \ No newline at end of file diff --git a/tests/savepoint/translate/translate_getMPIprop.py b/tests/savepoint/translate/translate_getMPIprop.py index ac4b29c0..ba5e1186 100644 --- a/tests/savepoint/translate/translate_getMPIprop.py +++ b/tests/savepoint/translate/translate_getMPIprop.py @@ -2,7 +2,8 @@ from ndsl.stencils.testing.grid import Grid from ndsl.stencils.testing import ParallelTranslate from ndsl.typing import Communicator -from mpi4py import MPI +from ndsl.quantity import Quantity +import numpy as np class TranslateGetMPIProp(ParallelTranslate): def __init__( @@ -24,9 +25,45 @@ def __init__( } } + len_k = 10 + + a = [1,2,3,4,5] + + self._testQuantity_1D = Quantity( + data=np.array(a,dtype=np.float32), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + self._testQuantity_2D = Quantity( + data=np.ones([5,5],dtype=np.float32), + dims=["I","J"], + units="dunno2", + gt4py_backend=stencil_factory.backend, + ) + + self._testQuantity_3D = Quantity( + data=np.ones([3,3,3],dtype=np.float32), + dims=["I","J","K"], + units="dunno3", + gt4py_backend=stencil_factory.backend, + ) + def compute_parallel(self, inputs, communicator: Communicator): print("Communicator rank = ", communicator.rank) print("Communicator size = ", communicator.size) - global_sum = communicator.comm.allreduce(communicator.rank, op=MPI.SUM) - print("global_sum of ranks = ", global_sum) + print("self._testQuantity = ", self._testQuantity_1D.data) + global_sum_q = communicator.all_reduce_sum(self._testQuantity_1D) + print("global_sum_q.data = ", global_sum_q.data) + print("global_sum_q.metadata = ", global_sum_q.metadata) + + global_sum_q = communicator.all_reduce_sum(self._testQuantity_2D) + print("global_sum_q.data = ", global_sum_q.data) + print("global_sum_q.metadata = ", global_sum_q.metadata) + + global_sum_q = communicator.all_reduce_sum(self._testQuantity_3D) + print("global_sum_q.data = ", global_sum_q.data) + print("global_sum_q.metadata = ", global_sum_q.metadata) + return inputs \ No newline at end of file diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py new file mode 100644 index 00000000..b90dc2bd --- /dev/null +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -0,0 +1,181 @@ +from ndsl import StencilFactory, Namelist +from ndsl.stencils.testing.grid import Grid +from ndsl.stencils.testing import TranslateFortranData2Py, ParallelTranslate +import numpy as np +from ndsl.typing import Communicator +from ndsl.quantity import Quantity + +# class TranslateMpp_global_sum(TranslateFortranData2Py): +class TranslateMpp_global_sum(ParallelTranslate): + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + # print("Base TranslateGetMPIProp is initialized") + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.grid = grid + + # self.in_vars["data_vars"] = { + # "inputArray" :{ + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # } + # } + # # self.in_vars["parameters"] = ["tesum"] + # self.out_vars = { + # "inputArray" :{ + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # } + # } + + self._base.in_vars["data_vars"] = { + "inputArray" :{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + } + } + # self.in_vars["parameters"] = ["tesum"] + self._base.out_vars = { + "inputArray" :{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + } + } + + self.NUMINT = 6 + self.NUMBIT = 46 + r_prec = 2.0 ** self.NUMBIT + self.prec = 2 ** self.NUMBIT + self.I_prec = 1.0 / (2.0 ** self.NUMBIT) + self.pr = [r_prec**2, r_prec, 1.0, 1.0/r_prec, 1.0/r_prec**2, 1.0/r_prec**3] + self.I_pr = [1.0/r_prec**2, 1.0/r_prec, 1.0, r_prec, r_prec**2, r_prec**3 ] + self.prec_error = (2**62 + (2**62 - 1)) / 6 + + + def compute_from_storage(self, inputs): + + print('shape of inputArray = ', inputs["inputArray"].shape) + print('prec_error = ', self.prec_error) + print('inputArray = ', inputs["inputArray"]) + # print('type(inputs["inputArray"]) = ', type(inputs["inputArray"])) + # self.array_manipulation(inputs["inputArray"][3:27,3:27]) + return inputs + + def compute_parallel(self, inputs, communicator: Communicator): + # print("Communicator rank = ", communicator.rank) + # print("Communicator size = ", communicator.size) + # print('inputArray = ', inputs["inputArray"]) + # print('shape of inputArray = ', inputs["inputArray"].shape) + # print("rank ", communicator.rank, ": sum(inputARray) = ", sum(sum(inputs["inputArray"]))) + # print('prec_error = ', self.prec_error) + + self.array_manipulation(inputs["inputArray"], communicator, self.stencil_factory) + return inputs + + def array_manipulation(self, inputArray, communicator=None, stencil_factory=None): + + print("rank ", communicator.rank, "sum(inputArray) = ", sum(sum(inputArray[0:24,0:24]))) + mag_max_term = 0.0 + # ints_sum = np.zeros((self.NUMINT)) + ints_sum = Quantity( + data=np.zeros((self.NUMINT),dtype=np.float64), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + for j in range(0,24): + for i in range(0,24): + self.increment_ints_faster(ints_sum.data, inputArray[i,j], mag_max_term) + + print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' before carry_over') + self.carry_overflow(ints_sum.data, self.prec_error) + print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' after carry_over') + + + if communicator == None: + print("No communcator passed, so global sum cannot be performed") + exit() + else: + ints_sum_reduce = communicator.all_reduce_sum(ints_sum) + + print("rank ", communicator.rank, "sum(ints_sum_reduce) = ", sum(ints_sum_reduce.data), ' after all_reduce') + self.regularize_ints(ints_sum_reduce.data) + print("rank ", communicator.rank,"ints_sum_reduce = ", sum(ints_sum_reduce.data), ' after regularize_ints') + + sum_ = self.ints_to_real(ints_sum_reduce.data) + + print("rank ", communicator.rank,"sum_ = ", sum_, ' after ints_to_real') + + + def increment_ints_faster(self, int_sum, r, max_mag_term): + if((r >= 1e30) == r < 1e30): + print("NaN_error") + return + sgn = 1 + if(r < 0.0): + sgn = -1 + + rs = abs(r) + if(rs > abs(max_mag_term)): + max_mag_term = r + + for i in range(0,self.NUMINT): + ival = int(rs*self.I_pr[i]) + rs = rs - ival*self.pr[i] + int_sum[i] = int_sum[i] + sgn*ival + + def carry_overflow(self, int_sum, prec_error): + for i in range(self.NUMINT-1,1,-1): + if abs(int_sum[i]) > self.prec: + num_carry = int(int_sum[i] * self.I_prec) + int_sum[i] = int_sum[i] - num_carry*self.prec + int_sum[i-1] = int_sum[i-1] + num_carry + if abs(int_sum[0]) > self.prec_error: + overflow_error = True + + def regularize_ints(self, int_sum): + for i in range(self.NUMINT-1, 0, -1): + if abs(int_sum[i]) > self.prec: + num_carry = int(int_sum[i] * self.I_prec) + int_sum[i] = int_sum[i] - num_carry*self.prec + int_sum[i-1] = int_sum[i-1] + num_carry + + positive = True + + for i in range(self.NUMINT): + if abs(int_sum[i]) > 0: + if int_sum[i] < 0: + positive = False + break + + if positive: + for i in range(self.NUMINT-1, 1, -1): + if int_sum[i] < 0: + int_sum[i] = int_sum[i] + self.prec + int_sum[i-1] = int_sum[i-1] - 1 + + else: + for i in range(self.NUMINT-1, 1, -1): + if int_sum[i] > 0: + int_sum[i] = int_sum[i] - self.prec + int_sum[i-1] = int_sum[i-1] + 1 + + def ints_to_real(self, ints): + r = 0.0 + + for i in range(self.NUMINT): + r = r + self.pr[i]*ints[i] + + return r \ No newline at end of file From 71bd36a9488cf4395b1e246d4b74f10bd4c4cc68 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 18 Dec 2024 09:30:07 -0500 Subject: [PATCH 057/252] Pass tracers exclusion to downstream --- pyFV3/stencils/remapping.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 591bb231..b81adcc6 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional +from typing import List, Optional from gt4py.cartesian.gtscript import ( __INLINED, @@ -13,7 +13,7 @@ region, ) -from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -295,6 +295,7 @@ def __init__( area_64, pfull, tracers: Tracers, + exclude_tracers: List[str], checkpointer: Optional[Checkpointer] = None, ): orchestrate( @@ -411,6 +412,7 @@ def __init__( abs(config.kord_tr), fill=config.fill, tracers=tracers, + exclude_tracers=exclude_tracers, ) self._map_single_w = MapSingle( From c808532baedc085a27f9e838cab48044e0704be3 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 18 Dec 2024 09:30:48 -0500 Subject: [PATCH 058/252] Push `use_logp` in `Fillz` --- pyFV3/stencils/dyn_core.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 530d7f26..5de8a5c9 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -460,6 +460,7 @@ def __init__( quantity_factory=quantity_factory, grid_data=grid_data, grid_type=config.grid_type, + use_logp=config.use_logp, ) ) self._akap = Float(constants.KAPPA) From 7e8c0f01829f804b8a7bb11c3da721c37c65c0ae Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 18 Dec 2024 09:32:06 -0500 Subject: [PATCH 059/252] Update branch list --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 04a8062a..b56beaa2 100644 --- a/README.md +++ b/README.md @@ -26,10 +26,10 @@ Branches: - ✅ `fix/F32/UpdateDzC`@Florian: Fix for fluxes gradient - ✅ `fix/F32/DivergenceDamping`@Florian: Fix for 32-bit scalars in DivergenceDamping +- ✅ `fix/F32/UpdateDzD`@Florian: Fix for fluxes gradient & python computation - 🔶 `fix/RayleighDamping_mixed_precision`@Florian: fix the Ray_Fast test - 🔶 `GEOS_update/yppm_xppm`@Florian: fix the YPPM/XPPM with `hord = -6` - 🔶 `fix/DelnFlux_f32_support`@Florian: Fix for f32 support for DelnFlux (partial pass) -- 🔶 `fix/F32/UpdateDzD`@Florian: Fix for fluxes gradient & python computation - ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation and column calculation (partial pass) - ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit A2B_Ord4 From b5b234700ad03c9c18469bff1dc605a7d7660297 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 18 Dec 2024 17:26:27 -0500 Subject: [PATCH 060/252] Fix orchestration (still unrolling) --- pyFV3/stencils/fillz.py | 17 +++++++------- pyFV3/stencils/mapn_tracer.py | 5 ++--- pyFV3/stencils/tracer_2d_1l.py | 41 +++++++++++++++++----------------- 3 files changed, 31 insertions(+), 32 deletions(-) diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index 6d86b20f..a0c0156f 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -154,12 +154,11 @@ def __call__( tracers (inout): tracers to fix negative masses in """ for name, tracer in tracers.items(): - if name in self._exclude_tracers: - continue - self._fix_tracer_stencil( - tracer, - dp2, - self._zfix, - self._sum0, - self._sum1, - ) + if name not in self._exclude_tracers: + self._fix_tracer_stencil( + tracer, + dp2, + self._zfix, + self._sum0, + self._sum1, + ) diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index ef4aafc7..a193a84a 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -79,9 +79,8 @@ def __call__( tracers (inout): tracers to be remapped """ for name in tracers.names(): - if name in self._exclude_tracers: - continue - self._map_single[name](tracers[name], pe1, pe2, self._qs) + if name not in self._exclude_tracers: + self._map_single[name](tracers[name], pe1, pe2, self._qs) if self._fill_negative_tracers is True: self._fillz(dp2, tracers) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index 2521a64c..d2fba074 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -399,26 +399,27 @@ def __call__( ) for name, q in tracers.items(): if name in self._exclude_tracers: - continue - self.finite_volume_transport( - q, - x_courant, - y_courant, - self._x_area_flux, - self._y_area_flux, - self._x_flux, - self._y_flux, - x_mass_flux=x_mass_flux, - y_mass_flux=y_mass_flux, - ) - self._apply_tracer_flux( - q, - dp1, - self._x_flux, - self._y_flux, - self.grid_data.rarea, - dp2, - ) + pass + else: + self.finite_volume_transport( + q, + x_courant, + y_courant, + self._x_area_flux, + self._y_area_flux, + self._x_flux, + self._y_flux, + x_mass_flux=x_mass_flux, + y_mass_flux=y_mass_flux, + ) + self._apply_tracer_flux( + q, + dp1, + self._x_flux, + self._y_flux, + self.grid_data.rarea, + dp2, + ) if not last_call: self._tracers_halo_updater.update() # we can't use variable assignment to avoid a data copy From 36a9612abe154236d4a60b6a4a79e092a58cd0c0 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 18 Dec 2024 17:28:35 -0500 Subject: [PATCH 061/252] Lint --- pyFV3/stencils/mapn_tracer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index a193a84a..2b6911a9 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -1,7 +1,5 @@ from typing import List -import dace - from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField From 81bcaebb66dc33e2d239c605a431969e1e6260d1 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Dec 2024 09:58:35 -0500 Subject: [PATCH 062/252] Default tracers for initializations --- pyFV3/initialization/test_cases/initialize_baroclinic.py | 4 +++- pyFV3/initialization/test_cases/initialize_tc.py | 5 +++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/pyFV3/initialization/test_cases/initialize_baroclinic.py b/pyFV3/initialization/test_cases/initialize_baroclinic.py index 17196057..1a4a92eb 100644 --- a/pyFV3/initialization/test_cases/initialize_baroclinic.py +++ b/pyFV3/initialization/test_cases/initialize_baroclinic.py @@ -339,9 +339,11 @@ def init_baroclinic_state( ) state = DycoreState.init_from_numpy_arrays( numpy_state.__dict__, - sizer=quantity_factory.sizer, + quantity_factory=quantity_factory, backend=sample_quantity.metadata.gt4py_backend, + tracer_list=["vapor", "liquid", "rain", "snow", "ice", "graupel", "cloud"], ) + state.tracers["vapor"].view[:] = numpy_state.qvapor[slice_3d] comm.halo_update(state.phis, n_points=NHALO) diff --git a/pyFV3/initialization/test_cases/initialize_tc.py b/pyFV3/initialization/test_cases/initialize_tc.py index 38a4a46c..a402bd48 100644 --- a/pyFV3/initialization/test_cases/initialize_tc.py +++ b/pyFV3/initialization/test_cases/initialize_tc.py @@ -561,7 +561,6 @@ def init_tc_state( numpy_state.pkz[:] = pkz numpy_state.ps[:] = pe[:, :, -1] numpy_state.pt[:] = pt - numpy_state.qvapor[:] = qvapor numpy_state.u[:] = ud numpy_state.ua[:] = ua numpy_state.v[:] = vd @@ -569,8 +568,10 @@ def init_tc_state( numpy_state.w[:] = w state = DycoreState.init_from_numpy_arrays( numpy_state.__dict__, - sizer=quantity_factory.sizer, + quantity_factory=quantity_factory, backend=sample_quantity.metadata.gt4py_backend, + tracer_list=["vapor", "liquid", "rain", "snow", "ice", "graupel", "cloud"], ) + state.tracers["vapor"].view[:] = qvapor return state From 1cea564867eb6c6fdbf4477deca3168999c052fc Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Dec 2024 09:59:35 -0500 Subject: [PATCH 063/252] Fix GEOS_WRAPPER tracer initialization, ingest, outgest Remove unused "make_mapping" --- pyFV3/tracers.py | 6 -- pyFV3/wrappers/geos_wrapper.py | 114 ++++++++++++--------------------- 2 files changed, 42 insertions(+), 78 deletions(-) diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py index 43802665..a2a6c04c 100644 --- a/pyFV3/tracers.py +++ b/pyFV3/tracers.py @@ -124,9 +124,3 @@ def make_from_4D_array( unit=cls.unit, ) return tracers - - @staticmethod - def make_mapping(tracer_data: np.ndarray): - if len(tracer_data.shape) != 4: - raise ValueError("Expected 4D field as input") - return tracer_data.shape[3] * [None] diff --git a/pyFV3/wrappers/geos_wrapper.py b/pyFV3/wrappers/geos_wrapper.py index 98feb79c..53ad91c5 100644 --- a/pyFV3/wrappers/geos_wrapper.py +++ b/pyFV3/wrappers/geos_wrapper.py @@ -10,6 +10,7 @@ from mpi4py import MPI import pyFV3 +import pyFV3.tracers from ndsl import ( CompilationConfig, CubedSphereCommunicator, @@ -36,6 +37,17 @@ from ndsl.utils import safe_assign_array +GEOS_TRACER_MAPPING = [ + "vapor", + "liquid", + "ice", + "rain", + "snow", + "graupel", + "cloud", +] + + class StencilBackendCompilerOverride: """Override the Pace global stencil JIT to allow for 9-rank build on any setup. @@ -104,8 +116,23 @@ def __init__( bdt: int, comm: Comm, backend: str, + water_tracers_count: int, + all_tracers_count: int, fortran_mem_space: MemorySpace = MemorySpace.HOST, ): + # Check for water species configuration not handled by the interface + if water_tracers_count != 6: + raise NotImplementedError( + "[pyFV3 Bridge] Bridge expect 6 water species," + f" got {water_tracers_count}." + ) + + # Build the full tracer mapping by appending None to the expected tracer list + # based on parameter + self._tracers_mapping = GEOS_TRACER_MAPPING + for i in range(all_tracers_count, len(GEOS_TRACER_MAPPING)): + self._tracers_mapping.append(f"tracer_#{i}") + # Look for an override to run on a single node gtfv3_single_rank_override = int(os.getenv("GTFV3_SINGLE_RANK_OVERRIDE", -1)) if gtfv3_single_rank_override >= 0: @@ -137,7 +164,7 @@ def __init__( metric_terms = MetricTerms( quantity_factory=quantity_factory, communicator=self.communicator, - eta_file=namelist["grid_config"]["config"]["eta_file"], + eta_file=namelist["grid_config"]["config"]["eta_file"], # type: ignore ) grid_data = GridData.new_from_metric_terms(metric_terms) @@ -173,7 +200,8 @@ def __init__( ) self.dycore_state = pyFV3.DycoreState.init_zeros( - quantity_factory=quantity_factory + quantity_factory=quantity_factory, + tracer_list=self._tracers_mapping, ) self.dycore_state.bdt = self.dycore_config.dt_atmos @@ -190,6 +218,7 @@ def __init__( timestep=timedelta(seconds=self.dycore_state.bdt), phis=self.dycore_state.phis, state=self.dycore_state, + exclude_tracers=[], ) self._fortran_mem_space = fortran_mem_space @@ -198,7 +227,6 @@ def __init__( ) self.output_dict: Dict[str, np.ndarray] = {} - self._allocate_output_dir() # Feedback information device_ordinal_info = ( @@ -368,15 +396,11 @@ def _put_fortran_data_in_dycore( safe_assign_array(state.omga.view[:], omga[isc:iec, jsc:jec, :]) safe_assign_array(state.diss_estd.view[:], diss_estd[isc:iec, jsc:jec, :]) - # tracer quantities should be a 4d array in order: - # vapor, liquid, ice, rain, snow, graupel, cloud - safe_assign_array(state.qvapor.view[:], q[isc:iec, jsc:jec, :, 0]) - safe_assign_array(state.qliquid.view[:], q[isc:iec, jsc:jec, :, 1]) - safe_assign_array(state.qice.view[:], q[isc:iec, jsc:jec, :, 2]) - safe_assign_array(state.qrain.view[:], q[isc:iec, jsc:jec, :, 3]) - safe_assign_array(state.qsnow.view[:], q[isc:iec, jsc:jec, :, 4]) - safe_assign_array(state.qgraupel.view[:], q[isc:iec, jsc:jec, :, 5]) - safe_assign_array(state.qcld.view[:], q[isc:iec, jsc:jec, :, 6]) + # Copy tracer data + for index, name in enumerate(self._tracers_mapping): + safe_assign_array( + state.tracers[name].view[:], q[isc:iec, jsc:jec, :, index] + ) return state @@ -388,6 +412,7 @@ def _prep_outputs_for_geos(self) -> Dict[str, np.ndarray]: jec = self._grid_indexing.jec + 1 if self._fortran_mem_space != self._pace_mem_space: + self._allocate_output_dir() safe_assign_array(output_dict["u"], self.dycore_state.u.data[:-1, :, :-1]) safe_assign_array(output_dict["v"], self.dycore_state.v.data[:, :-1, :-1]) safe_assign_array(output_dict["w"], self.dycore_state.w.data[:-1, :-1, :-1]) @@ -453,27 +478,8 @@ def _prep_outputs_for_geos(self) -> Dict[str, np.ndarray]: self.dycore_state.diss_estd.data[:-1, :-1, :-1], ) - safe_assign_array( - output_dict["qvapor"], self.dycore_state.qvapor.data[:-1, :-1, :-1] - ) - safe_assign_array( - output_dict["qliquid"], self.dycore_state.qliquid.data[:-1, :-1, :-1] - ) - safe_assign_array( - output_dict["qice"], self.dycore_state.qice.data[:-1, :-1, :-1] - ) - safe_assign_array( - output_dict["qrain"], self.dycore_state.qrain.data[:-1, :-1, :-1] - ) - safe_assign_array( - output_dict["qsnow"], self.dycore_state.qsnow.data[:-1, :-1, :-1] - ) - safe_assign_array( - output_dict["qgraupel"], self.dycore_state.qgraupel.data[:-1, :-1, :-1] - ) - safe_assign_array( - output_dict["qcld"], self.dycore_state.qcld.data[:-1, :-1, :-1] - ) + # Copy tracer data + safe_assign_array(output_dict["q"], self.dycore_state.tracers.as_4D_array()) else: output_dict["u"] = self.dycore_state.u.data[:-1, :, :-1] output_dict["v"] = self.dycore_state.v.data[:, :-1, :-1] @@ -504,23 +510,18 @@ def _prep_outputs_for_geos(self) -> Dict[str, np.ndarray]: output_dict["q_con"] = self.dycore_state.q_con.data[:-1, :-1, :-1] output_dict["omga"] = self.dycore_state.omga.data[:-1, :-1, :-1] output_dict["diss_estd"] = self.dycore_state.diss_estd.data[:-1, :-1, :-1] - output_dict["qvapor"] = self.dycore_state.qvapor.data[:-1, :-1, :-1] - output_dict["qliquid"] = self.dycore_state.qliquid.data[:-1, :-1, :-1] - output_dict["qice"] = self.dycore_state.qice.data[:-1, :-1, :-1] - output_dict["qrain"] = self.dycore_state.qrain.data[:-1, :-1, :-1] - output_dict["qsnow"] = self.dycore_state.qsnow.data[:-1, :-1, :-1] - output_dict["qgraupel"] = self.dycore_state.qgraupel.data[:-1, :-1, :-1] - output_dict["qcld"] = self.dycore_state.qcld.data[:-1, :-1, :-1] + output_dict["q"] = self.dycore_state.tracers.as_4D_array() return output_dict def _allocate_output_dir(self): + if len(self.output_dict) != 0: + return if self._fortran_mem_space != self._pace_mem_space: nhalo = self._grid_indexing.n_halo shape_centered = self._grid_indexing.domain_full(add=(0, 0, 0)) shape_x_interface = self._grid_indexing.domain_full(add=(1, 0, 0)) shape_y_interface = self._grid_indexing.domain_full(add=(0, 1, 0)) - shape_z_interface = self._grid_indexing.domain_full(add=(0, 0, 1)) shape_2d = shape_centered[:-1] self.output_dict["u"] = np.empty((shape_y_interface)) @@ -573,34 +574,3 @@ def _allocate_output_dir(self): self.output_dict["qsnow"] = np.empty((shape_centered)) self.output_dict["qgraupel"] = np.empty((shape_centered)) self.output_dict["qcld"] = np.empty((shape_centered)) - else: - self.output_dict["u"] = None - self.output_dict["v"] = None - self.output_dict["w"] = None - self.output_dict["ua"] = None - self.output_dict["va"] = None - self.output_dict["uc"] = None - self.output_dict["vc"] = None - self.output_dict["delz"] = None - self.output_dict["pt"] = None - self.output_dict["delp"] = None - self.output_dict["mfxd"] = None - self.output_dict["mfyd"] = None - self.output_dict["cxd"] = None - self.output_dict["cyd"] = None - self.output_dict["ps"] = None - self.output_dict["pe"] = None - self.output_dict["pk"] = None - self.output_dict["peln"] = None - self.output_dict["pkz"] = None - self.output_dict["phis"] = None - self.output_dict["q_con"] = None - self.output_dict["omga"] = None - self.output_dict["diss_estd"] = None - self.output_dict["qvapor"] = None - self.output_dict["qliquid"] = None - self.output_dict["qice"] = None - self.output_dict["qrain"] = None - self.output_dict["qsnow"] = None - self.output_dict["qgraupel"] = None - self.output_dict["qcld"] = None From eba2c8164c6fe8236536adf84a5f1355bfa5be64 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Dec 2024 11:51:56 -0500 Subject: [PATCH 064/252] Update README --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index b56beaa2..25ce0792 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ Branches: - ✅ `fix/F32/UpdateDzC`@Florian: Fix for fluxes gradient - ✅ `fix/F32/DivergenceDamping`@Florian: Fix for 32-bit scalars in DivergenceDamping - ✅ `fix/F32/UpdateDzD`@Florian: Fix for fluxes gradient & python computation +- ✅ `fix/F32/nh_p_grad` @ Florian: Fix for 32-bit NonHydrostaticPressureGradient - 🔶 `fix/RayleighDamping_mixed_precision`@Florian: fix the Ray_Fast test - 🔶 `GEOS_update/yppm_xppm`@Florian: fix the YPPM/XPPM with `hord = -6` - 🔶 `fix/DelnFlux_f32_support`@Florian: Fix for f32 support for DelnFlux (partial pass) From 470e188bc155c4a62716eb2172d83111a3129579 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Dec 2024 11:52:56 -0500 Subject: [PATCH 065/252] Fix test doubly periodic --- tests/mpi/test_doubly_periodic.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/mpi/test_doubly_periodic.py b/tests/mpi/test_doubly_periodic.py index a264ea40..5b357bbf 100644 --- a/tests/mpi/test_doubly_periodic.py +++ b/tests/mpi/test_doubly_periodic.py @@ -128,6 +128,7 @@ def setup_dycore() -> Tuple[DynamicalCore, List[Any]]: config=config, phis=state.phis, state=state, + exclude_tracers=[], timestep=timedelta(seconds=255), ) # TODO compute from namelist From c9a40e8b337d6bbd9d2aaedb636c189037372ed9 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Dec 2024 12:03:21 -0500 Subject: [PATCH 066/252] Fix translate_init_case --- .../savepoint/translate/translate_init_case.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/savepoint/translate/translate_init_case.py b/tests/savepoint/translate/translate_init_case.py index 9cd87f03..06d8fcc0 100644 --- a/tests/savepoint/translate/translate_init_case.py +++ b/tests/savepoint/translate/translate_init_case.py @@ -4,7 +4,6 @@ import pytest import ndsl.constants as constants -import ndsl.dsl.gt4py_utils as utils import pyFV3.initialization.analytic_init as analytic_init import pyFV3.initialization.init_utils as init_utils import pyFV3.initialization.test_cases.initialize_baroclinic as baroclinic_init @@ -20,7 +19,6 @@ ) from ndsl.grid import GridData, MetricTerms from ndsl.stencils.testing import ParallelTranslateBaseSlicing -from ndsl.stencils.testing.grid import TRACER_DIM # type: ignore from pyFV3.testing import TranslateDycoreFortranData2Py @@ -112,7 +110,7 @@ class TranslateInitCase(ParallelTranslateBaseSlicing): }, "q4d": { "name": "tracers", - "dims": [X_DIM, Y_DIM, Z_DIM, TRACER_DIM], + "dims": [X_DIM, Y_DIM, Z_DIM, "tracers"], "units": "kg/kg", }, } @@ -166,6 +164,10 @@ def __init__( self.ignore_near_zero_errors[var] = {"near_zero": 2e-13} self.namelist = namelist # type: ignore self.stencil_factory = stencil_factory + self._quantity_factory = QuantityFactory.from_backend( + sizer=stencil_factory.grid_indexing._sizer, + backend=stencil_factory.backend, + ) def compute_sequential(self, *args, **kwargs): pytest.skip( @@ -177,10 +179,8 @@ def outputs_from_state(self, state: dict): outputs = {} arrays = {} for name, properties in self.outputs.items(): - if isinstance(state[name], dict): - for tracer, quantity in state[name].items(): - state[name][tracer] = state[name][tracer].data - arrays[name] = state[name] + if name == "q4d": + arrays[name] = state["tracers"].as_4D_array() elif len(self.outputs[name]["dims"]) > 0: arrays[name] = state[name].data else: @@ -229,7 +229,6 @@ def compute_parallel(self, inputs, communicator): ) grid_data = GridData.new_from_metric_terms(metric_terms) - quantity_factory = QuantityFactory() state = analytic_init.init_analytic_state( analytic_init_case="baroclinic", @@ -241,9 +240,6 @@ def compute_parallel(self, inputs, communicator): comm=communicator, ) - state.q4d = {} - for tracer in utils.tracer_variables: - state.q4d[tracer] = getattr(state, tracer) return self.outputs_from_state(state.__dict__) From ac166a475d3d3b5a57c1cf7c183e6e8a5f8b8c1a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 19 Dec 2024 12:12:42 -0500 Subject: [PATCH 067/252] Update README.md --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 25ce0792..e7cc5749 100644 --- a/README.md +++ b/README.md @@ -36,3 +36,4 @@ Branches: - ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit - ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit +- ⚙️ `feature/tracer_rework_part1` @Florian: Allow for update of N Tracers From 65847d4a017598fc9a22c855cc22737cb2ad8ce9 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 19 Dec 2024 09:42:13 -0800 Subject: [PATCH 068/252] Creating savepoint --- .../translate/translate_mpp_global_sum.py | 50 ++++++++++--------- 1 file changed, 27 insertions(+), 23 deletions(-) diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index b90dc2bd..e6aef661 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -4,8 +4,6 @@ import numpy as np from ndsl.typing import Communicator from ndsl.quantity import Quantity - -# class TranslateMpp_global_sum(TranslateFortranData2Py): class TranslateMpp_global_sum(ParallelTranslate): def __init__( self, @@ -13,7 +11,6 @@ def __init__( namelist: Namelist, stencil_factory: StencilFactory, ): - # print("Base TranslateGetMPIProp is initialized") super().__init__(grid, namelist, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid @@ -42,15 +39,15 @@ def __init__( "iend": grid.ie, "jstart": grid.js, "jend": grid.je, + }, + "tesum":{ + } } - # self.in_vars["parameters"] = ["tesum"] + self._base.out_vars = { - "inputArray" :{ - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, + "tesum":{ + } } @@ -81,12 +78,17 @@ def compute_parallel(self, inputs, communicator: Communicator): # print("rank ", communicator.rank, ": sum(inputARray) = ", sum(sum(inputs["inputArray"]))) # print('prec_error = ', self.prec_error) - self.array_manipulation(inputs["inputArray"], communicator, self.stencil_factory) + # print("tesum from translate test 1 : ", inputs["tesum"], type(inputs["tesum"])) + + inputs["tesum"] = self.mpp_global_sum(inputs["inputArray"], communicator, self.stencil_factory) + + # print("tesum from translate test 2 : ", inputs["tesum"]) + return inputs - def array_manipulation(self, inputArray, communicator=None, stencil_factory=None): + def mpp_global_sum(self, inputArray, communicator, stencil_factory=None): - print("rank ", communicator.rank, "sum(inputArray) = ", sum(sum(inputArray[0:24,0:24]))) + # print("rank ", communicator.rank, "sum(inputArray) = ", sum(sum(inputArray[0:24,0:24]))) mag_max_term = 0.0 # ints_sum = np.zeros((self.NUMINT)) ints_sum = Quantity( @@ -95,28 +97,30 @@ def array_manipulation(self, inputArray, communicator=None, stencil_factory=None units="dunno", gt4py_backend=stencil_factory.backend, ) + + ints_sum_reduce = Quantity( + data=np.zeros((self.NUMINT),dtype=np.float64), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) for j in range(0,24): for i in range(0,24): self.increment_ints_faster(ints_sum.data, inputArray[i,j], mag_max_term) - print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' before carry_over') + # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' before carry_over') self.carry_overflow(ints_sum.data, self.prec_error) - print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' after carry_over') + # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' after carry_over') + communicator.all_reduce_sum(ints_sum, ints_sum_reduce) - if communicator == None: - print("No communcator passed, so global sum cannot be performed") - exit() - else: - ints_sum_reduce = communicator.all_reduce_sum(ints_sum) - - print("rank ", communicator.rank, "sum(ints_sum_reduce) = ", sum(ints_sum_reduce.data), ' after all_reduce') + # print("rank ", communicator.rank, "sum(ints_sum_reduce) = ", sum(ints_sum_reduce.data), ' after all_reduce') self.regularize_ints(ints_sum_reduce.data) - print("rank ", communicator.rank,"ints_sum_reduce = ", sum(ints_sum_reduce.data), ' after regularize_ints') + # print("rank ", communicator.rank,"ints_sum_reduce = ", sum(ints_sum_reduce.data), ' after regularize_ints') sum_ = self.ints_to_real(ints_sum_reduce.data) - print("rank ", communicator.rank,"sum_ = ", sum_, ' after ints_to_real') + return sum_ def increment_ints_faster(self, int_sum, r, max_mag_term): From 02a57c4e08afc3241d9423cdeca25c5c42cf0613 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 19 Dec 2024 12:41:11 -0800 Subject: [PATCH 069/252] Added mpp_global_sum and cooresponding routines into mpp_global_sum.py --- pyFV3/stencils/__init__.py | 1 + pyFV3/stencils/mpp_global_sum.py | 109 +++++++++++++ .../translate/translate_mpp_global_sum.py | 150 +----------------- 3 files changed, 113 insertions(+), 147 deletions(-) create mode 100644 pyFV3/stencils/mpp_global_sum.py diff --git a/pyFV3/stencils/__init__.py b/pyFV3/stencils/__init__.py index dbf4f1e3..40cbedd3 100644 --- a/pyFV3/stencils/__init__.py +++ b/pyFV3/stencils/__init__.py @@ -28,6 +28,7 @@ from .updatedzd import UpdateHeightOnDGrid from .xppm import XPiecewiseParabolic from .yppm import YPiecewiseParabolic +from .mpp_global_sum import mpp_global_sum """ diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py new file mode 100644 index 00000000..4688f823 --- /dev/null +++ b/pyFV3/stencils/mpp_global_sum.py @@ -0,0 +1,109 @@ +from ndsl.quantity import Quantity +import numpy as np + +def mpp_global_sum(inputArray, communicator, stencil_factory=None): + + NUMINT = 6 + NUMBIT = 46 + r_prec = 2.0 ** NUMBIT + prec = 2 ** NUMBIT + I_prec = 1.0 / (2.0 ** NUMBIT) + pr = [r_prec**2, r_prec, 1.0, 1.0/r_prec, 1.0/r_prec**2, 1.0/r_prec**3] + I_pr = [1.0/r_prec**2, 1.0/r_prec, 1.0, r_prec, r_prec**2, r_prec**3 ] + prec_error = (2**62 + (2**62 - 1)) / 6 + mag_max_term = 0.0 + + ints_sum = Quantity( + data=np.zeros((NUMINT),dtype=np.float64), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + ints_sum_reduce = Quantity( + data=np.zeros((NUMINT),dtype=np.float64), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + # Note: This loop range in i and j are for the TBC test case. + for j in range(inputArray.shape[1]): + for i in range(inputArray.shape[0]): + increment_ints_faster(ints_sum.data, pr, I_pr, inputArray[i,j], mag_max_term) + + # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' before carry_over') + carry_overflow(ints_sum.data, prec, I_prec, prec_error) + # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' after carry_over') + + communicator.all_reduce_sum(ints_sum, ints_sum_reduce) + + # print("rank ", communicator.rank, "sum(ints_sum_reduce) = ", sum(ints_sum_reduce.data), ' after all_reduce') + regularize_ints(ints_sum_reduce.data, prec, I_prec) + # print("rank ", communicator.rank,"ints_sum_reduce = ", sum(ints_sum_reduce.data), ' after regularize_ints') + + sum_ = ints_to_real(ints_sum_reduce.data, pr) + + return sum_ + + +def increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): + if((r >= 1e30) == r < 1e30): + print("NaN_error") + return + sgn = 1 + if(r < 0.0): + sgn = -1 + + rs = abs(r) + if(rs > abs(max_mag_term)): + max_mag_term = r + + for i in range(len(I_pr)): + ival = int(rs*I_pr[i]) + rs = rs - ival*pr[i] + int_sum[i] = int_sum[i] + sgn*ival + +def carry_overflow(int_sum, prec, I_prec, prec_error): + for i in range(len(int_sum)-1,1,-1): + if abs(int_sum[i]) > prec: + num_carry = int(int_sum[i] * I_prec) + int_sum[i] = int_sum[i] - num_carry*prec + int_sum[i-1] = int_sum[i-1] + num_carry + if abs(int_sum[0]) > prec_error: + overflow_error = True + +def regularize_ints(int_sum, prec, I_prec): + for i in range(len(int_sum)-1, 0, -1): + if abs(int_sum[i]) > prec: + num_carry = int(int_sum[i] * I_prec) + int_sum[i] = int_sum[i] - num_carry*prec + int_sum[i-1] = int_sum[i-1] + num_carry + + positive = True + + for i in range(len(int_sum)): + if abs(int_sum[i]) > 0: + if int_sum[i] < 0: + positive = False + break + + if positive: + for i in range(len(int_sum)-1, 1, -1): + if int_sum[i] < 0: + int_sum[i] = int_sum[i] + prec + int_sum[i-1] = int_sum[i-1] - 1 + + else: + for i in range(len(int_sum)-1, 1, -1): + if int_sum[i] > 0: + int_sum[i] = int_sum[i] - prec + int_sum[i-1] = int_sum[i-1] + 1 + +def ints_to_real(ints, pr): + r = 0.0 + + for i in range(len(ints)): + r = r + pr[i]*ints[i] + + return r \ No newline at end of file diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index e6aef661..6611daa7 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -4,6 +4,7 @@ import numpy as np from ndsl.typing import Communicator from ndsl.quantity import Quantity +from pyFV3.stencils.mpp_global_sum import mpp_global_sum class TranslateMpp_global_sum(ParallelTranslate): def __init__( self, @@ -15,24 +16,6 @@ def __init__( self.stencil_factory = stencil_factory self.grid = grid - # self.in_vars["data_vars"] = { - # "inputArray" :{ - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # } - # } - # # self.in_vars["parameters"] = ["tesum"] - # self.out_vars = { - # "inputArray" :{ - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # } - # } - self._base.in_vars["data_vars"] = { "inputArray" :{ "istart": grid.is_, @@ -50,136 +33,9 @@ def __init__( } } - - self.NUMINT = 6 - self.NUMBIT = 46 - r_prec = 2.0 ** self.NUMBIT - self.prec = 2 ** self.NUMBIT - self.I_prec = 1.0 / (2.0 ** self.NUMBIT) - self.pr = [r_prec**2, r_prec, 1.0, 1.0/r_prec, 1.0/r_prec**2, 1.0/r_prec**3] - self.I_pr = [1.0/r_prec**2, 1.0/r_prec, 1.0, r_prec, r_prec**2, r_prec**3 ] - self.prec_error = (2**62 + (2**62 - 1)) / 6 - - - def compute_from_storage(self, inputs): - - print('shape of inputArray = ', inputs["inputArray"].shape) - print('prec_error = ', self.prec_error) - print('inputArray = ', inputs["inputArray"]) - # print('type(inputs["inputArray"]) = ', type(inputs["inputArray"])) - # self.array_manipulation(inputs["inputArray"][3:27,3:27]) - return inputs def compute_parallel(self, inputs, communicator: Communicator): - # print("Communicator rank = ", communicator.rank) - # print("Communicator size = ", communicator.size) - # print('inputArray = ', inputs["inputArray"]) - # print('shape of inputArray = ', inputs["inputArray"].shape) - # print("rank ", communicator.rank, ": sum(inputARray) = ", sum(sum(inputs["inputArray"]))) - # print('prec_error = ', self.prec_error) - - # print("tesum from translate test 1 : ", inputs["tesum"], type(inputs["tesum"])) - - inputs["tesum"] = self.mpp_global_sum(inputs["inputArray"], communicator, self.stencil_factory) - - # print("tesum from translate test 2 : ", inputs["tesum"]) - - return inputs - - def mpp_global_sum(self, inputArray, communicator, stencil_factory=None): - - # print("rank ", communicator.rank, "sum(inputArray) = ", sum(sum(inputArray[0:24,0:24]))) - mag_max_term = 0.0 - # ints_sum = np.zeros((self.NUMINT)) - ints_sum = Quantity( - data=np.zeros((self.NUMINT),dtype=np.float64), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - - ints_sum_reduce = Quantity( - data=np.zeros((self.NUMINT),dtype=np.float64), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - for j in range(0,24): - for i in range(0,24): - self.increment_ints_faster(ints_sum.data, inputArray[i,j], mag_max_term) - - # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' before carry_over') - self.carry_overflow(ints_sum.data, self.prec_error) - # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' after carry_over') - - communicator.all_reduce_sum(ints_sum, ints_sum_reduce) - # print("rank ", communicator.rank, "sum(ints_sum_reduce) = ", sum(ints_sum_reduce.data), ' after all_reduce') - self.regularize_ints(ints_sum_reduce.data) - # print("rank ", communicator.rank,"ints_sum_reduce = ", sum(ints_sum_reduce.data), ' after regularize_ints') - - sum_ = self.ints_to_real(ints_sum_reduce.data) - - return sum_ - - - def increment_ints_faster(self, int_sum, r, max_mag_term): - if((r >= 1e30) == r < 1e30): - print("NaN_error") - return - sgn = 1 - if(r < 0.0): - sgn = -1 - - rs = abs(r) - if(rs > abs(max_mag_term)): - max_mag_term = r - - for i in range(0,self.NUMINT): - ival = int(rs*self.I_pr[i]) - rs = rs - ival*self.pr[i] - int_sum[i] = int_sum[i] + sgn*ival - - def carry_overflow(self, int_sum, prec_error): - for i in range(self.NUMINT-1,1,-1): - if abs(int_sum[i]) > self.prec: - num_carry = int(int_sum[i] * self.I_prec) - int_sum[i] = int_sum[i] - num_carry*self.prec - int_sum[i-1] = int_sum[i-1] + num_carry - if abs(int_sum[0]) > self.prec_error: - overflow_error = True - - def regularize_ints(self, int_sum): - for i in range(self.NUMINT-1, 0, -1): - if abs(int_sum[i]) > self.prec: - num_carry = int(int_sum[i] * self.I_prec) - int_sum[i] = int_sum[i] - num_carry*self.prec - int_sum[i-1] = int_sum[i-1] + num_carry - - positive = True - - for i in range(self.NUMINT): - if abs(int_sum[i]) > 0: - if int_sum[i] < 0: - positive = False - break - - if positive: - for i in range(self.NUMINT-1, 1, -1): - if int_sum[i] < 0: - int_sum[i] = int_sum[i] + self.prec - int_sum[i-1] = int_sum[i-1] - 1 + inputs["tesum"] = mpp_global_sum(inputs["inputArray"], communicator, self.stencil_factory) - else: - for i in range(self.NUMINT-1, 1, -1): - if int_sum[i] > 0: - int_sum[i] = int_sum[i] - self.prec - int_sum[i-1] = int_sum[i-1] + 1 - - def ints_to_real(self, ints): - r = 0.0 - - for i in range(self.NUMINT): - r = r + self.pr[i]*ints[i] - - return r \ No newline at end of file + return inputs \ No newline at end of file From 598e2cb7f43cd076f8e9921ee979a0bbf259f6ec Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 20 Dec 2024 13:50:02 -0800 Subject: [PATCH 070/252] Starting assembly of fv_mapz remapping translate test --- fv_mapz.F90.SER | 160 +++++++--- .../translate/translate_remapping_GEOS.py | 288 ++++++++++-------- 2 files changed, 277 insertions(+), 171 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index c2b2f28a..a58d6549 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -230,6 +230,8 @@ contains real, dimension(is:ie, js:je) :: ws_2d, gz_2d + logical :: serial_flag + !$ser verbatim real :: w_max_, w_min_ !$ser verbatim integer:: mode, abskord, iep1, iedp1, jedp1, js2d, o3mr, sgs_tke @@ -392,27 +394,34 @@ contains !$ser verbatim enddo !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d pe1_=pe1_3d pe2_=pe2_3d ptop=ptop -!$ser data qvapor_=q(:,j,:,sphum) -!$ser data qliquid_=q(:,j,:,liq_wat) -!$ser data qice_=q(:,j,:,ice_wat) -!$ser data qrain_=q(:,j,:,rainwat) -!$ser data qsnow_=q(:,j,:,snowwat) -!$ser data qgraupel_=q(:,j,:,graupel) -!$ser data q_con=q_con -!$ser data pt=pt +!$ser data qvapor_js=q(:,j,:,sphum) +!$ser data qliquid_js=q(:,j,:,liq_wat) +!$ser data qice_js=q(:,j,:,ice_wat) +!$ser data qrain_js=q(:,j,:,rainwat) +!$ser data qsnow_js=q(:,j,:,snowwat) +!$ser data qgraupel_js=q(:,j,:,graupel) +!$ser data qcld_js=q(:,j,:,cld_amt) +!$ser data gz1d=gz +!$ser data cvm=cvm +!$ser data r_vir=r_vir !$ser data cappa=cappa +!$ser data rrg=rrg !$ser data delp=delp !$ser data delz=delz -!$ser data r_vir=r_vir -!$ser data remap_t=remap_t -!$ser data ps=ps -!$ser data peln_3d=peln_3d -!$ser data pn2_3d=pn2_3d -!$ser data ak=ak -!$ser data bk=bk -!$ser data dp2_3d=dp2_3d -!$ser data pk=pk -!$ser data akap=akap +!$ser data pt=pt +!$ser data k1k=k1k +!$ser data j_2d=js2d +!$ser data q_con=q_con + +!!$ser data remap_t=remap_t +!!$ser data ps=ps +!!$ser data peln_3d=peln_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data ak=ak +!!$ser data bk=bk +!!$ser data dp2_3d=dp2_3d +!!$ser data pk=pk +!!$ser data akap=akap !$ser verbatim endif do k=1,km+1 @@ -442,6 +451,7 @@ contains else ! Transform "density pt" to "density temp" !$ser verbatim if(j == js2d) then + !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' !$ser savepoint MoistCVPlusPt_2d-In !$ser data qvapor_js=q(:,j,:,sphum) qliquid_js=q(:,j,:,liq_wat) qice_js=q(:,j,:,ice_wat) qrain_js=q(:,j,:,rainwat) qsnow_js=q(:,j,:,snowwat) qgraupel_js=q(:,j,:,graupel) qcld_js=q(:,j,:,cld_amt) gz1d=gz cvm=cvm r_vir=r_vir cappa=cappa rrg=rrg delp=delp delz=delz pt=pt k1k=k1k j_2d=js2d q_con=q_con !$ser verbatim endif @@ -464,6 +474,32 @@ contains !$ser savepoint MoistCVPlusPt_2d-Out !$ser data gz1d=gz cvm=cvm pt=pt cappa=cappa q_con=q_con !$ser verbatim endif + + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) + !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d + !$ser data pe2_=pe2_3d + !$ser data gz1d=gz + !$ser data cvm=cvm + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data q_con=q_con + !$ser verbatim endif + endif ! hydro test elseif (remap_pt) then print*,'CODE NOT TESTED HERE 3' @@ -583,32 +619,32 @@ contains enddo enddo -!$ser verbatim if(j == js) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe1_=pe1_3d pe2_=pe2_3d -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data dp2_3d=dp2_3d -!$ser data pn2_3d=pn2_3d -!$ser data pk=pk2_3d -!$ser verbatim endif +!!$ser verbatim if(j == js) then +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie +!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km +!!$ser verbatim do i = is, ie +!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe1_=pe1_3d pe2_=pe2_3d +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data dp2_3d=dp2_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data pk=pk2_3d +!!$ser verbatim endif if (remap_te) then print*,'CODE NOT TESTED HERE 5' @@ -1372,7 +1408,11 @@ contains !$OMP do_adiabatic_init,zsum1,zsum0,te0_2d,domain, & !$OMP ng,gridstruct,E_Flux,pdt,dtmp,reproduce_sum,q, & !$OMP mdt,cld_amt,cappa,dtdt,out_dt,rrg,akap,do_sat_adj, & -!$OMP fast_mp_consv,kord_tm) & +!$OMP fast_mp_consv,kord_tm, & +#ifdef SERIALIZE +!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, serial_flag & +#endif +!$OMP ) & !$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp) dtmp = 0. @@ -1453,9 +1493,33 @@ contains enddo ! j-loop !$OMP single - ! if(j==js) print*,"MPP GLOBAL SUM CODE EXECUTED 1" + !print*,"MPP GLOBAL SUM CODE EXECUTED 1" + !$ser savepoint Mpp_global_sum-In + !$ser verbatim serial_flag=.true. + !$ser verbatim j=24 + !$ser data x_compute_size=j + !$ser data y_compute_size=j + !$ser verbatim j=1 + !$ser data x_compute_begin=j + !$ser data y_compute_begin=j + !$ser data max_ntile_pe=j + !$ser data tile=j + !$ser verbatim j=0 + !$ser data ioff=j + !$ser data joff=j + !$ser data serial_flag=serial_flag + !$ser data tesum=tesum tesum = mpp_global_sum(domain, te_2d*gridstruct%area_64(is:ie,js:je), & flags=BITWISE_EFP_SUM) + + !$ser verbatim te_2d=te_2d*gridstruct%area_64(is:ie,js:je) + !!$ser verbatim print*,'Sum of input into mpp_global_sum = ', sum(te_2d) + !$ser data inputArray=te_2d + + !$ser savepoint Mpp_global_sum-Out + !$ser data tesum=tesum + + !print*,'tesum = ', tesum E_Flux = DBLE(consv)*tesum / DBLE(grav*pdt*4.*pi*radius**2) ! unit: W/m**2 ! Note pdt is "phys" time step if ( hydrostatic ) then @@ -1623,6 +1687,10 @@ contains endif !$OMP end parallel +!$ser savepoint GetMPIProp-In +!$ser data delz=delz +!$ser savepoint GetMPIProp-Out +!$ser data delz=delz end subroutine Lagrangian_to_Eulerian !>@brief The subroutine 'compute_total_energy' performs the FV3-consistent computation of the global total energy. diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 903fb817..2a54fb64 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -13,9 +13,39 @@ Z_DIM, Z_INTERFACE_DIM, ) -from ndsl.dsl.typing import Float +from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.mapn_tracer import MapNTracer +def moist_pt( + qvapor: FloatField, + qliquid: FloatField, + qrain: FloatField, + qsnow: FloatField, + qice: FloatField, + qgraupel: FloatField, + q_con: FloatField, + pt: FloatField, + cappa: FloatField, + delp: FloatField, + delz: FloatField, + r_vir: Float, +): + with computation(PARALLEL), interval(...): + cvm, gz, q_con, cappa, pt = moist_cv.moist_pt_func( + qvapor, + qliquid, + qrain, + qsnow, + qice, + qgraupel, + q_con, + pt, + cappa, + delp, + delz, + r_vir, + ) + class TranslateRemapping_GEOS(TranslateDycoreFortranData2Py): def __init__( self, @@ -28,12 +58,11 @@ def __init__( self.in_vars["data_vars"] = { "pe_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, + "istart": grid.is_-1, + "iend": grid.ie+1, + "jstart": grid.js-1, + "jend": grid.je+1, "kend": grid.npz + 1, - # "kaxis": 1, }, "pe1_": { "istart": grid.is_, @@ -41,7 +70,6 @@ def __init__( "jstart": grid.js, "jend": grid.je, "kend": grid.npz + 1, - # "kaxis": 1, }, "pe2_": { "istart": grid.is_, @@ -49,64 +77,53 @@ def __init__( "jstart": grid.js, "jend": grid.je, "kend": grid.npz + 1, - # "kaxis": 1, - }, - # Note that tracers are i x k shaped. - # Setting "axis" as 1 enables the translate test to read them properly - "qvapor_": { - "axis": 1 - }, - "qliquid_": { - "axis": 1 - }, - "qice_": { - "axis": 1 - }, - "qrain_": { - "axis": 1 - }, - "qsnow_": { - "axis": 1 - }, - "qgraupel_": { - "axis": 1 }, + "qvapor": {"serialname": "qvapor_js"}, + "qliquid": {"serialname": "qliquid_js"}, + "qice": {"serialname": "qice_js"}, + "qrain": {"serialname": "qrain_js"}, + "qsnow": {"serialname": "qsnow_js"}, + "qgraupel": {"serialname": "qgraupel_js"}, "delp": {}, "delz": {}, "q_con": {}, "pt": {}, "cappa": {}, - "ps": {}, - "pn2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "peln_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "ak": {}, - "bk": {}, - "dp2_3d": grid.compute_dict(), - "pk": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - } + # "ps": {}, + # "pn2_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "peln_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "ak": {}, + # "bk": {}, + # "dp2_3d": grid.compute_dict(), + # "pk": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # } } + self.write_vars = ["gz", "cvm"] + for k, v in self.in_vars["data_vars"].items(): + if k not in self.write_vars: + v["axis"] = 1 self.in_vars["parameters"] = [ "ptop", "r_vir", # "remap_t", # For some reason, translate test can't accept a logical variable - "akap", + # "akap", # "zvir", # "last_step", # "consv_te", @@ -120,7 +137,6 @@ def __init__( "jstart": grid.js, "jend": grid.je, "kend": grid.npz + 1, - # "kaxis": 1, }, "pe2_": { "istart": grid.is_, @@ -128,29 +144,28 @@ def __init__( "jstart": grid.js, "jend": grid.je, "kend": grid.npz + 1, - # "kaxis": 1, }, - "delp": {}, - "delz": {}, - "q_con": {}, "pt": {}, "cappa": {}, - "ps": {}, - "dp2_3d": grid.compute_dict(), - "pn2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pk": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - } + "q_con": {}, + # "delp": {}, + # "delz": {}, + # "ps": {}, + # "dp2_3d": grid.compute_dict(), + # "pn2_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "pk": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # } } self.stencil_factory = stencil_factory @@ -184,34 +199,40 @@ def __init__( self._init_pe = stencil_factory.from_origin_domain( init_pe, - # origin=(3,3,0), origin=grid_indexing.origin_compute(), - domain=(grid.nic,1,73), + # domain=(grid.nic,1,73), + domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1), ) - self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( - moist_cv_pt_pressure, - # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, - externals={"hydrostatic": hydrostatic}, - origin=grid_indexing.origin_compute(), - # domain=grid_indexing.domain_compute(add=(0, 0, 1)), - domain=(grid.nic, 1, grid.npz+1), # Note : Many intervals go from (0,-1) in this stencil + self._moist_cv_pt = stencil_factory.from_origin_domain( + moist_pt, + origin=grid.compute_origin(), + domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2]), ) - self._pn2_pk_delp = stencil_factory.from_origin_domain( - pn2_pk_delp, - origin=grid_indexing.origin_compute(), - domain=(grid.nic, 1, grid.npz+1), - ) + # self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( + # moist_cv_pt_pressure, + # # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, + # externals={"hydrostatic": hydrostatic}, + # origin=grid_indexing.origin_compute(), + # # domain=grid_indexing.domain_compute(add=(0, 0, 1)), + # domain=(grid.nic, 1, grid.npz+1), # Note : Many intervals go from (0,-1) in this stencil + # ) - self._compute_func = MapNTracer( - self.stencil_factory, - self.quantity_factory, - abs(self.kord), - self.nq, - fill=self.fill, - tracers=tracers, - ) + # self._pn2_pk_delp = stencil_factory.from_origin_domain( + # pn2_pk_delp, + # origin=grid_indexing.origin_compute(), + # domain=(grid.nic, 1, grid.npz+1), + # ) + + # self._compute_func = MapNTracer( + # self.stencil_factory, + # self.quantity_factory, + # abs(self.kord), + # self.nq, + # fill=self.fill, + # tracers=tracers, + # ) def compute_from_storage(self, inputs): @@ -223,6 +244,8 @@ def compute_from_storage(self, inputs): value, self.grid.njd, backend=self.stencil_factory.backend ) ) + # print("name = ", name) + # print("value.shape = ", value.shape) # print("inputs[qvapor].data.shape() 2 = ", inputs["qvapor"].data.shape) self._init_pe( inputs["pe_"], @@ -231,43 +254,58 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - self._moist_cv_pt_pressure( - inputs["qvapor_"], - inputs["qliquid_"], - inputs["qrain_"], - inputs["qsnow_"], - inputs["qice_"], - inputs["qgraupel_"], + self._moist_cv_pt( + inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], inputs["q_con"], inputs["pt"], inputs["cappa"], inputs["delp"], inputs["delz"], - inputs["pe_"], - inputs["pe2_"], - inputs["ak"], - inputs["bk"], - inputs["dp2_3d"], - inputs["ps"], - inputs["pn2_3d"], - inputs["peln_3d"], - True, - Float(inputs["r_vir"]), + inputs["r_vir"], ) - self._pn2_pk_delp( - inputs["pe2_"], - inputs["pn2_3d"], - inputs["pk"], - Float(inputs["akap"]), - ) + # self._moist_cv_pt_pressure( + # inputs["qvapor_"], + # inputs["qliquid_"], + # inputs["qrain_"], + # inputs["qsnow_"], + # inputs["qice_"], + # inputs["qgraupel_"], + # inputs["q_con"], + # inputs["pt"], + # inputs["cappa"], + # inputs["delp"], + # inputs["delz"], + # inputs["pe_"], + # inputs["pe2_"], + # inputs["ak"], + # inputs["bk"], + # inputs["dp2_3d"], + # inputs["ps"], + # inputs["pn2_3d"], + # inputs["peln_3d"], + # True, + # Float(inputs["r_vir"]), + # ) + + # self._pn2_pk_delp( + # inputs["pe2_"], + # inputs["pn2_3d"], + # inputs["pk"], + # Float(inputs["akap"]), + # ) - # now that we have the pressure profiles, we can start remapping - self._map_single_pt(pt, peln, self._pn2, qmin=self._t_min) + # # now that we have the pressure profiles, we can start remapping + # self._map_single_pt(pt, peln, self._pn2, qmin=self._t_min) - self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) + # self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) - self._map_single_w(w, self._pe1, self._pe2, qs=wsd) - self._map_single_delz(delz, self._pe1, self._pe2) + # self._map_single_w(w, self._pe1, self._pe2, qs=wsd) + # self._map_single_delz(delz, self._pe1, self._pe2) return inputs From f1cd30092fcbe686dedda6db2b0c65de14ddbd37 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 20 Dec 2024 17:02:38 -0500 Subject: [PATCH 071/252] Allow blind mapping from data --- pyFV3/tracers.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py index a2a6c04c..4b938c5f 100644 --- a/pyFV3/tracers.py +++ b/pyFV3/tracers.py @@ -101,6 +101,12 @@ def make( tracers._quantities[name] = qty return tracers + @staticmethod + def blind_mapping_from_data(tracer_data: np.ndarray): + if len(tracer_data.shape) != 4: + raise ValueError("Expected 4D field as input") + return [f"Tracer_{idx}" for idx in range(tracer_data.shape[3])] + @classmethod def make_from_4D_array( cls, From d2a3cf28d2f77146b1f6b52332cb9d637be25136 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 20 Dec 2024 17:05:05 -0500 Subject: [PATCH 072/252] Allow N tracers to be advected from translate test (based on input) Fix `cmax` and `nsplit` calculation with allreduce Fix f32 mass conservation calculation Allow for mass_fluxes and courant numbers to _not_ be udapted (GEOS default) --- pyFV3/stencils/tracer_2d_1l.py | 314 ++++++++++++------ tests/savepoint/translate/__init__.py | 1 + .../translate/translate_tracer2d1l.py | 40 +-- .../translate/translate_tracer2d1l_cmax.py | 87 +++++ 4 files changed, 322 insertions(+), 120 deletions(-) create mode 100644 tests/savepoint/translate/translate_tracer2d1l_cmax.py diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index d2fba074..26147e5b 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -1,10 +1,23 @@ -import math from typing import List import gt4py.cartesian.gtscript as gtscript -from gt4py.cartesian.gtscript import PARALLEL, computation, horizontal, interval, region +from gt4py.cartesian.gtscript import ( + PARALLEL, + computation, + horizontal, + interval, + region, + i32, +) -from ndsl import QuantityFactory, StencilFactory, WrappedHaloUpdater, orchestrate +from ndsl import ( + QuantityFactory, + StencilFactory, + WrappedHaloUpdater, + orchestrate, + Quantity, +) +from ndsl.grid import GridData from ndsl.constants import ( N_HALO_DEFAULT, X_DIM, @@ -13,7 +26,7 @@ Y_INTERFACE_DIM, Z_DIM, ) -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.typing import FloatField, FloatFieldIJ, FloatFieldK from ndsl.typing import Communicator from pyFV3.stencils.fvtp2d import FiniteVolumeTransport from pyFV3.tracers import Tracers @@ -82,10 +95,11 @@ def divide_fluxes_by_n_substeps( cyd: FloatField, yfx: FloatField, mfyd: FloatField, - n_split: int, + cmax: FloatFieldK, ): """ - Divide all inputs in-place by the number of substeps n_split. + Divide all inputs in-place by the number of substeps n_split computed + from the max courant number on the grid Args: cxd (inout): @@ -96,25 +110,15 @@ def divide_fluxes_by_n_substeps( mfyd (inout): """ with computation(PARALLEL), interval(...): - frac = 1.0 / n_split - cxd = cxd * frac - xfx = xfx * frac - mfxd = mfxd * frac - cyd = cyd * frac - yfx = yfx * frac - mfyd = mfyd * frac - - -def cmax_stencil1(cx: FloatField, cy: FloatField, cmax: FloatField): - with computation(PARALLEL), interval(...): - cmax = max(abs(cx), abs(cy)) - - -def cmax_stencil2( - cx: FloatField, cy: FloatField, sin_sg5: FloatField, cmax: FloatField -): - with computation(PARALLEL), interval(...): - cmax = max(abs(cx), abs(cy)) + 1.0 - sin_sg5 + n_split = i32(1.0 + cmax) + if n_split > 1: + frac = 1.0 / n_split + cxd = cxd * frac + xfx = xfx * frac + mfxd = mfxd * frac + cyd = cyd * frac + yfx = yfx * frac + mfyd = mfyd * frac def apply_mass_flux( @@ -135,7 +139,10 @@ def apply_mass_flux( with computation(PARALLEL), interval(...): dp2 = ( dp1 - + (x_mass_flux - x_mass_flux[1, 0, 0] + y_mass_flux - y_mass_flux[0, 1, 0]) + + ( + (x_mass_flux - x_mass_flux[1, 0, 0]) + + (y_mass_flux - y_mass_flux[0, 1, 0]) + ) * rarea ) @@ -147,6 +154,8 @@ def apply_tracer_flux( fy: FloatField, rarea: FloatFieldIJ, dp2: FloatField, + cmax: FloatFieldK, + current_nsplit: int, ): """ Args: @@ -158,7 +167,8 @@ def apply_tracer_flux( dp2 (in): """ with computation(PARALLEL), interval(...): - q = (q * dp1 + (fx - fx[1, 0, 0] + fy - fy[0, 1, 0]) * rarea) / dp2 + if current_nsplit < i32(1.0 + cmax): + q = (q * dp1 + ((fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) * rarea) / dp2 # Simple stencil replacing: @@ -178,6 +188,16 @@ class TracerAdvection: Performs horizontal advection on tracers. Corresponds to tracer_2D_1L in the Fortran code. + + Args: + stencil_factory: Stencil maker built on the required grid + quantity_factory: Quantity maker built on the required grid + transport: The Finite Volume to be applied to each tracers + grid_data: Metric Terms for the grid + comm: Communicator on the grid + tracers: Bundle of data of tracers to be advected + exclude_tracers: Tracers to not be advected + update_mass_courant: update the mass and courant numbers """ def __init__( @@ -185,10 +205,11 @@ def __init__( stencil_factory: StencilFactory, quantity_factory: QuantityFactory, transport: FiniteVolumeTransport, - grid_data, + grid_data: GridData, comm: Communicator, tracers: Tracers, exclude_tracers: List[str], + update_mass_courant: bool = True, ): orchestrate( obj=self, @@ -197,39 +218,55 @@ def __init__( ) grid_indexing = stencil_factory.grid_indexing self.grid_indexing = grid_indexing # needed for selective validation - self._tracer_count = tracers.count self.grid_data = grid_data self._exclude_tracers = exclude_tracers + self._update_mass_courant = update_mass_courant + + if not self._update_mass_courant: + self._tmp_mfx = quantity_factory.zeros( + [X_INTERFACE_DIM, Y_DIM, Z_DIM], + units="unknown", + ) + self._tmp_mfy = quantity_factory.zeros( + [X_DIM, Y_INTERFACE_DIM, Z_DIM], + units="unknown", + ) + self._tmp_cx = quantity_factory.zeros( + [X_INTERFACE_DIM, Y_DIM, Z_DIM], + units="unknown", + ) + self._tmp_cy = quantity_factory.zeros( + [X_DIM, Y_INTERFACE_DIM, Z_DIM], + units="unknown", + ) self._x_area_flux = quantity_factory.zeros( [X_INTERFACE_DIM, Y_DIM, Z_DIM], units="unknown", - dtype=Float, ) self._y_area_flux = quantity_factory.zeros( [X_DIM, Y_INTERFACE_DIM, Z_DIM], units="unknown", - dtype=Float, ) self._x_flux = quantity_factory.zeros( [X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], units="unknown", - dtype=Float, ) self._y_flux = quantity_factory.zeros( [X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], units="unknown", - dtype=Float, ) self._tmp_dp = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="Pa", - dtype=Float, ) self._tmp_dp2 = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="Pa", - dtype=Float, + ) + self._cmax = quantity_factory.zeros( + [Z_DIM], + units="unitless", ) ax_offsets = grid_indexing.axis_offsets( @@ -277,7 +314,6 @@ def __init__( tracer_halo_spec = quantity_factory.get_quantity_halo_spec( dims=[X_DIM, Y_DIM, Z_DIM], n_halo=N_HALO_DEFAULT, - dtype=Float, ) # We can exclude tracers from advecting and therefore also @@ -293,6 +329,14 @@ def __init__( [t for t in advected_tracers.keys()], ) + # Setup tracer courant max reduction calculation + self._compute_cmax = TracerCMax( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + grid_data=grid_data, + comm=comm, + ) + def __call__( self, tracers: Tracers, @@ -317,17 +361,25 @@ def __call__( x_courant (inout): accumulated courant number in x-direction y_courant (inout): accumulated courant number in y-direction """ - # DaCe parsing issue - # if len(tracers) != self._tracer_count: - # raise ValueError( - # f"incorrect number of tracers, {self._tracer_count} was " - # f"specified on init but {len(tracers)} were passed" - # ) - # start HALO update on q (in dyn_core in fortran -- just has started when - # this function is called...) + + if self._update_mass_courant: + working_x_mass_flux = x_mass_flux + working_y_mass_flux = x_mass_flux + working_x_courant = x_courant + working_y_courant = y_courant + else: + self._tmp_mfx.data[:] = x_mass_flux.data[:] + self._tmp_mfy.data[:] = y_mass_flux.data[:] + self._tmp_cx.data[:] = x_courant.data[:] + self._tmp_cy.data[:] = y_courant.data[:] + working_x_mass_flux = self._tmp_mfx + working_y_mass_flux = self._tmp_mfy + working_x_courant = self._tmp_cx + working_y_courant = self._tmp_cy + self._flux_compute( - x_courant, - y_courant, + working_x_courant, + working_y_courant, self.grid_data.dxa, self.grid_data.dya, self.grid_data.dx, @@ -336,64 +388,49 @@ def __call__( self.grid_data.sin_sg2, self.grid_data.sin_sg3, self.grid_data.sin_sg4, - # TODO: rename xfx/yfx to "area flux" self._x_area_flux, self._y_area_flux, ) - # # TODO for if we end up using the Allreduce and compute cmax globally - # (or locally). For now, hardcoded. - # split = int(grid_indexing.domain[2] / 6) - # self._cmax_1( - # cxd, cyd, self._tmp_cmax, origin=grid_indexing.origin_compute(), - # domain=(grid_indexing.domain[0], self.grid_indexing.domain[1], split) - # ) - # self._cmax_2( - # cxd, - # cyd, - # self.grid.sin_sg5, - # self._tmp_cmax, - # origin=(grid_indexing.isc, self.grid_indexing.jsc, split), - # domain=( - # grid_indexing.domain[0], - # self.grid_indexing.domain[1], - # grid_indexing.domain[2] - split + 1 - # ), - # ) - # cmax_flat = np.amax(self._tmp_cmax, axis=(0, 1)) - # # cmax_flat is a gt4py storage still, but of dimension [npz+1]... - - # cmax_max_all_ranks = cmax_flat.data - # # TODO mpi allreduce... - # # comm.Allreduce(cmax_flat, cmax_max_all_ranks, op=MPI.MAX) - - cmax_max_all_ranks = 2.0 - n_split = math.floor(1.0 + cmax_max_all_ranks) - # NOTE: cmax is not usually a single value, it varies with k, if return to - # that, make n_split a column as well - - if n_split > 1.0: - self._divide_fluxes_by_n_substeps( - x_courant, - self._x_area_flux, - x_mass_flux, - y_courant, - self._y_area_flux, - y_mass_flux, - n_split, - ) + self._compute_cmax( + cx=working_x_courant, + cy=working_y_courant, + cmax=self._cmax, + ) + + self._divide_fluxes_by_n_substeps( + working_x_courant, + self._x_area_flux, + working_x_mass_flux, + working_y_courant, + self._y_area_flux, + working_y_mass_flux, + self._cmax, + ) self._tracers_halo_updater.update() dp2 = self._tmp_dp - for it in range(n_split): - last_call = it == n_split - 1 + # The original algorithm works on K level independantly + # (from with a K loop) and therefore compute `nsplit` + # per K + # The stencil nature of the framework doesn't allow for it + # because after advection, an halo exchange need to be carried + # (or else we could just move the test within the stencil). + # We overcompute to retain true parallelization, by running + # a loop on the highest number of nsplit, but restraining + # actual update in `apply_tracer_flux` to only the valid + # K level for each tracers + max_n_split = i32(1.0 + self._cmax.view[:].max()) + + for current_nsplit in range(int(max_n_split)): + last_call = current_nsplit == max_n_split - 1 # tracer substep self._apply_mass_flux( dp1, - x_mass_flux, - y_mass_flux, + working_x_mass_flux, + working_y_mass_flux, self.grid_data.rarea, dp2, ) @@ -403,8 +440,8 @@ def __call__( else: self.finite_volume_transport( q, - x_courant, - y_courant, + working_x_courant, + working_y_courant, self._x_area_flux, self._y_area_flux, self._x_flux, @@ -419,9 +456,98 @@ def __call__( self._y_flux, self.grid_data.rarea, dp2, + cmax=self._cmax, + current_nsplit=current_nsplit, ) if not last_call: self._tracers_halo_updater.update() # we can't use variable assignment to avoid a data copy # because of current dace limitations self._swap_dp(dp1, dp2) + + +def cmax_stencil_low_k( + cx: FloatField, + cy: FloatField, + cmax: FloatField, +): + with computation(PARALLEL), interval(...): + cmax = max(abs(cx), abs(cy)) + + +def cmax_stencil_high_k( + cx: FloatField, + cy: FloatField, + sin_sg5: FloatFieldIJ, + cmax: FloatField, +): + with computation(PARALLEL), interval(...): + cmax = max(abs(cx), abs(cy)) + 1.0 - sin_sg5 + + +class TracerCMax: + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + grid_data: GridData, + comm: Communicator, + ): + """Perform global courant number max. + + The maximum courant number for every atmospheric level on the entire grid. + """ + self._grid_data = grid_data + self._comm = comm + grid_indexing = stencil_factory.grid_indexing + cmax_atmospheric_level_split = int(grid_indexing.domain[2] / 6) - 1 + self._cmax_low_k = stencil_factory.from_origin_domain( + func=cmax_stencil_low_k, + origin=grid_indexing.origin_compute(), + domain=( + grid_indexing.domain[0], + grid_indexing.domain[1], + cmax_atmospheric_level_split, + ), + ) + self._cmax_high_k = stencil_factory.from_origin_domain( + func=cmax_stencil_high_k, + origin=( + grid_indexing.origin_compute()[0], + grid_indexing.origin_compute()[1], + cmax_atmospheric_level_split, + ), + domain=( + grid_indexing.domain[0], + grid_indexing.domain[1], + grid_indexing.domain[2] - cmax_atmospheric_level_split, + ), + ) + self._tmp_cmax = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="unknown", + ) + self._tmp_cmax_in_K = quantity_factory.zeros( + [Z_DIM], + units="unknown", + ) + + def __call__(self, cx: Quantity, cy: Quantity, cmax: Quantity): + if __debug__: + if not isinstance(cmax, Quantity): + raise TypeError( + f"[pyFV3][Tracer]: cmax must be a quantity, got {type(cmax)}" + ) + self._cmax_low_k( + cx=cx, + cy=cy, + cmax=self._tmp_cmax, + ) + self._cmax_high_k( + cx=cx, + cy=cy, + sin_sg5=self._grid_data.sin_sg5, + cmax=self._tmp_cmax, + ) + self._tmp_cmax_in_K.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] + self._comm.buffer_all_reduce(self._tmp_cmax_in_K, cmax) diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 651b2551..e69bbd79 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -79,6 +79,7 @@ from .translate_riem_solver_c import TranslateRiem_Solver_C from .translate_satadjust3d import TranslateSatAdjust3d from .translate_tracer2d1l import TranslateTracer2D1L +from .translate_tracer2d1l_cmax import TranslateTracerCMax from .translate_updatedzc import TranslateUpdateDzC from .translate_updatedzd import TranslateUpdateDzD from .translate_xppm import TranslateXPPM, TranslateXPPM_2 diff --git a/tests/savepoint/translate/translate_tracer2d1l.py b/tests/savepoint/translate/translate_tracer2d1l.py index dcc9c104..90cfd600 100644 --- a/tests/savepoint/translate/translate_tracer2d1l.py +++ b/tests/savepoint/translate/translate_tracer2d1l.py @@ -26,10 +26,10 @@ def __init__( self._base.in_vars["data_vars"] = { "tracers": {}, "dp1": {}, - "mfxd": grid.x3d_compute_dict(), - "mfyd": grid.y3d_compute_dict(), - "cxd": grid.x3d_compute_domain_y_dict(), - "cyd": grid.y3d_compute_domain_x_dict(), + "mfxd_R4": grid.x3d_compute_dict(), + "mfyd_R4": grid.y3d_compute_dict(), + "cxd_R4": grid.x3d_compute_domain_y_dict(), + "cyd_R4": grid.y3d_compute_domain_x_dict(), } self._base.in_vars["parameters"] = ["nq"] self._base.out_vars = self._base.in_vars["data_vars"] @@ -52,17 +52,7 @@ def collect_input_data(self, serializer, savepoint): def compute_parallel(self, inputs, communicator): tracers = Tracers.make_from_4D_array( quantity_factory=self._quantity_factory, - tracer_mapping=[ - "vapor", - "liquid", - "rain", - "ice", - "snow", - "graupel", - "o3mr", - "sgs_tke", - "cloud", - ], + tracer_mapping=Tracers.blind_mapping_from_data(inputs["tracers"]), tracer_data=inputs["tracers"], ) self._base.make_storage_data_input_vars(inputs, dict_4d=False) @@ -85,19 +75,17 @@ def compute_parallel(self, inputs, communicator): communicator, tracers, exclude_tracers=["cloud"], + update_mass_courant=False, ) - inputs["x_mass_flux"] = inputs.pop("mfxd") - inputs["y_mass_flux"] = inputs.pop("mfyd") - inputs["x_courant"] = inputs.pop("cxd") - inputs["y_courant"] = inputs.pop("cyd") + inputs["x_mass_flux"] = inputs.pop("mfxd_R4") + inputs["y_mass_flux"] = inputs.pop("mfyd_R4") + inputs["x_courant"] = inputs.pop("cxd_R4") + inputs["y_courant"] = inputs.pop("cyd_R4") self.tracer_advection(tracers=tracers, **inputs) - inputs["mfxd"] = inputs.pop("x_mass_flux") - inputs["mfyd"] = inputs.pop("y_mass_flux") - inputs["cxd"] = inputs.pop("x_courant") - inputs["cyd"] = inputs.pop("y_courant") - # Put back un-advected tracers - # Tracers have -1 on all cartesian because of NDSL padding - # Dev note: qcld is not advected in Pace dataset for some reason + inputs["mfxd_R4"] = inputs.pop("x_mass_flux") + inputs["mfyd_R4"] = inputs.pop("y_mass_flux") + inputs["cxd_R4"] = inputs.pop("x_courant") + inputs["cyd_R4"] = inputs.pop("y_courant") inputs["tracers"] = tracers.as_4D_array() outputs = self._base.slice_output(inputs) outputs["tracers"] = self.subset_output("tracers", outputs["tracers"]) diff --git a/tests/savepoint/translate/translate_tracer2d1l_cmax.py b/tests/savepoint/translate/translate_tracer2d1l_cmax.py new file mode 100644 index 00000000..3eff9416 --- /dev/null +++ b/tests/savepoint/translate/translate_tracer2d1l_cmax.py @@ -0,0 +1,87 @@ +from ndsl import Namelist, StencilFactory, QuantityFactory, Quantity +from ndsl.stencils.testing import ParallelTranslate2Py +from pyFV3.stencils.tracer_2d_1l import TracerCMax +from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM + + +def _quantity_wrap(storage, dims, grid_indexing): + origin, extent = grid_indexing.get_origin_domain(dims) + return Quantity( + storage, + dims=dims, + units="unknown", + origin=origin, + extent=extent, + ) + + +class TranslateTracerCMax(ParallelTranslate2Py): + inputs = { + "cx_R4": { + "name": "cx_R4", + "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], + "units": "unitless", + }, + "cy_R4": { + "name": "cy_R4", + "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], + "units": "unitless", + }, + "cmax": { + "name": "cmaxgrid", + "dims": [Z_DIM], + "units": "unitless", + }, + } + + def __init__( + self, + grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, namelist, stencil_factory) + self._base.in_vars["data_vars"] = { + "cx_R4": grid.x3d_compute_domain_y_dict(), + "cy_R4": grid.y3d_compute_domain_x_dict(), + "cmax": {}, + } + self._base.out_vars = { + "cmax": {}, + } + self._stencil_factory = stencil_factory + self._grid_data = grid + self._quantity_factory = QuantityFactory.from_backend( + sizer=stencil_factory.grid_indexing._sizer, + backend=stencil_factory.backend, + ) + + def compute_parallel(self, inputs, communicator): + self._base.make_storage_data_input_vars(inputs) + tracer_cmax = TracerCMax( + stencil_factory=self._stencil_factory, + quantity_factory=self._quantity_factory, + grid_data=self._grid_data, + comm=communicator, + ) + cx_quantity = _quantity_wrap( + inputs["cx_R4"], + self.inputs["cx_R4"]["dims"], + self.grid.grid_indexing, + ) + cy_quantity = _quantity_wrap( + inputs["cy_R4"], + self.inputs["cy_R4"]["dims"], + self.grid.grid_indexing, + ) + cmax_quantity = _quantity_wrap( + inputs["cmax"], + self.inputs["cmax"]["dims"], + self.grid.grid_indexing, + ) + tracer_cmax( + cx=cx_quantity, + cy=cy_quantity, + cmax=cmax_quantity, + ) + return self._base.slice_output(inputs) From ca61aeeef8331b492f1dd0bc8e2c5f8f79e8e6f1 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 20 Dec 2024 17:07:10 -0500 Subject: [PATCH 073/252] Update readme --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index e7cc5749..1709e3ca 100644 --- a/README.md +++ b/README.md @@ -37,3 +37,4 @@ Branches: - ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit - ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit - ⚙️ `feature/tracer_rework_part1` @Florian: Allow for update of N Tracers +- ⚙️ `fix/GEOS/TracerAdvection` @Florian [BASED ON `tracer_rework_part1`]: Allow for non-update of mass fluxes and courant number, f32 fixes, correct computation of `cmax` and `nsplit`, overcomputation into the algorithm From 3de7c8e951d894aa0e649c05a129de73b0cd05e0 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 23 Dec 2024 18:22:53 -0800 Subject: [PATCH 074/252] Adding outline of fv_mapz translate test based on GEOS --- pyFV3/stencils/remapping.py | 6 +- .../translate/translate_remapping_GEOS.py | 323 ++++++++++++++---- 2 files changed, 255 insertions(+), 74 deletions(-) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index e245f8c0..7ec28c81 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -126,6 +126,8 @@ def moist_cv_pt_pressure( ps (out): pn2 (out): peln (in): + remap_t (in): + r_vir (in): """ from __externals__ import hydrostatic#, kord_tm @@ -149,8 +151,8 @@ def moist_cv_pt_pressure( ) # NOTE : GEOS does not perform the delz computation at this location # # delz_adjust - if __INLINED(not hydrostatic): - delz = -delz / delp + # if __INLINED(not hydrostatic): + # delz = -delz / delp # pressure_updates with computation(FORWARD): diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 2a54fb64..8b0f8553 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -2,6 +2,7 @@ from ndsl import Namelist, StencilFactory from pyFV3 import DynamicalCoreConfig from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure, pn2_pk_delp +from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils import moist_cv from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py @@ -89,23 +90,23 @@ def __init__( "q_con": {}, "pt": {}, "cappa": {}, - # "ps": {}, - # "pn2_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "peln_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "ak": {}, - # "bk": {}, + "ps": {}, + "pn2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, + "peln_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, + "ak": {'dumbass': 3, 'axis': 2}, + # "bk": {"kend": grid.npz + 1}, # "dp2_3d": grid.compute_dict(), # "pk": { # "istart": grid.is_, @@ -148,24 +149,24 @@ def __init__( "pt": {}, "cappa": {}, "q_con": {}, - # "delp": {}, - # "delz": {}, - # "ps": {}, - # "dp2_3d": grid.compute_dict(), - # "pn2_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "pk": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # } + "delp": {}, + "delz": {}, + "ps": {}, + "dp2_3d": grid.compute_dict(), + "pn2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + } } self.stencil_factory = stencil_factory @@ -210,14 +211,14 @@ def __init__( domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2]), ) - # self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( - # moist_cv_pt_pressure, - # # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, - # externals={"hydrostatic": hydrostatic}, - # origin=grid_indexing.origin_compute(), - # # domain=grid_indexing.domain_compute(add=(0, 0, 1)), - # domain=(grid.nic, 1, grid.npz+1), # Note : Many intervals go from (0,-1) in this stencil - # ) + self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( + moist_cv_pt_pressure, + # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, + externals={"hydrostatic": hydrostatic}, + origin=grid_indexing.origin_compute(), + # domain=grid_indexing.domain_compute(add=(0, 0, 1)), + domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2]+1), # Note : Many intervals go from (0,-1) in this stencil + ) # self._pn2_pk_delp = stencil_factory.from_origin_domain( # pn2_pk_delp, @@ -225,7 +226,15 @@ def __init__( # domain=(grid.nic, 1, grid.npz+1), # ) - # self._compute_func = MapNTracer( + # self._map_scalar = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # self._kord_tm, + # self.mode, + # dims=[X_DIM, Y_DIM, Z_DIM], + # ) + + # self._mapn_tracer = MapNTracer( # self.stencil_factory, # self.quantity_factory, # abs(self.kord), @@ -234,6 +243,66 @@ def __init__( # tracers=tracers, # ) + # self._map_single_w = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # self._kord_wz, + # -2, + # dims=[X_DIM, Y_DIM, Z_DIM], + # ) + + # self._map_single_delz = MapSingle( + # self.stencil_factory, + # quantity_factory, + # self._kord_wz, + # 1, + # dims=[X_DIM, Y_DIM, Z_DIM], + # ) + + # self._w_fix_consrv_moment = stencil_factory.from_origin_domain( + # func=W_fix_consrv_moment, + # origin=(3,3,0), + # domain=(24,1,72), + # ) + + # self._pressures_mapu = stencil_factory.from_origin_domain( + # pressures_mapu, + # origin=grid_indexing.origin_compute(), + # domain=self._domain_jextra, + # ) + + # self._map1_ppm_u = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # inputs["kord_mt"], + # -1, + # dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + # ) + + # self._pressures_mapv = stencil_factory.from_origin_domain( + # pressures_mapv, + # origin=grid_indexing.origin_compute(), + # domain=( + # grid_indexing.domain[0] + 1, + # grid_indexing.domain[1], + # grid_indexing.domain[2] + 1, + # ), + # ) + + # self._map1_ppm_v = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # inputs["kord_mt"], + # -1, + # dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + # ) + + # self._pe_pk_delp_peln = stencil_factory.from_origin_domain( + # pe_pk_delp_peln, + # origin=grid_indexing.origin_compute(), + # domain=self._domain_kextra, + # ) + def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction @@ -254,7 +323,22 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - self._moist_cv_pt( + # self._moist_cv_pt( + # inputs["qvapor"], + # inputs["qliquid"], + # inputs["qrain"], + # inputs["qsnow"], + # inputs["qice"], + # inputs["qgraupel"], + # inputs["q_con"], + # inputs["pt"], + # inputs["cappa"], + # inputs["delp"], + # inputs["delz"], + # inputs["r_vir"], + # ) + + self._moist_cv_pt_pressure( inputs["qvapor"], inputs["qliquid"], inputs["qrain"], @@ -266,33 +350,18 @@ def compute_from_storage(self, inputs): inputs["cappa"], inputs["delp"], inputs["delz"], - inputs["r_vir"], + inputs["pe_"], + inputs["pe2_"], + inputs["ak"], + inputs["bk"], + inputs["dp2_3d"], + inputs["ps"], + inputs["pn2_3d"], + inputs["peln_3d"], + True, + Float(inputs["r_vir"]), ) - # self._moist_cv_pt_pressure( - # inputs["qvapor_"], - # inputs["qliquid_"], - # inputs["qrain_"], - # inputs["qsnow_"], - # inputs["qice_"], - # inputs["qgraupel_"], - # inputs["q_con"], - # inputs["pt"], - # inputs["cappa"], - # inputs["delp"], - # inputs["delz"], - # inputs["pe_"], - # inputs["pe2_"], - # inputs["ak"], - # inputs["bk"], - # inputs["dp2_3d"], - # inputs["ps"], - # inputs["pn2_3d"], - # inputs["peln_3d"], - # True, - # Float(inputs["r_vir"]), - # ) - # self._pn2_pk_delp( # inputs["pe2_"], # inputs["pn2_3d"], @@ -301,11 +370,121 @@ def compute_from_storage(self, inputs): # ) # # now that we have the pressure profiles, we can start remapping - # self._map_single_pt(pt, peln, self._pn2, qmin=self._t_min) + # self._map_scalar(pt, self._pn1, self._pn2, qmin=self._t_min, interp=True) # self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) - # self._map_single_w(w, self._pe1, self._pe2, qs=wsd) + # self._map_single_w(w, self._pe1, self._pe2, qs=wsd, interp=False) # self._map_single_delz(delz, self._pe1, self._pe2) + # self._w_fix_consrv_moment( + # inputs["w"], + # inputs["w2"], + # inputs["dp2_W"], + # self._gz, + # inputs["w_max"], + # inputs["w_min"], + # self._compute_performed + # ) + + # self._pressures_mapu( + # inputs["pe_"], + # inputs["ak"], + # inputs["bk"], + # inputs["pe0_"], + # inputs["pe3_"], + # inputs["ptop"], + # ) + + # self._map1_ppm_u( + # inputs["u_"], + # inputs["pe0_"], + # inputs["pe3_"], + # interp=False, + # ) + + # self._map1_ppm_u( + # inputs["mfy_"], + # inputs["pe0_"], + # inputs["pe3_"], + # interp=False, + # ) + + # self._map1_ppm_u( + # inputs["cy_"], + # inputs["pe0_"], + # inputs["pe3_"], + # interp=False, + # ) + + # self._pressures_mapv( + # inputs["pe_"], + # inputs["ak"], + # inputs["bk"], + # inputs["pe0_v"], + # inputs["pe3_v"], + # ) + + # self._map1_ppm_v( + # inputs["v_"], + # inputs["pe0_v"], + # inputs["pe3_v"], + # interp=False, + # ) + + # self._map1_ppm_v( + # inputs["mfx_"], + # inputs["pe0_v"], + # inputs["pe3_v"], + # interp=False, + # ) + + # self._map1_ppm_v( + # inputs["cx_"], + # inputs["pe0_v"], + # inputs["pe3_v"], + # interp=False, + # ) + + # self._pe_pk_delp_peln(inputs["pe_"], + # inputs["pk"], + # inputs["delp"], + # inputs["peln_"], + # inputs["pe2_"], + # inputs["pk2_"], + # inputs["pn2_"], + # inputs["ak"], + # inputs["bk"], + # inputs["akap"], + # inputs["ptop"], + # ) + + # NOTE : THERE WILL BE ADJUSTMENTS TO ACCOUNT FOR PKZ + # self._moist_cv_pt( + # inputs["qvapor"], + # inputs["qliquid"], + # inputs["qrain"], + # inputs["qsnow"], + # inputs["qice"], + # inputs["qgraupel"], + # inputs["q_con"], + # inputs["pt"], + # inputs["cappa"], + # inputs["delp"], + # inputs["delz"], + # inputs["r_vir"], + # ) + + # If loop based on if( last_step .and. (.not.do_adiabatic_init) ) then + # PHIS computation + # Some variation of moist_cv_pt + # zsum1 computation + + # MPP GLOBAL SUM + # E_FLUX calcuation + + # If loop based on if ( last_step .and. (.not. adiabatic) ) then + # Some variation of moist_cv_pt + # Condensation update + return inputs From e6230b1a6c6b2faca329685cb1a865ae9d30cbeb Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 24 Dec 2024 16:24:24 -0500 Subject: [PATCH 075/252] `dpx` calculation (bad validation) --- pyFV3/stencils/d_sw.py | 23 ++++++++ pyFV3/stencils/dyn_core.py | 63 +++++++++++++-------- tests/savepoint/translate/translate_d_sw.py | 3 +- 3 files changed, 63 insertions(+), 26 deletions(-) diff --git a/pyFV3/stencils/d_sw.py b/pyFV3/stencils/d_sw.py index ceb67302..9e32de6f 100644 --- a/pyFV3/stencils/d_sw.py +++ b/pyFV3/stencils/d_sw.py @@ -203,6 +203,16 @@ def apply_pt_delp_fluxes_stencil_defn( pt, delp = apply_pt_delp_fluxes(gx, gy, rarea, fx, fy, pt, delp) +def delp_increment_accumulation( + dpx: FloatField64, + fx: FloatField, + fy: FloatField, + rarea: FloatFieldIJ, +): + with computation(PARALLEL), interval(...): + dpx = dpx + ((fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) * rarea + + def compute_kinetic_energy( vc: FloatField, uc: FloatField, @@ -1025,6 +1035,10 @@ def make_quantity(): da_min=damping_coefficients.da_min_c, nord=self._column_namelist["nord_w"], ) + self._accumulate_delp = stencil_factory.from_dims_halo( + func=delp_increment_accumulation, + compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], + ) def __call__( self, @@ -1043,6 +1057,7 @@ def __call__( mfy: FloatField64, cx: FloatField64, cy: FloatField64, + dpx: FloatField64, crx: FloatField, cry: FloatField, xfx: FloatField, @@ -1079,6 +1094,7 @@ def __call__( mfy (inout): accumulated y mass flux cx (inout): accumulated Courant number in the x direction cy (inout): accumulated Courant number in the y direction + dpx (inout): accumulated delp export for Dry Mass Roundoff Control crx (out): local courant number in the x direction cry (out): local courant number in the y direction xfx (out): flux of area in x-direction, in units of m^2 @@ -1217,6 +1233,13 @@ def __call__( self._adjust_w_and_qcon_stencil( w, delp, self._tmp_dw, q_con, self._column_namelist["damp_w"] ) + + self._accumulate_delp( + dpx=dpx, + fx=self._tmp_fx, + fy=self._tmp_fy, + rarea=self.grid_data.rarea, + ) # at this point, pt, delp, w and q_con have been stepped forward in time # the rest of this function updates the winds self._compute_kinetic_energy( diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 530d7f26..c1670c23 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -203,7 +203,18 @@ def dyncore_temporaries( quantity_factory: QuantityFactory, ) -> Mapping[str, Quantity]: temporaries: Dict[str, Quantity] = {} - for name in ["ut", "vt", "gz", "zh", "pem", "pkc", "pk3", "heat_source", "cappa"]: + for name in [ + "ut", + "vt", + "gz", + "zh", + "pem", + "pkc", + "pk3", + "heat_source", + "cappa", + "dpx", + ]: # TODO: the dimensions of ut and vt may not be correct, # because they are not used. double-check and correct as needed. temporaries[name] = quantity_factory.zeros( @@ -480,6 +491,7 @@ def __init__( self._xfx = temporaries["xfx"] self._yfx = temporaries["yfx"] self._ws3 = temporaries["ws3"] + self._dpx = temporaries["dpx"] if not config.hydrostatic: self._pk3.data[:] = HUGE_R @@ -873,30 +885,31 @@ def __call__( # by 1 timestep self._checkpoint_dsw_in(state) self.dgrid_shallow_water_lagrangian_dynamics( - self._vt, - state.delp, - state.pt, - state.u, - state.v, - state.w, - state.uc, - state.vc, - state.ua, - state.va, - self._divgd, - state.mfxd, - state.mfyd, - state.cxd, - state.cyd, - self._crx, - self._cry, - self._xfx, - self._yfx, - state.q_con, - self._zh, - self._heat_source, - state.diss_estd, - dt_acoustic_substep, + delpc=self._vt, + delp=state.delp, + pt=state.pt, + u=state.u, + v=state.v, + w=state.w, + uc=state.uc, + vc=state.vc, + ua=state.ua, + va=state.va, + divgd=self._divgd, + mfx=state.mfxd, + mfy=state.mfyd, + cx=state.cxd, + cy=state.cyd, + dpx=self._dpx, + crx=self._crx, + cry=self._cry, + xfx=self._xfx, + yfx=self._yfx, + q_con=state.q_con, + zh=self._zh, + heat_source=self._heat_source, + diss_est=state.diss_estd, + dt=dt_acoustic_substep, ) self._checkpoint_dsw_out(state) # note that uc and vc are not needed at all past this point. diff --git a/tests/savepoint/translate/translate_d_sw.py b/tests/savepoint/translate/translate_d_sw.py index 350e245d..ea042712 100644 --- a/tests/savepoint/translate/translate_d_sw.py +++ b/tests/savepoint/translate/translate_d_sw.py @@ -48,6 +48,7 @@ def __init__( "mfy": grid.y3d_compute_dict() | {"serialname": "mfyd_R8"}, "cx": grid.x3d_compute_domain_y_dict() | {"serialname": "cxd_R8"}, "cy": grid.y3d_compute_domain_x_dict() | {"serialname": "cyd_R8"}, + "dpx": grid.compute_dict(), "heat_source": {}, "diss_est": {}, "q_con": {}, @@ -58,7 +59,7 @@ def __init__( "divgd": grid.default_dict_buffer_2d(), } for name, info in self.in_vars["data_vars"].items(): - if name not in ["mfx", "mfy", "cx", "cy"]: + if name not in ["mfx", "mfy", "cx", "cy", "dpx"]: info["serialname"] = name + "d" self.in_vars["parameters"] = ["dt"] self.out_vars = self.in_vars["data_vars"].copy() From b6a1ea6b0e58ea7c77ecc6f48dedf7b322063a83 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 24 Dec 2024 16:31:24 -0500 Subject: [PATCH 076/252] Readme: update `D_SW` to mention new `dpx` field --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 1709e3ca..87c0edaa 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ Branches: - 🔶 `fix/RayleighDamping_mixed_precision`@Florian: fix the Ray_Fast test - 🔶 `GEOS_update/yppm_xppm`@Florian: fix the YPPM/XPPM with `hord = -6` - 🔶 `fix/DelnFlux_f32_support`@Florian: Fix for f32 support for DelnFlux (partial pass) -- ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation and column calculation (partial pass) +- ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation, column calculation and new `dpx` accumulation (partial pass) - ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit From bef0b4f3c8a1a0da52c34bcd1721f92401d550b4 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 24 Dec 2024 14:27:15 -0800 Subject: [PATCH 077/252] Modified translate test that address issue with reading in array for FloatFieldK and FloatFieldIJ data types --- .../translate/translate_remapping_GEOS.py | 126 +++++++++--------- 1 file changed, 65 insertions(+), 61 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 8b0f8553..dd2f8cdb 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -105,20 +105,22 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, - "ak": {'dumbass': 3, 'axis': 2}, - # "bk": {"kend": grid.npz + 1}, - # "dp2_3d": grid.compute_dict(), - # "pk": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # } + "ak": {}, + "bk": {}, + "dp2_3d": grid.compute_dict(), + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + } } - self.write_vars = ["gz", "cvm"] + # self.write_vars = ["gz", "cvm"] + self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel"] for k, v in self.in_vars["data_vars"].items(): - if k not in self.write_vars: + # if k not in self.write_vars: + if k in self.write_vars: v["axis"] = 1 self.in_vars["parameters"] = [ "ptop", @@ -149,24 +151,29 @@ def __init__( "pt": {}, "cappa": {}, "q_con": {}, - "delp": {}, - "delz": {}, - "ps": {}, - "dp2_3d": grid.compute_dict(), - "pn2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pk": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - } + # "delp": {}, + # "delz": {}, + # "ps": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed, + # }, + # "dp2_3d": grid.compute_dict(), + # "pn2_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "pk": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # } } self.stencil_factory = stencil_factory @@ -200,8 +207,7 @@ def __init__( self._init_pe = stencil_factory.from_origin_domain( init_pe, - origin=grid_indexing.origin_compute(), - # domain=(grid.nic,1,73), + origin=grid_indexing.origin_compute(), domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1), ) @@ -313,9 +319,7 @@ def compute_from_storage(self, inputs): value, self.grid.njd, backend=self.stencil_factory.backend ) ) - # print("name = ", name) - # print("value.shape = ", value.shape) - # print("inputs[qvapor].data.shape() 2 = ", inputs["qvapor"].data.shape) + self._init_pe( inputs["pe_"], inputs["pe1_"], @@ -323,22 +327,7 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - # self._moist_cv_pt( - # inputs["qvapor"], - # inputs["qliquid"], - # inputs["qrain"], - # inputs["qsnow"], - # inputs["qice"], - # inputs["qgraupel"], - # inputs["q_con"], - # inputs["pt"], - # inputs["cappa"], - # inputs["delp"], - # inputs["delz"], - # inputs["r_vir"], - # ) - - self._moist_cv_pt_pressure( + self._moist_cv_pt( inputs["qvapor"], inputs["qliquid"], inputs["qrain"], @@ -350,18 +339,33 @@ def compute_from_storage(self, inputs): inputs["cappa"], inputs["delp"], inputs["delz"], - inputs["pe_"], - inputs["pe2_"], - inputs["ak"], - inputs["bk"], - inputs["dp2_3d"], - inputs["ps"], - inputs["pn2_3d"], - inputs["peln_3d"], - True, - Float(inputs["r_vir"]), + inputs["r_vir"], ) + # self._moist_cv_pt_pressure( + # inputs["qvapor"], + # inputs["qliquid"], + # inputs["qrain"], + # inputs["qsnow"], + # inputs["qice"], + # inputs["qgraupel"], + # inputs["q_con"], + # inputs["pt"], + # inputs["cappa"], + # inputs["delp"], + # inputs["delz"], + # inputs["pe_"], + # inputs["pe2_"], + # inputs["ak_3d"], + # inputs["bk_3d"], + # inputs["dp2_3d"], + # inputs["ps_"], + # inputs["pn2_3d"], + # inputs["peln_3d"], + # True, + # Float(inputs["r_vir"]), + # ) + # self._pn2_pk_delp( # inputs["pe2_"], # inputs["pn2_3d"], From 7f51a6fed9a73c556129565847be4dde27db9ba9 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 26 Dec 2024 11:45:30 -0800 Subject: [PATCH 078/252] Added initial moistcv plus pressure adjustment stencils. Verified with C24 test case --- .../translate/translate_remapping_GEOS.py | 145 ++++++++++-------- 1 file changed, 82 insertions(+), 63 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index dd2f8cdb..06078ee3 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -107,14 +107,27 @@ def __init__( }, "ak": {}, "bk": {}, - "dp2_3d": grid.compute_dict(), + "dp2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, "pk": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, "kend": grid.npz + 1, - } + }, + "pk2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, } # self.write_vars = ["gz", "cvm"] self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel"] @@ -126,7 +139,7 @@ def __init__( "ptop", "r_vir", # "remap_t", # For some reason, translate test can't accept a logical variable - # "akap", + "akap", # "zvir", # "last_step", # "consv_te", @@ -151,29 +164,30 @@ def __init__( "pt": {}, "cappa": {}, "q_con": {}, - # "delp": {}, - # "delz": {}, - # "ps": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.jsd, - # "jend": grid.jed, - # }, - # "dp2_3d": grid.compute_dict(), - # "pn2_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "pk": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # } + "delp": {}, + "delz": {}, + "ps": {}, + "dp2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + "pn2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, + "pk2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + } } self.stencil_factory = stencil_factory @@ -226,11 +240,11 @@ def __init__( domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2]+1), # Note : Many intervals go from (0,-1) in this stencil ) - # self._pn2_pk_delp = stencil_factory.from_origin_domain( - # pn2_pk_delp, - # origin=grid_indexing.origin_compute(), - # domain=(grid.nic, 1, grid.npz+1), - # ) + self._pn2_pk_delp = stencil_factory.from_origin_domain( + pn2_pk_delp, + origin=grid_indexing.origin_compute(add=(0,0,1)), + domain=(grid.nic, 1, grid.npz-1), + ) # self._map_scalar = MapSingle( # self.stencil_factory, @@ -309,6 +323,11 @@ def __init__( # domain=self._domain_kextra, # ) + print("grid.compute_dict() = ", grid.compute_dict()) + print("istard iend jstart jend kend = ", grid.is_, grid.ie, grid.js, grid.je, grid.npz) + print("grid.nic, 1, grid.npz+1 : ", grid.nic, 1, grid.npz+1) + print("grid_indexing.origin_compute() = ", grid_indexing.origin_compute()) + def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction @@ -327,22 +346,7 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - self._moist_cv_pt( - inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - inputs["q_con"], - inputs["pt"], - inputs["cappa"], - inputs["delp"], - inputs["delz"], - inputs["r_vir"], - ) - - # self._moist_cv_pt_pressure( + # self._moist_cv_pt( # inputs["qvapor"], # inputs["qliquid"], # inputs["qrain"], @@ -354,24 +358,39 @@ def compute_from_storage(self, inputs): # inputs["cappa"], # inputs["delp"], # inputs["delz"], - # inputs["pe_"], - # inputs["pe2_"], - # inputs["ak_3d"], - # inputs["bk_3d"], - # inputs["dp2_3d"], - # inputs["ps_"], - # inputs["pn2_3d"], - # inputs["peln_3d"], - # True, - # Float(inputs["r_vir"]), + # inputs["r_vir"], # ) - # self._pn2_pk_delp( - # inputs["pe2_"], - # inputs["pn2_3d"], - # inputs["pk"], - # Float(inputs["akap"]), - # ) + self._moist_cv_pt_pressure( + inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + inputs["q_con"], + inputs["pt"], + inputs["cappa"], + inputs["delp"], + inputs["delz"], + inputs["pe_"], + inputs["pe2_"], + inputs["ak"], + inputs["bk"], + inputs["dp2_3d"], + inputs["ps"], + inputs["pn2_3d"], + inputs["peln_3d"], + True, + Float(inputs["r_vir"]), + ) + + self._pn2_pk_delp( + inputs["pe2_"], + inputs["pn2_3d"], + inputs["pk2_3d"], + Float(inputs["akap"]), + ) # # now that we have the pressure profiles, we can start remapping # self._map_scalar(pt, self._pn1, self._pn2, qmin=self._t_min, interp=True) From 7027a4a20e80b38d39b8ca0ebac6759c8a791c8c Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 26 Dec 2024 11:47:11 -0800 Subject: [PATCH 079/252] Added fv_mapz serialization file --- fv_mapz.F90.SER | 100 +++++++++++++++++++++++++++++++----------------- 1 file changed, 64 insertions(+), 36 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index a58d6549..4fc0c893 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -393,35 +393,36 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser savepoint Remapping_GEOS-In -!$ser data pe_=pe_3d pe1_=pe1_3d pe2_=pe2_3d ptop=ptop +!$ser data pe_=pe_3d +!$ser data pe1_=pe1_3d +!$ser data pe2_=pe2_3d +!$ser data ptop=ptop !$ser data qvapor_js=q(:,j,:,sphum) !$ser data qliquid_js=q(:,j,:,liq_wat) -!$ser data qice_js=q(:,j,:,ice_wat) +!$ser data qice_js=q(:,j,:,ice_wat) !$ser data qrain_js=q(:,j,:,rainwat) !$ser data qsnow_js=q(:,j,:,snowwat) !$ser data qgraupel_js=q(:,j,:,graupel) -!$ser data qcld_js=q(:,j,:,cld_amt) -!$ser data gz1d=gz -!$ser data cvm=cvm -!$ser data r_vir=r_vir +!!$ser data qcld_js=q(:,j,:,cld_amt) +!$ser data q_con=q_con +!$ser data pt=pt !$ser data cappa=cappa -!$ser data rrg=rrg !$ser data delp=delp !$ser data delz=delz -!$ser data pt=pt -!$ser data k1k=k1k -!$ser data j_2d=js2d -!$ser data q_con=q_con - +!$ser data ak=ak +!$ser data bk=bk +!$ser data dp2_3d=dp2_3d +!$ser data ps=ps +!$ser data pn2_3d=pn2_3d +!$ser data peln_3d=peln_3d +!$ser data pk2_3d=pk2_3d +!$ser data r_vir=r_vir +!$ser data gz=gz +!$ser data cvm=cvm +!$ser data pk=pk +!$ser data akap=akap +!!$ser data k1k=k1k !!$ser data remap_t=remap_t -!!$ser data ps=ps -!!$ser data peln_3d=peln_3d -!!$ser data pn2_3d=pn2_3d -!!$ser data ak=ak -!!$ser data bk=bk -!!$ser data dp2_3d=dp2_3d -!!$ser data pk=pk -!!$ser data akap=akap !$ser verbatim endif do k=1,km+1 @@ -490,14 +491,15 @@ contains !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) !$ser verbatim enddo !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d - !$ser data pe2_=pe2_3d - !$ser data gz1d=gz - !$ser data cvm=cvm - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data q_con=q_con + !!$ser savepoint Remapping_GEOS-Out + !!$ser data pe1_=pe1_3d + !!$ser data pe2_=pe2_3d + !!$ser data gz1d=gz + !!$ser data cvm=cvm + !!$ser data pt=pt + !!$ser data cappa=cappa + !!$ser data q_con=q_con + !!$ser data dp2_3d=dp2_3d !$ser verbatim endif endif ! hydro test @@ -612,13 +614,6 @@ contains pk2(i,km+1) = pk1(i,km+1) enddo - do k=2,km - do i=is,ie - pn2(i,k) = log(pe2(i,k)) - pk2(i,k) = exp(akap*pn2(i,k)) - enddo - enddo - !!$ser verbatim if(j == js) then !!$ser verbatim do k = 1, km+1 !!$ser verbatim do i = is, ie @@ -643,9 +638,42 @@ contains !!$ser data ps=ps !!$ser data dp2_3d=dp2_3d !!$ser data pn2_3d=pn2_3d -!!$ser data pk=pk2_3d !!$ser verbatim endif + do k=2,km + do i=is,ie + pn2(i,k) = log(pe2(i,k)) + pk2(i,k) = exp(akap*pn2(i,k)) + enddo + enddo + +!$ser verbatim if(j == js) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe1_=pe1_3d pe2_=pe2_3d +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data dp2_3d=dp2_3d +!$ser data pn2_3d=pn2_3d +!$ser data pk2_3d=pk2_3d +!$ser verbatim endif + if (remap_te) then print*,'CODE NOT TESTED HERE 5' !---------------------------------- From 954866cd39fe13c16d298ddbecd9ffdfe379385f Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 26 Dec 2024 14:46:38 -0800 Subject: [PATCH 080/252] Added pn1 variable computation --- pyFV3/stencils/remapping.py | 4 ++++ .../translate/translate_remapping_GEOS.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 7ec28c81..7832b6c9 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -98,6 +98,7 @@ def moist_cv_pt_pressure( bk: FloatFieldK, dp2: FloatField, ps: FloatFieldIJ, + pn1: FloatField, pn2: FloatField, peln: FloatField, remap_t: bool, @@ -161,13 +162,16 @@ def moist_cv_pt_pressure( with computation(PARALLEL): with interval(0, 1): pn2 = peln + pn1 = peln # TODO: refactor the pe2 = ptop assignment from # previous stencil into this one, and remove # pe2 from the other stencil with interval(1, -1): pe2 = ak + bk * ps + pn1 = peln with interval(-1, None): pn2 = peln + pn1 = peln with computation(BACKWARD), interval(0, -1): dp2 = pe2[0, 0, 1] - pe2 diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 06078ee3..5d05dc5e 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -91,6 +91,13 @@ def __init__( "pt": {}, "cappa": {}, "ps": {}, + "pn1_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, "pn2_3d": { "istart": grid.is_, "iend": grid.ie, @@ -174,6 +181,13 @@ def __init__( "jend": grid.je, "kend": grid.npz-1, }, + "pn1_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, "pn2_3d": { "istart": grid.is_, "iend": grid.ie, @@ -379,6 +393,7 @@ def compute_from_storage(self, inputs): inputs["bk"], inputs["dp2_3d"], inputs["ps"], + inputs["pn1_3d"], inputs["pn2_3d"], inputs["peln_3d"], True, From 6ec892b1fea6d3d4f1ff2a20011ce992a999162b Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 26 Dec 2024 14:47:42 -0800 Subject: [PATCH 081/252] Added fv_mapz serialization file --- fv_mapz.F90.SER | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 4fc0c893..56c84e9d 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -413,14 +413,15 @@ contains !$ser data bk=bk !$ser data dp2_3d=dp2_3d !$ser data ps=ps +!$ser data pn1_3d=pn1_3d !$ser data pn2_3d=pn2_3d !$ser data peln_3d=peln_3d !$ser data pk2_3d=pk2_3d !$ser data r_vir=r_vir -!$ser data gz=gz !$ser data cvm=cvm !$ser data pk=pk !$ser data akap=akap +!!$ser data gz=gz !!$ser data k1k=k1k !!$ser data remap_t=remap_t !$ser verbatim endif @@ -652,6 +653,7 @@ contains !$ser verbatim do i = is, ie !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) !$ser verbatim enddo @@ -670,6 +672,7 @@ contains !$ser data delz=delz !$ser data ps=ps !$ser data dp2_3d=dp2_3d +!$ser data pn1_3d=pn1_3d !$ser data pn2_3d=pn2_3d !$ser data pk2_3d=pk2_3d !$ser verbatim endif @@ -721,9 +724,8 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser savepoint Map_Scalar-In - !$ser verbatim mode=1 !$ser verbatim qmin=184.0 - !$ser data pe1_=pn1_3d pe2_=pn2_3d q1=pt mode=mode j_2d=js2d q_min=qmin + !$ser data pe1_=pn1_3d pe2_=pn2_3d q1=pt j_2d=js2d q_min=qmin !$ser verbatim endif call map_scalar(km, pn1, pt, gz, & km, pn2, pt, & From 9737eb84cdf6e5d207c00962f5eab4e33afbb0e2 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 26 Dec 2024 15:33:39 -0800 Subject: [PATCH 082/252] Added map_scalar and verified with TBC C24 data --- fv_mapz.F90.SER | 87 +++++++++++++------ .../translate/translate_remapping_GEOS.py | 30 +++++-- 2 files changed, 80 insertions(+), 37 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 56c84e9d..3ca3b7f2 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -392,6 +392,7 @@ contains !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) !$ser verbatim enddo !$ser verbatim enddo +!$ser verbatim qmin=t_min !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d !$ser data pe1_=pe1_3d @@ -421,6 +422,7 @@ contains !$ser data cvm=cvm !$ser data pk=pk !$ser data akap=akap +!$ser data t_min=qmin !!$ser data gz=gz !!$ser data k1k=k1k !!$ser data remap_t=remap_t @@ -648,34 +650,34 @@ contains enddo enddo -!$ser verbatim if(j == js) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe1_=pe1_3d pe2_=pe2_3d -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data dp2_3d=dp2_3d -!$ser data pn1_3d=pn1_3d -!$ser data pn2_3d=pn2_3d -!$ser data pk2_3d=pk2_3d -!$ser verbatim endif +!!$ser verbatim if(j == js) then +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie +!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km +!!$ser verbatim do i = is, ie +!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe1_=pe1_3d pe2_=pe2_3d +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data dp2_3d=dp2_3d +!!$ser data pn1_3d=pn1_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data pk2_3d=pk2_3d +!!$ser verbatim endif if (remap_te) then print*,'CODE NOT TESTED HERE 5' @@ -734,6 +736,35 @@ contains !$ser savepoint Map_Scalar-Out !$ser data q1=pt !$ser verbatim endif + + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + !$ser verbatim endif endif endif diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 5d05dc5e..b3739c4f 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -147,6 +147,7 @@ def __init__( "r_vir", # "remap_t", # For some reason, translate test can't accept a logical variable "akap", + "t_min", # "zvir", # "last_step", # "consv_te", @@ -222,6 +223,12 @@ def __init__( grid_indexing.domain[2] + 1, ) + # Value from GEOS + self._kord_tm = 9 + + # mode / iv set to 1 from GEOS + self.mode = 1 + # self._pe1 = self.quantity_factory.zeros( # [X_DIM, Y_DIM, Z_INTERFACE_DIM], # units="Pa", @@ -260,13 +267,13 @@ def __init__( domain=(grid.nic, 1, grid.npz-1), ) - # self._map_scalar = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # self._kord_tm, - # self.mode, - # dims=[X_DIM, Y_DIM, Z_DIM], - # ) + self._map_scalar = MapSingle( + self.stencil_factory, + self.quantity_factory, + self._kord_tm, + self.mode, + dims=[X_DIM, Y_DIM, Z_DIM], + ) # self._mapn_tracer = MapNTracer( # self.stencil_factory, @@ -407,8 +414,13 @@ def compute_from_storage(self, inputs): Float(inputs["akap"]), ) - # # now that we have the pressure profiles, we can start remapping - # self._map_scalar(pt, self._pn1, self._pn2, qmin=self._t_min, interp=True) + self._map_scalar( + inputs["pt"], + inputs["pn1_3d"], + inputs["pn2_3d"], + qmin=inputs["t_min"], + interp=True, + ) # self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) From 7e79aafafc34ca9a48e37ae65b35b068c5057545 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 27 Dec 2024 09:58:07 -0800 Subject: [PATCH 083/252] Corrected serialization for MapNTracers --- fv_mapz.F90.SER | 103 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 72 insertions(+), 31 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 3ca3b7f2..20c1131c 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -404,7 +404,9 @@ contains !$ser data qrain_js=q(:,j,:,rainwat) !$ser data qsnow_js=q(:,j,:,snowwat) !$ser data qgraupel_js=q(:,j,:,graupel) -!!$ser data qcld_js=q(:,j,:,cld_amt) +!$ser data qcld_js=q(:,j,:,cld_amt) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) !$ser data q_con=q_con !$ser data pt=pt !$ser data cappa=cappa @@ -737,34 +739,34 @@ contains !$ser data q1=pt !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d - !$ser verbatim endif + !!$ser verbatim if(j == js) then + !!$ser verbatim do k = 1, km+1 + !!$ser verbatim do i = is, ie + !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser verbatim do k = 1, km + !!$ser verbatim do i = is, ie + !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser savepoint Remapping_GEOS-Out + !!$ser data pe1_=pe1_3d pe2_=pe2_3d + !!$ser data q_con=q_con + !!$ser data pt=pt + !!$ser data cappa=cappa + !!$ser data delp=delp + !!$ser data delz=delz + !!$ser data ps=ps + !!$ser data dp2_3d=dp2_3d + !!$ser data pn1_3d=pn1_3d + !!$ser data pn2_3d=pn2_3d + !!$ser data pk2_3d=pk2_3d + !!$ser verbatim endif endif endif @@ -775,8 +777,8 @@ contains !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) !$ser verbatim !$ser verbatim if (k < km+1) dp2_3d(i,j,k) = dp2(i,k) !$ser verbatim enddo @@ -807,6 +809,45 @@ contains !$ser data qo3mr=q(:,:,:,8) !$ser data qsgs_tke=q(:,:,:,9) !$ser verbatim endif + + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + !$ser verbatim endif elseif ( nq > 0 ) then print*,'CODE NOT TESTED HERE 7' ! Remap one tracer at a time From 5363d360cc3d9d666f713e5ac84f9270536f2965 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 27 Dec 2024 10:54:32 -0800 Subject: [PATCH 084/252] Added MapNTracers to remapping translate test. Verified with C24 TBC data --- .../translate/translate_remapping_GEOS.py | 73 +++++++++++++++++-- 1 file changed, 65 insertions(+), 8 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index b3739c4f..529cb69e 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -85,6 +85,13 @@ def __init__( "qrain": {"serialname": "qrain_js"}, "qsnow": {"serialname": "qsnow_js"}, "qgraupel": {"serialname": "qgraupel_js"}, + "qcld": {"serialname": "qcld_js"}, + "qo3mr": { + "kend": grid.npz-1, + }, + "qsgs_tke": { + "kend": grid.npz-1, + }, "delp": {}, "delz": {}, "q_con": {}, @@ -137,7 +144,7 @@ def __init__( }, } # self.write_vars = ["gz", "cvm"] - self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel"] + self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] for k, v in self.in_vars["data_vars"].items(): # if k not in self.write_vars: if k in self.write_vars: @@ -202,7 +209,35 @@ def __init__( "jstart": grid.js, "jend": grid.je, "kend": grid.npz + 1, - } + }, + + "qvapor": { + "kend": grid.npz-1, + }, + "qliquid": { + "kend": grid.npz-1, + }, + "qice": { + "kend": grid.npz-1, + }, + "qrain": { + "kend": grid.npz-1, + }, + "qsnow": { + "kend": grid.npz-1, + }, + "qgraupel": { + "kend": grid.npz-1, + }, + "qcld": { + "kend": grid.npz-1, + }, + "qo3mr": { + "kend": grid.npz-1, + }, + "qsgs_tke": { + "kend": grid.npz-1, + }, } self.stencil_factory = stencil_factory @@ -223,12 +258,19 @@ def __init__( grid_indexing.domain[2] + 1, ) + # Value from GEOS + self.kord = 9 + # Value from GEOS self._kord_tm = 9 # mode / iv set to 1 from GEOS self.mode = 1 + self.nq = 9 + + self.fill = True + # self._pe1 = self.quantity_factory.zeros( # [X_DIM, Y_DIM, Z_INTERFACE_DIM], # units="Pa", @@ -344,11 +386,6 @@ def __init__( # domain=self._domain_kextra, # ) - print("grid.compute_dict() = ", grid.compute_dict()) - print("istard iend jstart jend kend = ", grid.is_, grid.ie, grid.js, grid.je, grid.npz) - print("grid.nic, 1, grid.npz+1 : ", grid.nic, 1, grid.npz+1) - print("grid_indexing.origin_compute() = ", grid_indexing.origin_compute()) - def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction @@ -422,7 +459,27 @@ def compute_from_storage(self, inputs): interp=True, ) - # self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) + tracers = { "qvapor": inputs["qvapor"], + "qliquid": inputs["qliquid"], + "qice": inputs["qice"], + "qrain": inputs["qrain"], + "qsnow": inputs["qsnow"], + "qgraupel": inputs["qgraupel"], + "qcld": inputs["qcld"], + "qo3mr": inputs["qo3mr"], + "qsgs_tke": inputs["qsgs_tke"], + } + + self._mapn_tracer = MapNTracer( + self.stencil_factory, + self.quantity_factory, + abs(self.kord), + self.nq, + fill=self.fill, + tracers=tracers, + ) + + self._mapn_tracer(inputs["pe1_"], inputs["pe2_"], inputs["dp2_3d"], tracers) # self._map_single_w(w, self._pe1, self._pe2, qs=wsd, interp=False) # self._map_single_delz(delz, self._pe1, self._pe2) From 7c22b1e37a3e6c10df54ea1ff4f9f4af455a2dc7 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Sun, 29 Dec 2024 09:14:49 -0800 Subject: [PATCH 085/252] Added map_single_w computation. Verified with C24 TBC example --- .../translate/translate_remapping_GEOS.py | 40 ++++++++++--------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 529cb69e..b20e8b36 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -142,6 +142,15 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, + "w": { + "kend": grid.npz-1, + }, + "ws_":{ + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, } # self.write_vars = ["gz", "cvm"] self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] @@ -155,6 +164,7 @@ def __init__( # "remap_t", # For some reason, translate test can't accept a logical variable "akap", "t_min", + "kord_wz", # "zvir", # "last_step", # "consv_te", @@ -238,6 +248,9 @@ def __init__( "qsgs_tke": { "kend": grid.npz-1, }, + "w": { + "kend": grid.npz-1, + }, } self.stencil_factory = stencil_factory @@ -317,23 +330,6 @@ def __init__( dims=[X_DIM, Y_DIM, Z_DIM], ) - # self._mapn_tracer = MapNTracer( - # self.stencil_factory, - # self.quantity_factory, - # abs(self.kord), - # self.nq, - # fill=self.fill, - # tracers=tracers, - # ) - - # self._map_single_w = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # self._kord_wz, - # -2, - # dims=[X_DIM, Y_DIM, Z_DIM], - # ) - # self._map_single_delz = MapSingle( # self.stencil_factory, # quantity_factory, @@ -481,7 +477,15 @@ def compute_from_storage(self, inputs): self._mapn_tracer(inputs["pe1_"], inputs["pe2_"], inputs["dp2_3d"], tracers) - # self._map_single_w(w, self._pe1, self._pe2, qs=wsd, interp=False) + self._map_single_w = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_wz"], + -2, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._map_single_w(inputs["w"], inputs["pe1_"], inputs["pe2_"], qs=inputs["ws_"], interp=False) # self._map_single_delz(delz, self._pe1, self._pe2) # self._w_fix_consrv_moment( From 630593fa9f9400e95ec4079c5d56fc5a74f1affb Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Sun, 29 Dec 2024 09:16:32 -0800 Subject: [PATCH 086/252] Updated fv_mapz serialization file --- fv_mapz.F90.SER | 130 +++++++++++++++++++++++++++++++----------------- 1 file changed, 85 insertions(+), 45 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 20c1131c..4d96d22c 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -228,7 +228,7 @@ contains real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d - real, dimension(is:ie, js:je) :: ws_2d, gz_2d + real, dimension(is:ie, js:je) :: gz_2d logical :: serial_flag @@ -249,7 +249,6 @@ contains !$ser verbatim pn1_3d = 0.0 !$ser verbatim pn2_3d = 0.0 !$ser verbatim pk2_3d = pk - !$ser verbatim ws_2d = 0.0 !$ser verbatim gz_2d = 0.0 remap_t = .false. @@ -372,7 +371,7 @@ contains !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & !$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,& !$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & -!$OMP peln_3d, ws_2d, gz_2d, pe0_3d, pe3_3d) & +!$OMP peln_3d, gz_2d, pe0_3d, pe3_3d) & !$OMP private(gz,cvm,bkh,dp2, & !$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_) do 1000 j=js,je+1 @@ -425,6 +424,9 @@ contains !$ser data pk=pk !$ser data akap=akap !$ser data t_min=qmin +!$ser data kord_wz=kord_wz +!$ser data ws_=ws +!$ser data w=w !!$ser data gz=gz !!$ser data k1k=k1k !!$ser data remap_t=remap_t @@ -810,44 +812,44 @@ contains !$ser data qsgs_tke=q(:,:,:,9) !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d - - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) - !$ser verbatim endif + !!$ser verbatim if(j == js) then + !!$ser verbatim do k = 1, km+1 + !!$ser verbatim do i = is, ie + !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser verbatim do k = 1, km + !!$ser verbatim do i = is, ie + !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser savepoint Remapping_GEOS-Out + !!$ser data pe1_=pe1_3d pe2_=pe2_3d + !!$ser data q_con=q_con + !!$ser data pt=pt + !!$ser data cappa=cappa + !!$ser data delp=delp + !!$ser data delz=delz + !!$ser data ps=ps + !!$ser data dp2_3d=dp2_3d + !!$ser data pn1_3d=pn1_3d + !!$ser data pn2_3d=pn2_3d + !!$ser data pk2_3d=pk2_3d + + !!$ser data qvapor=q(:,:,:,1) + !!$ser data qliquid=q(:,:,:,2) + !!$ser data qice=q(:,:,:,3) + !!$ser data qrain=q(:,:,:,4) + !!$ser data qsnow=q(:,:,:,5) + !!$ser data qgraupel=q(:,:,:,6) + !!$ser data qcld=q(:,:,:,7) + !!$ser data qo3mr=q(:,:,:,8) + !!$ser data qsgs_tke=q(:,:,:,9) + !!$ser verbatim endif elseif ( nq > 0 ) then print*,'CODE NOT TESTED HERE 7' ! Remap one tracer at a time @@ -873,13 +875,10 @@ contains !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) !$ser verbatim enddo !$ser verbatim enddo - !$ser verbatim do i = is, ie - !$ser verbatim ws_2d(i,j) = ws(i,j) - !$ser verbatim enddo !$ser savepoint Map1_PPM_W-In !$ser data pe1_=pe1_3d !$ser data pe2_=pe2_3d - !$ser data ws_=ws_2d + !$ser data ws_=ws !$ser data w_=w !$ser data kord_wz=kord_wz !$ser verbatim endif @@ -890,6 +889,47 @@ contains !$ser savepoint Map1_PPM_W-Out !$ser data w_=w !$ser verbatim endif + + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + + !$ser data w=w + !$ser verbatim endif ! Remap delz for hybrid sigma-p coordinate do k=1,km do i=is,ie From afa4165c498b8253040629ec00963818e2f52452 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 30 Dec 2024 15:08:46 -0500 Subject: [PATCH 087/252] Update to newer comms --- pyFV3/stencils/tracer_2d_1l.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index 26147e5b..ce81c915 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -27,7 +27,7 @@ Z_DIM, ) from ndsl.dsl.typing import FloatField, FloatFieldIJ, FloatFieldK -from ndsl.typing import Communicator +from ndsl.comm.communicator import Communicator, ReductionOperator from pyFV3.stencils.fvtp2d import FiniteVolumeTransport from pyFV3.tracers import Tracers @@ -549,5 +549,5 @@ def __call__(self, cx: Quantity, cy: Quantity, cmax: Quantity): sin_sg5=self._grid_data.sin_sg5, cmax=self._tmp_cmax, ) - self._tmp_cmax_in_K.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] - self._comm.buffer_all_reduce(self._tmp_cmax_in_K, cmax) + cmax.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] + self._comm.all_reduce_per_element_in_place(cmax, ReductionOperator.MAX) From 1deef1ce1c229ba8ac5e71e0c98b71402a673323 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 7 Jan 2025 11:04:53 -0800 Subject: [PATCH 088/252] Added scaling factor for delz computation before and after map1_ppm calcuation --- fv_mapz.F90.SER | 52 +++++++++++-------- .../translate/translate_map1_ppm_delz.py | 47 ++++++++++++++++- 2 files changed, 77 insertions(+), 22 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 4d96d22c..59425171 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -930,41 +930,51 @@ contains !$ser data w=w !$ser verbatim endif + + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do i = is, ie + !$ser verbatim gz_2d(i,j) = gz(i) + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Map1_PPM_delz-In + !$ser data pe1_=pe1_3d + !$ser data pe2_=pe2_3d + !$ser data dp2_3d=dp2_3d + !$ser data gz_=gz_2d + !$ser data delz_=delz + !$ser data delp=delp + !$ser data kord_wz=kord_wz + !$ser verbatim endif ! Remap delz for hybrid sigma-p coordinate do k=1,km do i=is,ie delz(i,j,k) = -delz(i,j,k) / delp(i,j,k) ! ="specific volume"/grav enddo enddo - !$ser verbatim if(j == js2d) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do i = is, ie - !$ser verbatim gz_2d(i,j) = gz(i) - !$ser verbatim enddo - !$ser savepoint Map1_PPM_delz-In - !$ser data pe1_=pe1_3d - !$ser data pe2_=pe2_3d - !$ser data gz_=gz_2d - !$ser data delz_=delz - !$ser data kord_wz=kord_wz - !$ser verbatim endif + call map1_ppm (km, pe1, delz, gz, & km, pe2, delz, & is, ie, j, isd, ied, jsd, jed, 1, abs(kord_wz)) - !$ser verbatim if(j == js2d) then - !$ser savepoint Map1_PPM_delz-Out - !$ser data delz_=delz - !$ser verbatim endif + do k=1,km do i=is,ie delz(i,j,k) = -delz(i,j,k)*dp2(i,k) enddo enddo + !$ser verbatim if(j == js2d) then + !$ser savepoint Map1_PPM_delz-Out + !$ser data delz_=delz + !$ser verbatim endif !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py index 400ad750..ebee7e1b 100644 --- a/tests/savepoint/translate/translate_map1_ppm_delz.py +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -1,9 +1,24 @@ -from ndsl import StencilFactory, Namelist +from ndsl import StencilFactory, Namelist, QuantityFactory from ndsl.stencils.testing.grid import Grid from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils.map_single import MapSingle +from ndsl.dsl.typing import FloatField +from gt4py.cartesian.gtscript import FORWARD, PARALLEL, computation, interval +def rescale_delz_1( + delz: FloatField, + delp: FloatField, +): + with computation(PARALLEL), interval(...): + delz = -delz / delp + +def rescale_delz_2( + delz: FloatField, + dp: FloatField, +): + with computation(PARALLEL), interval(...): + delz = -delz * dp class TranslateMap1_PPM_delz(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) @@ -29,12 +44,20 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "jend": grid.je, "kend": grid.npz }, + "dp2_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + }, "gz_":{ "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, }, + "delp": {}, } self.in_vars["parameters"] = [ @@ -53,6 +76,18 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.dims=[X_DIM, Y_DIM, Z_DIM] + self._rescale_delz_1 = stencil_factory.from_origin_domain( + rescale_delz_1, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + + self._rescale_delz_2 = stencil_factory.from_origin_domain( + rescale_delz_2, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + def compute_from_storage(self, inputs): self._compute_func = MapSingle( self.stencil_factory, @@ -62,6 +97,11 @@ def compute_from_storage(self, inputs): dims=[X_DIM, Y_DIM, Z_DIM], ) + self._rescale_delz_1( + inputs["delz_"], + inputs["delp"], + ) + self._compute_func( inputs["delz_"], inputs["pe1_"], @@ -69,4 +109,9 @@ def compute_from_storage(self, inputs): qs=inputs["gz_"], interp=False, ) + + self._rescale_delz_2( + inputs["delz_"], + inputs["dp2_3d"], + ) return inputs From 2f186590d26b1bc814cc9b25d938ded24dd8484a Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 7 Jan 2025 14:53:49 -0800 Subject: [PATCH 089/252] Added map1_ppm operation on delz to Remapping translate test. Verifies with C24 TBC example on numpy backend --- fv_mapz.F90.SER | 119 ++++++++++++------ pyFV3/stencils/scale_delz.py | 16 +++ .../translate/translate_remapping_GEOS.py | 54 ++++---- 3 files changed, 127 insertions(+), 62 deletions(-) create mode 100644 pyFV3/stencils/scale_delz.py diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 59425171..b3fae4b9 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -890,46 +890,46 @@ contains !$ser data w_=w !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d - - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) + !!$ser verbatim if(j == js) then + !!$ser verbatim do k = 1, km+1 + !!$ser verbatim do i = is, ie + !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser verbatim do k = 1, km + !!$ser verbatim do i = is, ie + !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser savepoint Remapping_GEOS-Out + !!$ser data pe1_=pe1_3d pe2_=pe2_3d + !!$ser data q_con=q_con + !!$ser data pt=pt + !!$ser data cappa=cappa + !!$ser data delp=delp + !!$ser data delz=delz + !!$ser data ps=ps + !!$ser data dp2_3d=dp2_3d + !!$ser data pn1_3d=pn1_3d + !!$ser data pn2_3d=pn2_3d + !!$ser data pk2_3d=pk2_3d - !$ser data w=w - !$ser verbatim endif + !!$ser data qvapor=q(:,:,:,1) + !!$ser data qliquid=q(:,:,:,2) + !!$ser data qice=q(:,:,:,3) + !!$ser data qrain=q(:,:,:,4) + !!$ser data qsnow=q(:,:,:,5) + !!$ser data qgraupel=q(:,:,:,6) + !!$ser data qcld=q(:,:,:,7) + !!$ser data qo3mr=q(:,:,:,8) + !!$ser data qsgs_tke=q(:,:,:,9) + + !!$ser data w=w + !!$ser verbatim endif !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km+1 @@ -976,6 +976,47 @@ contains !$ser data delz_=delz !$ser verbatim endif + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + + !$ser data w=w + !$ser verbatim endif + !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array if ( w_limiter ) then diff --git a/pyFV3/stencils/scale_delz.py b/pyFV3/stencils/scale_delz.py new file mode 100644 index 00000000..aae04f30 --- /dev/null +++ b/pyFV3/stencils/scale_delz.py @@ -0,0 +1,16 @@ +from ndsl.dsl.typing import FloatField +from gt4py.cartesian.gtscript import FORWARD, PARALLEL, computation, interval + +def rescale_delz_1( + delz: FloatField, + delp: FloatField, +): + with computation(PARALLEL), interval(...): + delz = -delz / delp + +def rescale_delz_2( + delz: FloatField, + dp: FloatField, +): + with computation(PARALLEL), interval(...): + delz = -delz * dp \ No newline at end of file diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index b20e8b36..e4c0771a 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -4,6 +4,7 @@ from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure, pn2_pk_delp from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils import moist_cv +from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py from ndsl.constants import ( @@ -330,13 +331,17 @@ def __init__( dims=[X_DIM, Y_DIM, Z_DIM], ) - # self._map_single_delz = MapSingle( - # self.stencil_factory, - # quantity_factory, - # self._kord_wz, - # 1, - # dims=[X_DIM, Y_DIM, Z_DIM], - # ) + self._rescale_delz_1 = stencil_factory.from_origin_domain( + rescale_delz_1, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + + self._rescale_delz_2 = stencil_factory.from_origin_domain( + rescale_delz_2, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) # self._w_fix_consrv_moment = stencil_factory.from_origin_domain( # func=W_fix_consrv_moment, @@ -400,21 +405,6 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - # self._moist_cv_pt( - # inputs["qvapor"], - # inputs["qliquid"], - # inputs["qrain"], - # inputs["qsnow"], - # inputs["qice"], - # inputs["qgraupel"], - # inputs["q_con"], - # inputs["pt"], - # inputs["cappa"], - # inputs["delp"], - # inputs["delz"], - # inputs["r_vir"], - # ) - self._moist_cv_pt_pressure( inputs["qvapor"], inputs["qliquid"], @@ -485,9 +475,27 @@ def compute_from_storage(self, inputs): dims=[X_DIM, Y_DIM, Z_DIM], ) + self._map_single_delz = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_wz"], + 1, + dims=[X_DIM, Y_DIM, Z_DIM], + ) self._map_single_w(inputs["w"], inputs["pe1_"], inputs["pe2_"], qs=inputs["ws_"], interp=False) - # self._map_single_delz(delz, self._pe1, self._pe2) + + self._rescale_delz_1( + inputs["delz"], + inputs["delp"], + ) + + self._map_single_delz(inputs["delz"], inputs["pe1_"], inputs["pe2_"]) + self._rescale_delz_2( + inputs["delz"], + inputs["dp2_3d"], + ) + # self._w_fix_consrv_moment( # inputs["w"], # inputs["w2"], From 79ffa5104917d0f2a74ece85bb67aacea7044797 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 7 Jan 2025 15:46:13 -0800 Subject: [PATCH 090/252] Created independent stencil file for w_fix_consrv_moment and updated transate test to test w_fix_consrv_moment --- fv_mapz.F90.SER | 89 +++++------ pyFV3/stencils/w_fix_consrv_moment.py | 81 ++++++++++ .../translate_w_fix_consrv_moment.py | 146 ++---------------- 3 files changed, 139 insertions(+), 177 deletions(-) create mode 100644 pyFV3/stencils/w_fix_consrv_moment.py diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index b3fae4b9..f27d7f61 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -976,46 +976,46 @@ contains !$ser data delz_=delz !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d - - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) - - !$ser data w=w - !$ser verbatim endif + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + + !$ser data w=w + !$ser verbatim endif !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array @@ -1030,7 +1030,7 @@ contains !$ser verbatim w_min_ = w_min !$ser savepoint W_fix_consrv_moment-In -!$ser data w=w w2=w2_3d dp2_W=dp2_3d w_max=w_max_ w_min=w_min_ +!$ser data w=w dp2_W=dp2_3d w_max=w_max_ w_min=w_min_ !$ser verbatim endif do k=1,km do i=is,ie @@ -1101,12 +1101,7 @@ contains enddo !$ser verbatim if(j == js2d) then !$ser savepoint W_fix_consrv_moment-Out -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim w2_3d(i,:,k) = w2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser data w=w w2=w2_3d dp2_W=dp2_3d +!$ser data w=w !$ser verbatim endif endif endif diff --git a/pyFV3/stencils/w_fix_consrv_moment.py b/pyFV3/stencils/w_fix_consrv_moment.py new file mode 100644 index 00000000..ed28cb77 --- /dev/null +++ b/pyFV3/stencils/w_fix_consrv_moment.py @@ -0,0 +1,81 @@ +from ndsl.dsl.typing import FloatField, FloatFieldIJ, Float, BoolFieldIJ +from gt4py.cartesian.gtscript import FORWARD, BACKWARD, PARALLEL, computation, interval + +def W_fix_consrv_moment( + w: FloatField, + w2: FloatField, + dp2: FloatField, + gz: FloatFieldIJ, + w_max: Float, + w_min: Float, + compute_performed: BoolFieldIJ, +): + """ + Args: + w (in/out): + w2 (in?): + dp2(in): + w_max(in): + w_min(in): + compute_performed: (Internal Temporary), + """ + + with computation(PARALLEL), interval(...): + w2 = w + + with computation(FORWARD): + with interval(0,1): + compute_performed = False + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + compute_performed = True + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + compute_performed = True + with interval(1,-1): + if(compute_performed): + w2 = w2 + gz / dp2 + compute_performed = False + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + compute_performed = True + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + compute_performed = True + + with computation(BACKWARD): + with interval(-1,None): + compute_performed = False + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + compute_performed = True + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + compute_performed = True + with interval(1,-1): + if(compute_performed): + w2 = w2 + gz / dp2 + compute_performed = False + if(w2 > w_max): + gz = (w2 - w_max) * dp2 + w2 = w_max + compute_performed = True + elif(w2 < w_min): + gz = (w2 - w_min) * dp2 + w2 = w_min + compute_performed = True + + with computation(FORWARD), interval(0,1): + if(w2 > (w_max * 2.0)): + w2 = w_max * 2.0 + elif(w2 < (w_min * 2.0)): + w2 = w_min * 2.0 + + with computation(PARALLEL), interval(...): + w = w2 \ No newline at end of file diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py index 1a07efe6..6a298bc2 100644 --- a/tests/savepoint/translate/translate_w_fix_consrv_moment.py +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -1,142 +1,26 @@ -from gt4py.cartesian.gtscript import PARALLEL, FORWARD, BACKWARD, computation, interval - -from ndsl import StencilFactory, QuantityFactory -from ndsl.dsl.typing import FloatField, FloatFieldIJ, Float, BoolFieldIJ +from ndsl.dsl.typing import Float from ndsl.stencils.testing.grid import Grid from ndsl.stencils.testing import TranslateFortranData2Py - -def W_fix_consrv_moment( - w: FloatField, - w2: FloatField, - dp2: FloatField, - gz: FloatFieldIJ, - w_max: Float, - w_min: Float, - compute_performed: BoolFieldIJ, -): - """ - Args: - w (in/out): - w2 (in?): - dp2(in): - w_max(in): - w_min(in): - compute_performed: (Internal Temporary), - """ - - with computation(PARALLEL), interval(...): - w2 = w - - with computation(FORWARD): - with interval(0,1): - compute_performed = False - if(w2 > w_max): - gz = (w2 - w_max) * dp2 - w2 = w_max - compute_performed = True - elif(w2 < w_min): - gz = (w2 - w_min) * dp2 - w2 = w_min - compute_performed = True - with interval(1,-1): - if(compute_performed): - w2 = w2 + gz / dp2 - compute_performed = False - if(w2 > w_max): - gz = (w2 - w_max) * dp2 - w2 = w_max - compute_performed = True - elif(w2 < w_min): - gz = (w2 - w_min) * dp2 - w2 = w_min - compute_performed = True - - with computation(BACKWARD): - with interval(-1,None): - compute_performed = False - if(w2 > w_max): - gz = (w2 - w_max) * dp2 - w2 = w_max - compute_performed = True - elif(w2 < w_min): - gz = (w2 - w_min) * dp2 - w2 = w_min - compute_performed = True - with interval(1,-1): - if(compute_performed): - w2 = w2 + gz / dp2 - compute_performed = False - if(w2 > w_max): - gz = (w2 - w_max) * dp2 - w2 = w_max - compute_performed = True - elif(w2 < w_min): - gz = (w2 - w_min) * dp2 - w2 = w_min - compute_performed = True - - with computation(FORWARD), interval(0,1): - if(w2 > (w_max * 2.0)): - w2 = w_max * 2.0 - elif(w2 < (w_min * 2.0)): - w2 = w_min * 2.0 - - with computation(PARALLEL), interval(...): - w = w2 -class testClass: - - def __init__( - self, - stencil_factory: StencilFactory, - grid: Grid, - ): - print("In testClass") - print("Grid.nid Grid.njd Grid.npz", grid.nid, grid.njd, grid.npz) - - self._w_fix_consrv_moment = stencil_factory.from_origin_domain( - func=W_fix_consrv_moment, - origin=(3,3,0), - domain=(24,1,72), - ) - - def __call__( - self, - w: FloatField, - w2: FloatField, - dp2: FloatField, - gz: FloatFieldIJ, - w_max: Float, - w_min: Float, - compute_performed_bool: BoolFieldIJ - ): - self._w_fix_consrv_moment( - w, - w2, - dp2, - gz, - w_max, - w_min, - compute_performed_bool, - ) - +from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment class TranslateW_fix_consrv_moment(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist, stencil_factory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid - self.compute_func = testClass(self.stencil_factory, self.grid) # type: ignore self.quantity_factory = grid.quantity_factory - print("Running W_fix_consrv_moment translate test") + self.compute_func = stencil_factory.from_origin_domain( + func=W_fix_consrv_moment, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) self.in_vars["data_vars"] = { "w": { "kend": grid.npz-1, }, - "w2": grid.compute_dict(), "dp2_W": grid.compute_dict(), - } self.in_vars["parameters"] = ["w_max", "w_min"] @@ -145,8 +29,6 @@ def __init__(self, grid: Grid, namelist, stencil_factory): "w": { "kend": grid.npz-1, }, - "w2": grid.compute_dict(), - "dp2_W": grid.compute_dict(), } self._gz = self.quantity_factory._numpy.zeros( ( @@ -155,6 +37,14 @@ def __init__(self, grid: Grid, namelist, stencil_factory): ), dtype=Float, ) + self._w2 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz, + ), dtype=Float, + ) + self._compute_performed = self.quantity_factory._numpy.zeros( ( grid.nid, @@ -166,7 +56,7 @@ def compute_from_storage(self, inputs): self.compute_func( inputs["w"], - inputs["w2"], + self._w2, inputs["dp2_W"], self._gz, inputs["w_max"], @@ -174,8 +64,4 @@ def compute_from_storage(self, inputs): self._compute_performed ) - for i in range(self.grid.is_,self.grid.ie+1): - for k in range(self.grid.npz): - temp = inputs["w2"][i,3,k] - inputs["w2"][i,:,k] = temp return inputs From 955e327644252ac67b7f6e522d18c2efb920edf8 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 8 Jan 2025 13:46:51 -0800 Subject: [PATCH 091/252] Added w_fix to Remapping translate test and corrected Pressures_mapU translate test by adding stencil to compute at ie+1 --- pyFV3/stencils/remapping.py | 4 + .../translate/translate_Pressures_mapU.py | 147 ++++++++++-------- .../translate/translate_remapping_GEOS.py | 62 +++++--- 3 files changed, 121 insertions(+), 92 deletions(-) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 7832b6c9..b0d2ce4b 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -204,6 +204,10 @@ def pn2_pk_delp( pn2 = log(pe2) pk = exp(akap * pn2) +def pe0_ptop_xmax(pe0: FloatField, + ptop: Float): + with computation(PARALLEL), interval(0,1): + pe0 = ptop def pressures_mapu( pe: FloatField, diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index f944959d..441d4454 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -2,7 +2,7 @@ from ndsl.stencils.testing.grid import Grid from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Y_INTERFACE_DIM from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils.remapping import pressures_mapu +from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax from pyFV3.stencils.map_single import MapSingle class TranslatePressures_mapU(TranslateFortranData2Py): @@ -20,26 +20,27 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "jend": grid.je+1, "kend": grid.npz }, + "ak":{ + + }, + "bk":{ + + }, "pe0_": { "istart": grid.is_, - "iend": grid.ie, + "iend": grid.ie+1, "jstart": grid.js, - "jend": grid.je, + "jend": grid.je+1, "kend": grid.npz }, "pe3_": { "istart": grid.is_, - "iend": grid.ie, + "iend": grid.ie+1, "jstart": grid.js, - "jend": grid.je, + "jend": grid.je+1, "kend": grid.npz }, - "ak":{ - - }, - "bk":{ - }, "u_": { "istart": grid.isd, "iend": grid.ied, @@ -48,44 +49,44 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "kend": grid.npz-1, }, - "mfy_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, - }, + # "mfy_": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, - "cy_": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, - }, + # "cy_": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, } self.in_vars["parameters"] = [ - "kord_mt", "ptop", + "kord_mt", ] self.out_vars = { - "pe0_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz - }, - "pe3_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz - }, + # "pe0_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz + # }, + # "pe3_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz + # }, "u_": { "istart": grid.isd, "iend": grid.ied, @@ -94,21 +95,21 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "kend": grid.npz-1, }, - "mfy_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, - }, - - "cy_": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, - }, + # "mfy_": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, + + # "cy_": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, } grid_indexing = stencil_factory.grid_indexing @@ -117,14 +118,20 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self._domain_jextra = ( grid_indexing.domain[0], - grid_indexing.domain[1] + 1, + 1, grid_indexing.domain[2] + 1, ) self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=self._domain_jextra, + domain=(grid_indexing.domain[1],1,grid_indexing.domain[2] + 1) + ) + + self._pe0_ptop_xmax = stencil_factory.from_origin_domain( + pe0_ptop_xmax, + origin=(grid_indexing.domain[0]+3,3,0), + domain=(1,1,grid_indexing.domain[2] + 1) ) def compute_from_storage(self, inputs): @@ -145,24 +152,28 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - self._map1_ppm_u( - inputs["u_"], + self._pe0_ptop_xmax( inputs["pe0_"], - inputs["pe3_"], - interp=False, + inputs["ptop"], ) - + self._map1_ppm_u( - inputs["mfy_"], + inputs["u_"], inputs["pe0_"], inputs["pe3_"], interp=False, ) + # self._map1_ppm_u( + # inputs["mfy_"], + # inputs["pe0_"], + # inputs["pe3_"], + # interp=False, + # ) - self._map1_ppm_u( - inputs["cy_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + # self._map1_ppm_u( + # inputs["cy_"], + # inputs["pe0_"], + # inputs["pe3_"], + # interp=False, + # ) return inputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index e4c0771a..ae664d99 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -5,6 +5,7 @@ from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils import moist_cv from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 +from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py from ndsl.constants import ( @@ -166,6 +167,8 @@ def __init__( "akap", "t_min", "kord_wz", + "w_max", + "w_min", # "zvir", # "last_step", # "consv_te", @@ -285,16 +288,27 @@ def __init__( self.fill = True - # self._pe1 = self.quantity_factory.zeros( - # [X_DIM, Y_DIM, Z_INTERFACE_DIM], - # units="Pa", - # dtype=Float, - # ) - # self._pe2 = self.quantity_factory.zeros( - # [X_DIM, Y_DIM, Z_INTERFACE_DIM], - # units="Pa", - # dtype=Float, - # ) + self._gz = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + ), dtype=Float, + ) + + self._w2 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz, + ), dtype=Float, + ) + + self._compute_performed = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + ), dtype=bool, + ) self._init_pe = stencil_factory.from_origin_domain( init_pe, @@ -343,11 +357,11 @@ def __init__( domain=(grid.nic, 1, grid.npz), ) - # self._w_fix_consrv_moment = stencil_factory.from_origin_domain( - # func=W_fix_consrv_moment, - # origin=(3,3,0), - # domain=(24,1,72), - # ) + self._w_fix_consrv_moment = stencil_factory.from_origin_domain( + func=W_fix_consrv_moment, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) # self._pressures_mapu = stencil_factory.from_origin_domain( # pressures_mapu, @@ -496,15 +510,15 @@ def compute_from_storage(self, inputs): inputs["dp2_3d"], ) - # self._w_fix_consrv_moment( - # inputs["w"], - # inputs["w2"], - # inputs["dp2_W"], - # self._gz, - # inputs["w_max"], - # inputs["w_min"], - # self._compute_performed - # ) + self._w_fix_consrv_moment( + inputs["w"], + self._w2, + inputs["dp2_3d"], + self._gz, + inputs["w_max"], + inputs["w_min"], + self._compute_performed + ) # self._pressures_mapu( # inputs["pe_"], From 517da152e628726634d5083f777af4944102d3fd Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 8 Jan 2025 16:47:01 -0800 Subject: [PATCH 092/252] Added map1_ppm computation with mfy and cy. Also added fv_mapz serialization file --- fv_mapz.F90.SER | 154 ++++++++++-------- .../translate/translate_Pressures_mapU.py | 112 ++++++------- 2 files changed, 143 insertions(+), 123 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index f27d7f61..e384cacd 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -392,6 +392,8 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser verbatim qmin=t_min +!$ser verbatim w_max_ = w_max +!$ser verbatim w_min_ = w_min !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d !$ser data pe1_=pe1_3d @@ -427,6 +429,8 @@ contains !$ser data kord_wz=kord_wz !$ser data ws_=ws !$ser data w=w +!$ser data w_max=w_max_ +!$ser data w_min=w_min_ !!$ser data gz=gz !!$ser data k1k=k1k !!$ser data remap_t=remap_t @@ -976,46 +980,46 @@ contains !$ser data delz_=delz !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d + !!$ser verbatim if(j == js) then + !!$ser verbatim do k = 1, km+1 + !!$ser verbatim do i = is, ie + !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser verbatim do k = 1, km + !!$ser verbatim do i = is, ie + !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser savepoint Remapping_GEOS-Out + !!$ser data pe1_=pe1_3d pe2_=pe2_3d + !!$ser data q_con=q_con + !!$ser data pt=pt + !!$ser data cappa=cappa + !!$ser data delp=delp + !!$ser data delz=delz + !!$ser data ps=ps + !!$ser data dp2_3d=dp2_3d + !!$ser data pn1_3d=pn1_3d + !!$ser data pn2_3d=pn2_3d + !!$ser data pk2_3d=pk2_3d - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) + !!$ser data qvapor=q(:,:,:,1) + !!$ser data qliquid=q(:,:,:,2) + !!$ser data qice=q(:,:,:,3) + !!$ser data qrain=q(:,:,:,4) + !!$ser data qsnow=q(:,:,:,5) + !!$ser data qgraupel=q(:,:,:,6) + !!$ser data qcld=q(:,:,:,7) + !!$ser data qo3mr=q(:,:,:,8) + !!$ser data qsgs_tke=q(:,:,:,9) - !$ser data w=w - !$ser verbatim endif + !!$ser data w=w + !!$ser verbatim endif !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array @@ -1103,15 +1107,53 @@ contains !$ser savepoint W_fix_consrv_moment-Out !$ser data w=w !$ser verbatim endif + + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + + !$ser data w=w + !$ser verbatim endif endif endif endif !(j < je+1) !$ser verbatim if(j == js2d) then -!$ser verbatim pe_3d = 0 -!$ser verbatim pe0_3d = 0 -!$ser verbatim pe3_3d = 0 !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is-1, ie+1 !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) @@ -1129,8 +1171,8 @@ contains !$ser verbatim enddo !$ser savepoint Pressures_mapU-In !$ser data pe_=pe_3d -!$ser data pe0_=pe0_3d(is:ie,js:je,:) -!$ser data pe3_=pe3_3d(is:ie,js:je,:) +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d !$ser data ak=ak !$ser data bk=bk !$ser data ptop=ptop @@ -1163,28 +1205,6 @@ contains enddo enddo -!!$ser verbatim if(j == js2d) then -!!$ser verbatim pe0_3d = 0.0 -!!$ser verbatim pe3_3d = 0.0 -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie+1 -!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Pressures_mapU-In -!!$ser data pe0_=pe0_3d(is:ie,:,:) -!!$ser data pe3_=pe3_3d(is:ie,:,:) -!!$ser data u_=u -!!$ser data kord_mt=kord_mt -!!$ser verbatim if (present(mfy)) then -!!$ser data mfy_=mfy -!!$ser verbatim endif -!!$ser verbatim if (present(cy)) then -!!$ser data cy_=cy -!!$ser verbatim endif -!!$ser verbatim endif - call map1_ppm( km, pe0(is:ie,:), u, gz, & km, pe3(is:ie,:), u, & is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) @@ -1209,8 +1229,8 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser savepoint Pressures_mapU-Out -!$ser data pe0_=pe0_3d(is:ie,js:je,:) -!$ser data pe3_=pe3_3d(is:ie,js:je,:) +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d !$ser data u_=u !$ser verbatim if (present(mfy)) then !$ser data mfy_=mfy diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index 441d4454..c11e3ae1 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -49,21 +49,21 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "kend": grid.npz-1, }, - # "mfy_": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, - - # "cy_": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, + "mfy_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + + "cy_": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, } @@ -73,20 +73,20 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto ] self.out_vars = { - # "pe0_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz - # }, - # "pe3_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz - # }, + "pe0_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz + }, + "pe3_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz + }, "u_": { "istart": grid.isd, "iend": grid.ied, @@ -95,21 +95,21 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "kend": grid.npz-1, }, - # "mfy_": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, - - # "cy_": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, + "mfy_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + + "cy_": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, } grid_indexing = stencil_factory.grid_indexing @@ -163,17 +163,17 @@ def compute_from_storage(self, inputs): inputs["pe3_"], interp=False, ) - # self._map1_ppm_u( - # inputs["mfy_"], - # inputs["pe0_"], - # inputs["pe3_"], - # interp=False, - # ) + self._map1_ppm_u( + inputs["mfy_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) - # self._map1_ppm_u( - # inputs["cy_"], - # inputs["pe0_"], - # inputs["pe3_"], - # interp=False, - # ) + self._map1_ppm_u( + inputs["cy_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) return inputs From e3fd2ef0fae55f3dba456f27fc43f17b4b977cec Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 8 Jan 2025 17:53:11 -0800 Subject: [PATCH 093/252] Addressed small change in Pressures_mapU translate test and incorporated pressure changes before mapU in Remapping translate test --- .../translate/translate_Pressures_mapU.py | 8 +- .../translate/translate_remapping_GEOS.py | 132 +++++++++++++++--- 2 files changed, 110 insertions(+), 30 deletions(-) diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index c11e3ae1..6ea80092 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -116,16 +116,10 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.dims=[X_DIM, Y_DIM, Z_DIM] - self._domain_jextra = ( - grid_indexing.domain[0], - 1, - grid_indexing.domain[2] + 1, - ) - self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[1],1,grid_indexing.domain[2] + 1) + domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1) ) self._pe0_ptop_xmax = stencil_factory.from_origin_domain( diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index ae664d99..35c78430 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -6,6 +6,7 @@ from pyFV3.stencils import moist_cv from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment +from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py from ndsl.constants import ( @@ -81,6 +82,20 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, + "pe0_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz + }, + "pe3_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz + }, "qvapor": {"serialname": "qvapor_js"}, "qliquid": {"serialname": "qliquid_js"}, "qice": {"serialname": "qice_js"}, @@ -153,6 +168,28 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, + "u": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed+1, + "kend": grid.npz-1, + }, + "mfy": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + + "cy": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, } # self.write_vars = ["gz", "cvm"] self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] @@ -169,6 +206,7 @@ def __init__( "kord_wz", "w_max", "w_min", + "kord_mt", # "zvir", # "last_step", # "consv_te", @@ -190,6 +228,20 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, + "pe0_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz + }, + "pe3_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz + }, "pt": {}, "cappa": {}, "q_con": {}, @@ -255,6 +307,29 @@ def __init__( "w": { "kend": grid.npz-1, }, + # "u": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed+1, + # "kend": grid.npz-1, + # }, + + # "mfy": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, + + # "cy": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, } self.stencil_factory = stencil_factory @@ -363,19 +438,17 @@ def __init__( domain=(grid.nic, 1, grid.npz), ) - # self._pressures_mapu = stencil_factory.from_origin_domain( - # pressures_mapu, - # origin=grid_indexing.origin_compute(), - # domain=self._domain_jextra, - # ) + self._pressures_mapu = stencil_factory.from_origin_domain( + pressures_mapu, + origin=grid_indexing.origin_compute(), + domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1) + ) - # self._map1_ppm_u = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # inputs["kord_mt"], - # -1, - # dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], - # ) + self._pe0_ptop_xmax = stencil_factory.from_origin_domain( + pe0_ptop_xmax, + origin=(grid_indexing.domain[0]+3,3,0), + domain=(1,1,grid_indexing.domain[2] + 1) + ) # self._pressures_mapv = stencil_factory.from_origin_domain( # pressures_mapv, @@ -520,31 +593,44 @@ def compute_from_storage(self, inputs): self._compute_performed ) - # self._pressures_mapu( - # inputs["pe_"], - # inputs["ak"], - # inputs["bk"], - # inputs["pe0_"], - # inputs["pe3_"], - # inputs["ptop"], - # ) + self._map1_ppm_u = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_mt"], + -1, + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + ) + + self._pressures_mapu( + inputs["pe_"], + inputs["ak"], + inputs["bk"], + inputs["pe0_"], + inputs["pe3_"], + inputs["ptop"], + ) + + self._pe0_ptop_xmax( + inputs["pe0_"], + inputs["ptop"], + ) # self._map1_ppm_u( - # inputs["u_"], + # inputs["u"], # inputs["pe0_"], # inputs["pe3_"], # interp=False, # ) # self._map1_ppm_u( - # inputs["mfy_"], + # inputs["mfy"], # inputs["pe0_"], # inputs["pe3_"], # interp=False, # ) # self._map1_ppm_u( - # inputs["cy_"], + # inputs["cy"], # inputs["pe0_"], # inputs["pe3_"], # interp=False, From d9b7abd88639caffdf021eaf11e869c09754b2ba Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 8 Jan 2025 17:53:43 -0800 Subject: [PATCH 094/252] Added fv_mapz.F90 serialization file --- fv_mapz.F90.SER | 203 ++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 162 insertions(+), 41 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index e384cacd..7a116cb5 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -382,6 +382,7 @@ contains !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim pe_3d(i,j-1,k) = pe(i,k,j-1) !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) !$ser verbatim enddo @@ -391,6 +392,12 @@ contains !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) !$ser verbatim enddo !$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo !$ser verbatim qmin=t_min !$ser verbatim w_max_ = w_max !$ser verbatim w_min_ = w_min @@ -398,6 +405,8 @@ contains !$ser data pe_=pe_3d !$ser data pe1_=pe1_3d !$ser data pe2_=pe2_3d +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d !$ser data ptop=ptop !$ser data qvapor_js=q(:,j,:,sphum) !$ser data qliquid_js=q(:,j,:,liq_wat) @@ -429,8 +438,16 @@ contains !$ser data kord_wz=kord_wz !$ser data ws_=ws !$ser data w=w +!$ser data u=u +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif !$ser data w_max=w_max_ !$ser data w_min=w_min_ +!$ser data kord_mt=kord_mt !!$ser data gz=gz !!$ser data k1k=k1k !!$ser data remap_t=remap_t @@ -1108,46 +1125,46 @@ contains !$ser data w=w !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d + !!$ser verbatim if(j == js) then + !!$ser verbatim do k = 1, km+1 + !!$ser verbatim do i = is, ie + !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser verbatim do k = 1, km + !!$ser verbatim do i = is, ie + !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !!$ser verbatim enddo + !!$ser verbatim enddo + !!$ser savepoint Remapping_GEOS-Out + !!$ser data pe1_=pe1_3d pe2_=pe2_3d + !!$ser data q_con=q_con + !!$ser data pt=pt + !!$ser data cappa=cappa + !!$ser data delp=delp + !!$ser data delz=delz + !!$ser data ps=ps + !!$ser data dp2_3d=dp2_3d + !!$ser data pn1_3d=pn1_3d + !!$ser data pn2_3d=pn2_3d + !!$ser data pk2_3d=pk2_3d - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) + !!$ser data qvapor=q(:,:,:,1) + !!$ser data qliquid=q(:,:,:,2) + !!$ser data qice=q(:,:,:,3) + !!$ser data qrain=q(:,:,:,4) + !!$ser data qsnow=q(:,:,:,5) + !!$ser data qgraupel=q(:,:,:,6) + !!$ser data qcld=q(:,:,:,7) + !!$ser data qo3mr=q(:,:,:,8) + !!$ser data qsgs_tke=q(:,:,:,9) - !$ser data w=w - !$ser verbatim endif + !!$ser data w=w + !!$ser verbatim endif endif endif @@ -1166,9 +1183,6 @@ contains !$ser verbatim pe3_3d(i,j,k) = pe3(i,k) !$ser verbatim enddo !$ser verbatim enddo -!$ser verbatim do i = is, ie -!$ser verbatim gz_2d(i,j) = gz(i) -!$ser verbatim enddo !$ser savepoint Pressures_mapU-In !$ser data pe_=pe_3d !$ser data pe0_=pe0_3d @@ -1205,6 +1219,56 @@ contains enddo enddo +!$ser verbatim if(j == js) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe1_=pe1_3d pe2_=pe2_3d +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data dp2_3d=dp2_3d +!$ser data pn1_3d=pn1_3d +!$ser data pn2_3d=pn2_3d +!$ser data pk2_3d=pk2_3d + +!$ser data qvapor=q(:,:,:,1) +!$ser data qliquid=q(:,:,:,2) +!$ser data qice=q(:,:,:,3) +!$ser data qrain=q(:,:,:,4) +!$ser data qsnow=q(:,:,:,5) +!$ser data qgraupel=q(:,:,:,6) +!$ser data qcld=q(:,:,:,7) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) + +!$ser data w=w + +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser verbatim endif + call map1_ppm( km, pe0(is:ie,:), u, gz, & km, pe3(is:ie,:), u, & is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) @@ -1240,6 +1304,63 @@ contains !$ser verbatim endif !$ser verbatim endif +!!$ser verbatim if(j == js) then +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie +!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km +!!$ser verbatim do i = is, ie +!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie+1 +!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe1_=pe1_3d pe2_=pe2_3d +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data dp2_3d=dp2_3d +!!$ser data pn1_3d=pn1_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data pk2_3d=pk2_3d + +!!$ser data qvapor=q(:,:,:,1) +!!$ser data qliquid=q(:,:,:,2) +!!$ser data qice=q(:,:,:,3) +!!$ser data qrain=q(:,:,:,4) +!!$ser data qsnow=q(:,:,:,5) +!!$ser data qgraupel=q(:,:,:,6) +!!$ser data qcld=q(:,:,:,7) +!!$ser data qo3mr=q(:,:,:,8) +!!$ser data qsgs_tke=q(:,:,:,9) + +!!$ser data w=w + +!!$ser data pe0_=pe0_3d +!!$ser data pe3_=pe3_3d +!!$ser data u=u +!!$ser verbatim if (present(mfy)) then +!!$ser data mfy=mfy +!!$ser verbatim endif +!!$ser verbatim if (present(cy)) then +!!$ser data cy=cy +!!$ser verbatim endif +!!$ser verbatim endif + ! Note : This serialization portion will test the update of pe0 ! and pe3 !$ser verbatim if(j == js2d) then From ab2a3122ff75ab7ce4ea2e73db0bb09291767da3 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 9 Jan 2025 08:05:01 -0800 Subject: [PATCH 095/252] Added mfy and cy remapping calculation --- .../translate/translate_remapping_GEOS.py | 82 +++++++++---------- 1 file changed, 41 insertions(+), 41 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 35c78430..1a828fa8 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -307,29 +307,29 @@ def __init__( "w": { "kend": grid.npz-1, }, - # "u": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.jsd, - # "jend": grid.jed+1, - # "kend": grid.npz-1, - # }, - - # "mfy": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, - - # "cy": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, + "u": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed+1, + "kend": grid.npz-1, + }, + + "mfy": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + + "cy": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, } self.stencil_factory = stencil_factory @@ -615,26 +615,26 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - # self._map1_ppm_u( - # inputs["u"], - # inputs["pe0_"], - # inputs["pe3_"], - # interp=False, - # ) + self._map1_ppm_u( + inputs["u"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) - # self._map1_ppm_u( - # inputs["mfy"], - # inputs["pe0_"], - # inputs["pe3_"], - # interp=False, - # ) + self._map1_ppm_u( + inputs["mfy"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) - # self._map1_ppm_u( - # inputs["cy"], - # inputs["pe0_"], - # inputs["pe3_"], - # interp=False, - # ) + self._map1_ppm_u( + inputs["cy"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) # self._pressures_mapv( # inputs["pe_"], From 6ed3602cc5723785b3c16128042c41cb9feb2e4d Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 9 Jan 2025 08:48:47 -0800 Subject: [PATCH 096/252] Editing and verifying translate test for V remapping --- fv_mapz.F90.SER | 220 ++++++++---------- .../translate/translate_Pressures_mapV.py | 166 ++++++------- 2 files changed, 184 insertions(+), 202 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 7a116cb5..297077a9 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -1219,6 +1219,91 @@ contains enddo enddo +!!$ser verbatim if(j == js) then +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie +!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km +!!$ser verbatim do i = is, ie +!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie+1 +!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe1_=pe1_3d pe2_=pe2_3d +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data dp2_3d=dp2_3d +!!$ser data pn1_3d=pn1_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data pk2_3d=pk2_3d + +!!$ser data qvapor=q(:,:,:,1) +!!$ser data qliquid=q(:,:,:,2) +!!$ser data qice=q(:,:,:,3) +!!$ser data qrain=q(:,:,:,4) +!!$ser data qsnow=q(:,:,:,5) +!!$ser data qgraupel=q(:,:,:,6) +!!$ser data qcld=q(:,:,:,7) +!!$ser data qo3mr=q(:,:,:,8) +!!$ser data qsgs_tke=q(:,:,:,9) + +!!$ser data w=w + +!!$ser data pe0_=pe0_3d +!!$ser data pe3_=pe3_3d +!!$ser verbatim endif + + call map1_ppm( km, pe0(is:ie,:), u, gz, & + km, pe3(is:ie,:), u, & + is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) + if (present(mfy)) then + ! print*, 'mfy present' + call map1_ppm( km, pe0(is:ie,:), mfy, gz, & + km, pe3(is:ie,:), mfy, & + is, ie, j, is, ie, js, je+1, -1, kord_mt) + endif + if (present(cy)) then + ! print*, 'cy present' + call map1_ppm( km, pe0(is:ie,:), cy, gz, & + km, pe3(is:ie,:), cy, & + is, ie, j, isd, ied, js, je+1, -1, kord_mt) + endif + +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Pressures_mapU-Out +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser data u_=u +!$ser verbatim if (present(mfy)) then +!$ser data mfy_=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy_=cy +!$ser verbatim endif +!$ser verbatim endif + !$ser verbatim if(j == js) then !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie @@ -1267,106 +1352,18 @@ contains !$ser data pe0_=pe0_3d !$ser data pe3_=pe3_3d -!$ser verbatim endif - - call map1_ppm( km, pe0(is:ie,:), u, gz, & - km, pe3(is:ie,:), u, & - is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) - if (present(mfy)) then - ! print*, 'mfy present' - call map1_ppm( km, pe0(is:ie,:), mfy, gz, & - km, pe3(is:ie,:), mfy, & - is, ie, j, is, ie, js, je+1, -1, kord_mt) - endif - if (present(cy)) then - ! print*, 'cy present' - call map1_ppm( km, pe0(is:ie,:), cy, gz, & - km, pe3(is:ie,:), cy, & - is, ie, j, isd, ied, js, je+1, -1, kord_mt) - endif - -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Pressures_mapU-Out -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data u_=u +!$ser data u=u !$ser verbatim if (present(mfy)) then -!$ser data mfy_=mfy +!$ser data mfy=mfy !$ser verbatim endif !$ser verbatim if (present(cy)) then -!$ser data cy_=cy +!$ser data cy=cy !$ser verbatim endif !$ser verbatim endif -!!$ser verbatim if(j == js) then -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie -!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km -!!$ser verbatim do i = is, ie -!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie+1 -!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe1_=pe1_3d pe2_=pe2_3d -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data dp2_3d=dp2_3d -!!$ser data pn1_3d=pn1_3d -!!$ser data pn2_3d=pn2_3d -!!$ser data pk2_3d=pk2_3d - -!!$ser data qvapor=q(:,:,:,1) -!!$ser data qliquid=q(:,:,:,2) -!!$ser data qice=q(:,:,:,3) -!!$ser data qrain=q(:,:,:,4) -!!$ser data qsnow=q(:,:,:,5) -!!$ser data qgraupel=q(:,:,:,6) -!!$ser data qcld=q(:,:,:,7) -!!$ser data qo3mr=q(:,:,:,8) -!!$ser data qsgs_tke=q(:,:,:,9) - -!!$ser data w=w - -!!$ser data pe0_=pe0_3d -!!$ser data pe3_=pe3_3d -!!$ser data u=u -!!$ser verbatim if (present(mfy)) then -!!$ser data mfy=mfy -!!$ser verbatim endif -!!$ser verbatim if (present(cy)) then -!!$ser data cy=cy -!!$ser verbatim endif -!!$ser verbatim endif - ! Note : This serialization portion will test the update of pe0 ! and pe3 !$ser verbatim if(j == js2d) then -!$ser verbatim pe_3d = 0 -!$ser verbatim pe0_3d = 0 -!$ser verbatim pe3_3d = 0 !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie+1 !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) @@ -1381,10 +1378,18 @@ contains !$ser verbatim enddo !$ser savepoint Pressures_mapV-In !$ser data pe_=pe_3d -!$ser data pe0_v=pe0_3d(:,js:je,:) -!$ser data pe3_v=pe3_3d(:,js:je,:) +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d !$ser data ak=ak -!$ser data bk=bk +!$ser data bk=bk +!$ser data v_=v +!$ser data kord_mt=kord_mt +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif !$ser verbatim endif !------ ! map v @@ -1403,29 +1408,6 @@ contains enddo enddo -! Note : This serialization portion will test the update of v, -! mfx, and cx -!!$ser verbatim if(j == js2d) then -!!$ser verbatim pe0_3d = 0.0 -!!$ser verbatim pe3_3d = 0.0 -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie+1 -!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Pressures_mapV-In -!!$ser data pe0_v=pe0_3d(:,js:je,:) -!!$ser data pe3_v=pe3_3d(:,js:je,:) -!!$ser data v_=v -!!$ser data kord_mt=kord_mt -!!$ser verbatim if (present(mfx)) then -!!$ser data mfx_=mfx -!!$ser verbatim endif -!!$ser verbatim if (present(cx)) then -!!$ser data cx_=cx -!!$ser verbatim endif -!!$ser verbatim endif call map1_ppm (km, pe0, v, gz, & km, pe3, v, is, ie+1, & j, isd, ied+1, jsd, jed, -1, kord_mt) @@ -1451,8 +1433,8 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser savepoint Pressures_mapV-Out -!$ser data pe0_v=pe0_3d(:,js:je,:) -!$ser data pe3_v=pe3_3d(:,js:je,:) +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d !$ser data v_=v !$ser verbatim if (present(mfx)) then !$ser data mfx_=mfx diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index f0904a9b..d233f669 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -20,18 +20,18 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "jend": grid.je+1, "kend": grid.npz+1 }, - "pe0_v": { + "pe0_": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, - "jend": grid.je, + "jend": grid.je+1, "kend": grid.npz+1 }, - "pe3_v": { + "pe3_": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, - "jend": grid.je, + "jend": grid.je+1, "kend": grid.npz+1 }, "ak":{ @@ -40,74 +40,74 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "bk":{ }, - # "v_": { - # "istart": grid.isd, - # "iend": grid.ied+1, - # "jstart": grid.jsd, - # "jend": grid.jed, - # "kend": grid.npz-1, - # }, - - # "mfx_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, - - # "cx_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.jsd, - # "jend": grid.jed, - # "kend": grid.npz-1, - # }, + "v_": { + "istart": grid.isd, + "iend": grid.ied+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, + + "mfx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + + "cx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, } self.in_vars["parameters"] = [ - # "kord_mt", + "kord_mt", ] self.out_vars = { - "pe0_v": { + "pe0_": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, - "jend": grid.je, + "jend": grid.je+1, "kend": grid.npz+1, }, - "pe3_v": { + "pe3_": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, - "jend": grid.je, + "jend": grid.je+1, "kend": grid.npz+1, }, - # "v_": { - # "istart": grid.isd, - # "iend": grid.ied+1, - # "jstart": grid.jsd, - # "jend": grid.jed, - # "kend": grid.npz-1, - # }, - - # "mfx_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, - - # "cx_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.jsd, - # "jend": grid.jed, - # "kend": grid.npz-1, - # }, + "v_": { + "istart": grid.isd, + "iend": grid.ied+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, + + "mfx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + + "cx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, } grid_indexing = stencil_factory.grid_indexing @@ -119,46 +119,46 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto origin=grid_indexing.origin_compute(), domain=( grid_indexing.domain[0] + 1, - grid_indexing.domain[1], + 1, grid_indexing.domain[2] + 1, ), ) def compute_from_storage(self, inputs): - # self._map1_ppm_v = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # inputs["kord_mt"], - # -1, - # dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], - # ) + self._map1_ppm_v = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_mt"], + -1, + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + ) self._pressures_mapv( inputs["pe_"], inputs["ak"], inputs["bk"], - inputs["pe0_v"], - inputs["pe3_v"], + inputs["pe0_"], + inputs["pe3_"], ) - # self._map1_ppm_v( - # inputs["v_"], - # inputs["pe0_v"], - # inputs["pe3_v"], - # interp=False, - # ) + self._map1_ppm_v( + inputs["v_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) - # self._map1_ppm_v( - # inputs["mfx_"], - # inputs["pe0_v"], - # inputs["pe3_v"], - # interp=False, - # ) + self._map1_ppm_v( + inputs["mfx_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) - # self._map1_ppm_v( - # inputs["cx_"], - # inputs["pe0_v"], - # inputs["pe3_v"], - # interp=False, - # ) + self._map1_ppm_v( + inputs["cx_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) return inputs From cf66fb5ca7e09ae8f19ff524c34c4a4174f178d6 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 9 Jan 2025 14:02:12 -0800 Subject: [PATCH 097/252] U remapping added to Remapping translate test and updated tracer serialization --- fv_mapz.F90.SER | 198 ++++++++++++------ .../translate/translate_remapping_GEOS.py | 177 +++++++++++----- 2 files changed, 257 insertions(+), 118 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 297077a9..cdfdca6e 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -381,12 +381,16 @@ contains !$ser verbatim do i = is, ie !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim pe_3d(i,j-1,k) = pe(i,k,j-1) !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) !$ser verbatim enddo !$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim pe_3d(i,j-1,k) = pe(i,k,j-1) +!$ser verbatim enddo +!$ser verbatim enddo !$ser verbatim do k = 1, km !$ser verbatim do i = is, ie !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) @@ -408,13 +412,13 @@ contains !$ser data pe0_=pe0_3d !$ser data pe3_=pe3_3d !$ser data ptop=ptop -!$ser data qvapor_js=q(:,j,:,sphum) -!$ser data qliquid_js=q(:,j,:,liq_wat) -!$ser data qice_js=q(:,j,:,ice_wat) -!$ser data qrain_js=q(:,j,:,rainwat) -!$ser data qsnow_js=q(:,j,:,snowwat) -!$ser data qgraupel_js=q(:,j,:,graupel) -!$ser data qcld_js=q(:,j,:,cld_amt) +!$ser data qvapor=q(:,:,:,sphum) +!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qice=q(:,:,:,ice_wat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qgraupel=q(:,:,:,graupel) +!$ser data qcld=q(:,:,:,cld_amt) !$ser data qo3mr=q(:,:,:,8) !$ser data qsgs_tke=q(:,:,:,9) !$ser data q_con=q_con @@ -439,12 +443,19 @@ contains !$ser data ws_=ws !$ser data w=w !$ser data u=u +!$ser data v=v !$ser verbatim if (present(mfy)) then !$ser data mfy=mfy !$ser verbatim endif !$ser verbatim if (present(cy)) then !$ser data cy=cy !$ser verbatim endif +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif !$ser data w_max=w_max_ !$ser data w_min=w_min_ !$ser data kord_mt=kord_mt @@ -1304,62 +1315,62 @@ contains !$ser verbatim endif !$ser verbatim endif -!$ser verbatim if(j == js) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe1_=pe1_3d pe2_=pe2_3d -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data dp2_3d=dp2_3d -!$ser data pn1_3d=pn1_3d -!$ser data pn2_3d=pn2_3d -!$ser data pk2_3d=pk2_3d +!!$ser verbatim if(j == js) then +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie +!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km +!!$ser verbatim do i = is, ie +!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie+1 +!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe1_=pe1_3d pe2_=pe2_3d +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data dp2_3d=dp2_3d +!!$ser data pn1_3d=pn1_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data pk2_3d=pk2_3d -!$ser data qvapor=q(:,:,:,1) -!$ser data qliquid=q(:,:,:,2) -!$ser data qice=q(:,:,:,3) -!$ser data qrain=q(:,:,:,4) -!$ser data qsnow=q(:,:,:,5) -!$ser data qgraupel=q(:,:,:,6) -!$ser data qcld=q(:,:,:,7) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) +!!$ser data qvapor=q(:,:,:,1) +!!$ser data qliquid=q(:,:,:,2) +!!$ser data qice=q(:,:,:,3) +!!$ser data qrain=q(:,:,:,4) +!!$ser data qsnow=q(:,:,:,5) +!!$ser data qgraupel=q(:,:,:,6) +!!$ser data qcld=q(:,:,:,7) +!!$ser data qo3mr=q(:,:,:,8) +!!$ser data qsgs_tke=q(:,:,:,9) -!$ser data w=w +!!$ser data w=w -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data u=u -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser verbatim endif +!!$ser data pe0_=pe0_3d +!!$ser data pe3_=pe3_3d +!!$ser data u=u +!!$ser verbatim if (present(mfy)) then +!!$ser data mfy=mfy +!!$ser verbatim endif +!!$ser verbatim if (present(cy)) then +!!$ser data cy=cy +!!$ser verbatim endif +!!$ser verbatim endif ! Note : This serialization portion will test the update of pe0 ! and pe3 @@ -1444,6 +1455,69 @@ contains !$ser verbatim endif !$ser verbatim endif +!$ser verbatim if(j == js) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe1_=pe1_3d pe2_=pe2_3d +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data dp2_3d=dp2_3d +!$ser data pn1_3d=pn1_3d +!$ser data pn2_3d=pn2_3d +!$ser data pk2_3d=pk2_3d + +!$ser data qvapor=q(:,:,:,1) +!$ser data qliquid=q(:,:,:,2) +!$ser data qice=q(:,:,:,3) +!$ser data qrain=q(:,:,:,4) +!$ser data qsnow=q(:,:,:,5) +!$ser data qgraupel=q(:,:,:,6) +!$ser data qcld=q(:,:,:,7) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) + +!$ser data w=w + +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser data u=u +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser data v=v +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser verbatim endif 1000 continue ! Update pressure variables and get new pkz, T_v, and omega diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 1a828fa8..44fb4bfd 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -6,7 +6,7 @@ from pyFV3.stencils import moist_cv from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment -from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax +from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax, pressures_mapv from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py from ndsl.constants import ( @@ -87,22 +87,43 @@ def __init__( "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je+1, - "kend": grid.npz + "kend": grid.npz+1 }, "pe3_": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je+1, - "kend": grid.npz - }, - "qvapor": {"serialname": "qvapor_js"}, - "qliquid": {"serialname": "qliquid_js"}, - "qice": {"serialname": "qice_js"}, - "qrain": {"serialname": "qrain_js"}, - "qsnow": {"serialname": "qsnow_js"}, - "qgraupel": {"serialname": "qgraupel_js"}, - "qcld": {"serialname": "qcld_js"}, + "kend": grid.npz+1 + }, + # "qvapor": {"serialname": "qvapor_js"}, + # "qliquid": {"serialname": "qliquid_js"}, + # "qice": {"serialname": "qice_js"}, + # "qrain": {"serialname": "qrain_js"}, + # "qsnow": {"serialname": "qsnow_js"}, + # "qgraupel": {"serialname": "qgraupel_js"}, + # "qcld": {"serialname": "qcld_js"}, + "qvapor": { + "kend": grid.npz-1, + }, + "qliquid": { + "kend": grid.npz-1, + }, + "qice": { + "kend": grid.npz-1, + }, + "qrain": { + "kend": grid.npz-1, + }, + "qsnow": { + "kend": grid.npz-1, + }, + "qgraupel": { + "kend": grid.npz-1, + }, + "qcld": { + "kend": grid.npz-1, + }, "qo3mr": { "kend": grid.npz-1, }, @@ -175,6 +196,13 @@ def __init__( "jend": grid.jed+1, "kend": grid.npz-1, }, + "v": { + "istart": grid.isd, + "iend": grid.ied+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, "mfy": { "istart": grid.is_, "iend": grid.ie, @@ -190,6 +218,21 @@ def __init__( "jend": grid.je+1, "kend": grid.npz-1, }, + "mfx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + + "cx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, } # self.write_vars = ["gz", "cvm"] self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] @@ -233,14 +276,14 @@ def __init__( "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je+1, - "kend": grid.npz + "kend": grid.npz+1 }, "pe3_": { "istart": grid.is_, "iend": grid.ie+1, "jstart": grid.js, "jend": grid.je+1, - "kend": grid.npz + "kend": grid.npz+1 }, "pt": {}, "cappa": {}, @@ -314,6 +357,13 @@ def __init__( "jend": grid.jed+1, "kend": grid.npz-1, }, + "v": { + "istart": grid.isd, + "iend": grid.ied+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, "mfy": { "istart": grid.is_, @@ -330,6 +380,21 @@ def __init__( "jend": grid.je+1, "kend": grid.npz-1, }, + "mfx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, + + "cx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, } self.stencil_factory = stencil_factory @@ -450,23 +515,15 @@ def __init__( domain=(1,1,grid_indexing.domain[2] + 1) ) - # self._pressures_mapv = stencil_factory.from_origin_domain( - # pressures_mapv, - # origin=grid_indexing.origin_compute(), - # domain=( - # grid_indexing.domain[0] + 1, - # grid_indexing.domain[1], - # grid_indexing.domain[2] + 1, - # ), - # ) - - # self._map1_ppm_v = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # inputs["kord_mt"], - # -1, - # dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], - # ) + self._pressures_mapv = stencil_factory.from_origin_domain( + pressures_mapv, + origin=grid_indexing.origin_compute(), + domain=( + grid_indexing.domain[0] + 1, + 1, + grid_indexing.domain[2] + 1, + ), + ) # self._pe_pk_delp_peln = stencil_factory.from_origin_domain( # pe_pk_delp_peln, @@ -635,35 +692,43 @@ def compute_from_storage(self, inputs): inputs["pe3_"], interp=False, ) + + self._map1_ppm_v = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_mt"], + -1, + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + ) + + self._pressures_mapv( + inputs["pe_"], + inputs["ak"], + inputs["bk"], + inputs["pe0_"], + inputs["pe3_"], + ) - # self._pressures_mapv( - # inputs["pe_"], - # inputs["ak"], - # inputs["bk"], - # inputs["pe0_v"], - # inputs["pe3_v"], - # ) - - # self._map1_ppm_v( - # inputs["v_"], - # inputs["pe0_v"], - # inputs["pe3_v"], - # interp=False, - # ) + self._map1_ppm_v( + inputs["v"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) - # self._map1_ppm_v( - # inputs["mfx_"], - # inputs["pe0_v"], - # inputs["pe3_v"], - # interp=False, - # ) + self._map1_ppm_v( + inputs["mfx_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) - # self._map1_ppm_v( - # inputs["cx_"], - # inputs["pe0_v"], - # inputs["pe3_v"], - # interp=False, - # ) + self._map1_ppm_v( + inputs["cx_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) # self._pe_pk_delp_peln(inputs["pe_"], # inputs["pk"], From cac66338237cbc8b4bb1eddedb76356ccee3b2bd Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 9 Jan 2025 14:51:08 -0800 Subject: [PATCH 098/252] Added pe_pk_delp_peln stencil to Remapping translate test --- fv_mapz.F90.SER | 194 ++++++++++++------ .../translate/translate_remapping_GEOS.py | 61 ++++-- 2 files changed, 175 insertions(+), 80 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index cdfdca6e..7456aaf4 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -371,7 +371,7 @@ contains !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & !$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,& !$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & -!$OMP peln_3d, gz_2d, pe0_3d, pe3_3d) & +!$OMP peln_3d, gz_2d, pe0_3d, pe3_3d) & !$OMP private(gz,cvm,bkh,dp2, & !$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_) do 1000 j=js,je+1 @@ -1455,69 +1455,69 @@ contains !$ser verbatim endif !$ser verbatim endif -!$ser verbatim if(j == js) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe1_=pe1_3d pe2_=pe2_3d -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data dp2_3d=dp2_3d -!$ser data pn1_3d=pn1_3d -!$ser data pn2_3d=pn2_3d -!$ser data pk2_3d=pk2_3d +!!$ser verbatim if(j == js) then +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie +!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km +!!$ser verbatim do i = is, ie +!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie+1 +!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe1_=pe1_3d pe2_=pe2_3d +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data dp2_3d=dp2_3d +!!$ser data pn1_3d=pn1_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data pk2_3d=pk2_3d -!$ser data qvapor=q(:,:,:,1) -!$ser data qliquid=q(:,:,:,2) -!$ser data qice=q(:,:,:,3) -!$ser data qrain=q(:,:,:,4) -!$ser data qsnow=q(:,:,:,5) -!$ser data qgraupel=q(:,:,:,6) -!$ser data qcld=q(:,:,:,7) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) +!!$ser data qvapor=q(:,:,:,1) +!!$ser data qliquid=q(:,:,:,2) +!!$ser data qice=q(:,:,:,3) +!!$ser data qrain=q(:,:,:,4) +!!$ser data qsnow=q(:,:,:,5) +!!$ser data qgraupel=q(:,:,:,6) +!!$ser data qcld=q(:,:,:,7) +!!$ser data qo3mr=q(:,:,:,8) +!!$ser data qsgs_tke=q(:,:,:,9) -!$ser data w=w +!!$ser data w=w -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data u=u -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser data v=v -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser verbatim endif +!!$ser data pe0_=pe0_3d +!!$ser data pe3_=pe3_3d +!!$ser data u=u +!!$ser verbatim if (present(mfy)) then +!!$ser data mfy=mfy +!!$ser verbatim endif +!!$ser verbatim if (present(cy)) then +!!$ser data cy=cy +!!$ser verbatim endif +!!$ser data v=v +!!$ser verbatim if (present(mfx)) then +!!$ser data mfx_=mfx +!!$ser verbatim endif +!!$ser verbatim if (present(cx)) then +!!$ser data cx_=cx +!!$ser verbatim endif +!!$ser verbatim endif 1000 continue ! Update pressure variables and get new pkz, T_v, and omega @@ -1532,7 +1532,8 @@ contains #ifdef SERIALIZE !$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, js2d, & #endif -!$OMP pe2_3d, pe_3d, peln_3d, pn2_3d, pk2_3d) & +!$OMP pe2_3d, pe_3d, peln_3d, pn2_3d, pk2_3d, pe0, pe0_3d, pe3_3d, & +!$OMP pe1_3d, pn1, dp2_3d, q_con, pe1, pn1_3d) & !$OMP private(gz,cvm,kp,k_next,bkh,dp2, & !$OMP pe2,pe3,pk2,pn2,phis,tpe,dlnp,tmp) do 2000 j=js,je @@ -1635,6 +1636,75 @@ contains !$ser data pk2_=pk2_3d !$ser data delp=delp !$ser data pk=pk +!$ser verbatim endif + +!$ser verbatim if(j == js) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe1_=pe1_3d pe2_=pe2_3d +!$ser data pe_=pe_3d +!$ser data peln_3d=peln_3d +!$ser data pk=pk +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data dp2_3d=dp2_3d +!$ser data pn1_3d=pn1_3d +!$ser data pn2_3d=pn2_3d +!$ser data pk2_3d=pk2_3d + +!$ser data qvapor=q(:,:,:,1) +!$ser data qliquid=q(:,:,:,2) +!$ser data qice=q(:,:,:,3) +!$ser data qrain=q(:,:,:,4) +!$ser data qsnow=q(:,:,:,5) +!$ser data qgraupel=q(:,:,:,6) +!$ser data qcld=q(:,:,:,7) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) + +!$ser data w=w + +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser data u=u +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser data v=v +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif !$ser verbatim endif !--------------------- diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 44fb4bfd..1671e44c 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -6,7 +6,7 @@ from pyFV3.stencils import moist_cv from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment -from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax, pressures_mapv +from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax, pressures_mapv, pe_pk_delp_peln from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py from ndsl.constants import ( @@ -395,6 +395,30 @@ def __init__( "jend": grid.jed, "kend": grid.npz-1, }, + + "peln_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, + + "pe_": { + "istart": grid.is_-1, + "iend": grid.ie+1, + "jstart": grid.js-1, + "jend": grid.je+1, + "kend": grid.npz + 1, + }, + + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, } self.stencil_factory = stencil_factory @@ -525,11 +549,12 @@ def __init__( ), ) - # self._pe_pk_delp_peln = stencil_factory.from_origin_domain( - # pe_pk_delp_peln, - # origin=grid_indexing.origin_compute(), - # domain=self._domain_kextra, - # ) + self._pe_pk_delp_peln = stencil_factory.from_origin_domain( + pe_pk_delp_peln, + origin=grid_indexing.origin_compute(), + domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2] + 1, + ), + ) def compute_from_storage(self, inputs): @@ -730,18 +755,18 @@ def compute_from_storage(self, inputs): interp=False, ) - # self._pe_pk_delp_peln(inputs["pe_"], - # inputs["pk"], - # inputs["delp"], - # inputs["peln_"], - # inputs["pe2_"], - # inputs["pk2_"], - # inputs["pn2_"], - # inputs["ak"], - # inputs["bk"], - # inputs["akap"], - # inputs["ptop"], - # ) + self._pe_pk_delp_peln(inputs["pe_"], + inputs["pk"], + inputs["delp"], + inputs["peln_3d"], + inputs["pe2_"], + inputs["pk2_3d"], + inputs["pn2_3d"], + inputs["ak"], + inputs["bk"], + inputs["akap"], + inputs["ptop"], + ) # NOTE : THERE WILL BE ADJUSTMENTS TO ACCOUNT FOR PKZ # self._moist_cv_pt( From eada51087bab446eb935e4a67b62a54f58b75454 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 10 Jan 2025 08:59:40 -0800 Subject: [PATCH 099/252] Verified translate_moistcvpluspkz_2d.py test --- fv_mapz.F90.SER | 27 ++++++++++++++++++- pyFV3/stencils/moist_cv.py | 3 +-- .../translate/translate_moistcvpluspkz_2d.py | 4 --- .../translate/translate_remapping_GEOS.py | 4 +-- 4 files changed, 29 insertions(+), 9 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 7456aaf4..86d4c73d 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -1525,7 +1525,7 @@ contains !$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & !$OMP remap_pt,remap_te,mfy,mfx,cx,cy,hydrostatic, & !$OMP pt,pk,rg,peln,q,nwat,liq_wat,rainwat,ice_wat,snowwat, & -!$OMP graupel,sphum,cappa,r_vir,rcp,cp,k1k,delp, & +!$OMP graupel,sphum,cappa,r_vir,rcp,cp,k1k,delp, cld_amt, & !$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, & !$OMP hs,w,ws,kord_wz,do_omega,omga,rrg,kord_mt, & @@ -1789,6 +1789,27 @@ contains if (remap_t) then ! print*,'CODE EXECUTED' ! Note: pt at this stage is T_v or T_m + !$ser verbatim if(j == js2d) then + !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' + !$ser savepoint MoistCVPlusPkz_2d-In + !$ser data qvapor_js=q(:,j,:,sphum) + !$ser data qliquid_js=q(:,j,:,liq_wat) + !$ser data qice_js=q(:,j,:,ice_wat) + !$ser data qrain_js=q(:,j,:,rainwat) + !$ser data qsnow_js=q(:,j,:,snowwat) + !$ser data qgraupel_js=q(:,j,:,graupel) + !$ser data qcld_js=q(:,j,:,cld_amt) + !$ser data gz1d=gz + !$ser data cvm=cvm + !$ser data r_vir=r_vir + !$ser data cappa=cappa + !$ser data rrg=rrg + !$ser data delp=delp + !$ser data delz=delz + !$ser data pkz=pkz + !$ser data pt=pt + !$ser data k1k=k1k + !$ser verbatim endif do k=1,km #ifdef MOIST_CAPPA call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & @@ -1803,6 +1824,10 @@ contains enddo #endif enddo + !$ser verbatim if(j == js2d) then + !$ser savepoint MoistCVPlusPkz_2d-Out + !$ser data gz1d=gz cvm=cvm pkz=pkz cappa=cappa + !$ser verbatim endif endif if (remap_pt) then print*,'CODE NOT TESTED HERE 10' diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index d22f0a26..a912d31c 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -136,7 +136,6 @@ def moist_pkz( qsnow: FloatField, qice: FloatField, qgraupel: FloatField, - q_con: FloatField, gz: FloatField, cvm: FloatField, pkz: FloatField, @@ -169,7 +168,7 @@ def moist_pkz( cvm, gz = moist_cv_nwat6_fn( qvapor, qliquid, qrain, qsnow, qice, qgraupel ) # if (nwat == 6) else moist_cv_default_fn(constants.CV_AIR) - q_con[0, 0, 0] = gz + # q_con[0, 0, 0] = gz cappa = set_cappa(qvapor, cvm, r_vir) pkz = compute_pkz_func(delp, delz, pt, cappa) diff --git a/tests/savepoint/translate/translate_moistcvpluspkz_2d.py b/tests/savepoint/translate/translate_moistcvpluspkz_2d.py index 54577da2..f26d3f5b 100644 --- a/tests/savepoint/translate/translate_moistcvpluspkz_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluspkz_2d.py @@ -29,7 +29,6 @@ def __call__( qsnow: FloatField, qice: FloatField, qgraupel: FloatField, - q_con: FloatField, gz: FloatField, cvm: FloatField, pkz: FloatField, @@ -46,7 +45,6 @@ def __call__( qsnow, qice, qgraupel, - q_con, gz, cvm, pkz, @@ -80,7 +78,6 @@ def __init__( "cvm": {"kstart": grid.is_, "axis": 0}, "delp": {}, "delz": {}, - "q_con": {}, "pkz": {"istart": grid.is_, "jstart": grid.js}, "pt": {}, "cappa": {}, @@ -116,7 +113,6 @@ def __init__( "jend": grid.je, }, "cappa": {}, - "q_con": {}, } def compute_from_storage(self, inputs): diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 1671e44c..0c9f9323 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -768,7 +768,7 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - # NOTE : THERE WILL BE ADJUSTMENTS TO ACCOUNT FOR PKZ + # NOTE :USE THE STENCIL moist_pkz to do the moist_cv and pkz calculation # self._moist_cv_pt( # inputs["qvapor"], # inputs["qliquid"], @@ -777,7 +777,7 @@ def compute_from_storage(self, inputs): # inputs["qice"], # inputs["qgraupel"], # inputs["q_con"], - # inputs["pt"], + # inputs["pkz"], # inputs["cappa"], # inputs["delp"], # inputs["delz"], From bc58a794d41d101c4e8a96c833d0c59445913511 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 10 Jan 2025 09:46:09 -0800 Subject: [PATCH 100/252] Took out serialized gz and cvm variables in moist_cv pkz translate test since those are computed internally --- pyFV3/stencils/moist_cv.py | 5 ---- .../translate/translate_moistcvpluspkz_2d.py | 24 ------------------- 2 files changed, 29 deletions(-) diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index a912d31c..4e38ef51 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -136,8 +136,6 @@ def moist_pkz( qsnow: FloatField, qice: FloatField, qgraupel: FloatField, - gz: FloatField, - cvm: FloatField, pkz: FloatField, pt: FloatField, cappa: FloatField, @@ -153,9 +151,6 @@ def moist_pkz( qsnow (in): qice (in): qgraupel (in): - q_con (out): - gz (out): - cvm (out): pkz (out): pt (in): cappa (out): diff --git a/tests/savepoint/translate/translate_moistcvpluspkz_2d.py b/tests/savepoint/translate/translate_moistcvpluspkz_2d.py index f26d3f5b..657d27dc 100644 --- a/tests/savepoint/translate/translate_moistcvpluspkz_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluspkz_2d.py @@ -29,8 +29,6 @@ def __call__( qsnow: FloatField, qice: FloatField, qgraupel: FloatField, - gz: FloatField, - cvm: FloatField, pkz: FloatField, pt: FloatField, cappa: FloatField, @@ -45,8 +43,6 @@ def __call__( qsnow, qice, qgraupel, - gz, - cvm, pkz, pt, cappa, @@ -74,38 +70,18 @@ def __init__( "qrain": {"serialname": "qrain_js"}, "qsnow": {"serialname": "qsnow_js"}, "qgraupel": {"serialname": "qgraupel_js"}, - "gz": {"serialname": "gz1d", "kstart": grid.is_, "axis": 0}, - "cvm": {"kstart": grid.is_, "axis": 0}, "delp": {}, "delz": {}, "pkz": {"istart": grid.is_, "jstart": grid.js}, "pt": {}, "cappa": {}, } - self.write_vars = ["gz", "cvm"] for k, v in self.in_vars["data_vars"].items(): if k not in self.write_vars: v["axis"] = 1 self.in_vars["parameters"] = ["r_vir"] self.out_vars = { - "gz": { - "serialname": "gz1d", - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.js, - "kstart": grid.npz - 1, - "kend": grid.npz - 1, - }, - "cvm": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.js, - "kstart": grid.npz - 1, - "kend": grid.npz - 1, - }, "pkz": { "istart": grid.is_, "iend": grid.ie, From e7377c23c2d31123b229e969254b13c762d8b556 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 10 Jan 2025 10:53:28 -0800 Subject: [PATCH 101/252] Added moist_cv_pkz stencil to Remapping translate test --- fv_mapz.F90.SER | 210 ++++++++++++------ .../translate/translate_remapping_GEOS.py | 53 +++-- 2 files changed, 172 insertions(+), 91 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 86d4c73d..91546df1 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -437,6 +437,7 @@ contains !$ser data r_vir=r_vir !$ser data cvm=cvm !$ser data pk=pk +!$ser data pkz=pkz !$ser data akap=akap !$ser data t_min=qmin !$ser data kord_wz=kord_wz @@ -1638,74 +1639,74 @@ contains !$ser data pk=pk !$ser verbatim endif -!$ser verbatim if(j == js) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe1_=pe1_3d pe2_=pe2_3d -!$ser data pe_=pe_3d -!$ser data peln_3d=peln_3d -!$ser data pk=pk -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data dp2_3d=dp2_3d -!$ser data pn1_3d=pn1_3d -!$ser data pn2_3d=pn2_3d -!$ser data pk2_3d=pk2_3d +!!$ser verbatim if(j == js) then +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie +!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) +!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) +!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km +!!$ser verbatim do i = is, ie +!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do i = is, ie+1 +!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe1_=pe1_3d pe2_=pe2_3d +!!$ser data pe_=pe_3d +!!$ser data peln_3d=peln_3d +!!$ser data pk=pk +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data dp2_3d=dp2_3d +!!$ser data pn1_3d=pn1_3d +!!$ser data pn2_3d=pn2_3d +!!$ser data pk2_3d=pk2_3d -!$ser data qvapor=q(:,:,:,1) -!$ser data qliquid=q(:,:,:,2) -!$ser data qice=q(:,:,:,3) -!$ser data qrain=q(:,:,:,4) -!$ser data qsnow=q(:,:,:,5) -!$ser data qgraupel=q(:,:,:,6) -!$ser data qcld=q(:,:,:,7) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) +!!$ser data qvapor=q(:,:,:,1) +!!$ser data qliquid=q(:,:,:,2) +!!$ser data qice=q(:,:,:,3) +!!$ser data qrain=q(:,:,:,4) +!!$ser data qsnow=q(:,:,:,5) +!!$ser data qgraupel=q(:,:,:,6) +!!$ser data qcld=q(:,:,:,7) +!!$ser data qo3mr=q(:,:,:,8) +!!$ser data qsgs_tke=q(:,:,:,9) -!$ser data w=w +!!$ser data w=w -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data u=u -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser data v=v -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser verbatim endif +!!$ser data pe0_=pe0_3d +!!$ser data pe3_=pe3_3d +!!$ser data u=u +!!$ser verbatim if (present(mfy)) then +!!$ser data mfy=mfy +!!$ser verbatim endif +!!$ser verbatim if (present(cy)) then +!!$ser data cy=cy +!!$ser verbatim endif +!!$ser data v=v +!!$ser verbatim if (present(mfx)) then +!!$ser data mfx_=mfx +!!$ser verbatim endif +!!$ser verbatim if (present(cx)) then +!!$ser data cx_=cx +!!$ser verbatim endif +!!$ser verbatim endif !--------------------- ! Compute pkz and T_v @@ -1799,8 +1800,6 @@ contains !$ser data qsnow_js=q(:,j,:,snowwat) !$ser data qgraupel_js=q(:,j,:,graupel) !$ser data qcld_js=q(:,j,:,cld_amt) - !$ser data gz1d=gz - !$ser data cvm=cvm !$ser data r_vir=r_vir !$ser data cappa=cappa !$ser data rrg=rrg @@ -1825,9 +1824,80 @@ contains #endif enddo !$ser verbatim if(j == js2d) then - !$ser savepoint MoistCVPlusPkz_2d-Out - !$ser data gz1d=gz cvm=cvm pkz=pkz cappa=cappa - !$ser verbatim endif + !$ser savepoint MoistCVPlusPkz_2d-Out + !$ser data pkz=pkz + !$ser data cappa=cappa + !$ser verbatim endif + + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) + !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie+1 + !$ser verbatim pe0_3d(i,j,k) = pe0(i,k) + !$ser verbatim pe3_3d(i,j,k) = pe3(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data pe_=pe_3d + !$ser data peln_3d=peln_3d + !$ser data pk=pk + !$ser data pkz=pkz + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + + !$ser data w=w + + !$ser data pe0_=pe0_3d + !$ser data pe3_=pe3_3d + !$ser data u=u + !$ser verbatim if (present(mfy)) then + !$ser data mfy=mfy + !$ser verbatim endif + !$ser verbatim if (present(cy)) then + !$ser data cy=cy + !$ser verbatim endif + !$ser data v=v + !$ser verbatim if (present(mfx)) then + !$ser data mfx_=mfx + !$ser verbatim endif + !$ser verbatim if (present(cx)) then + !$ser data cx_=cx + !$ser verbatim endif + !$ser verbatim endif endif if (remap_pt) then print*,'CODE NOT TESTED HERE 10' diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 0c9f9323..31776a2e 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -173,6 +173,12 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, + "pkz": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, "pk2_3d": { "istart": grid.is_, "iend": grid.ie, @@ -419,6 +425,12 @@ def __init__( "jend": grid.je, "kend": grid.npz+1, }, + "pkz": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, } self.stencil_factory = stencil_factory @@ -480,12 +492,6 @@ def __init__( domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1), ) - self._moist_cv_pt = stencil_factory.from_origin_domain( - moist_pt, - origin=grid.compute_origin(), - domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2]), - ) - self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( moist_cv_pt_pressure, # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, @@ -556,6 +562,12 @@ def __init__( ), ) + self._moist_cv_pkz = stencil_factory.from_origin_domain( + moist_cv.moist_pkz, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction @@ -768,21 +780,20 @@ def compute_from_storage(self, inputs): inputs["ptop"], ) - # NOTE :USE THE STENCIL moist_pkz to do the moist_cv and pkz calculation - # self._moist_cv_pt( - # inputs["qvapor"], - # inputs["qliquid"], - # inputs["qrain"], - # inputs["qsnow"], - # inputs["qice"], - # inputs["qgraupel"], - # inputs["q_con"], - # inputs["pkz"], - # inputs["cappa"], - # inputs["delp"], - # inputs["delz"], - # inputs["r_vir"], - # ) + self._moist_cv_pkz( + inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + inputs["pkz"], + inputs["pt"], + inputs["cappa"], + inputs["delp"], + inputs["delz"], + Float(inputs["r_vir"]), + ) # If loop based on if( last_step .and. (.not.do_adiabatic_init) ) then # PHIS computation From 973d0ba0bb927e4689f0572853115ca50cef65c5 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 13 Jan 2025 10:12:51 -0800 Subject: [PATCH 102/252] Added moist_te translate test and stencil --- pyFV3/stencils/moist_cv.py | 48 +++++++- tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_moistcvpluste_2d.py | 107 ++++++++++++++++++ 3 files changed, 156 insertions(+), 2 deletions(-) create mode 100644 tests/savepoint/translate/translate_moistcvpluste_2d.py diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index 4e38ef51..f318d603 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -2,6 +2,7 @@ from gt4py.cartesian.gtscript import ( __INLINED, PARALLEL, + FORWARD, computation, exp, interval, @@ -9,7 +10,7 @@ ) import ndsl.constants as constants -from ndsl.dsl.typing import Float, FloatField +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ @gtscript.function @@ -167,6 +168,51 @@ def moist_pkz( cappa = set_cappa(qvapor, cvm, r_vir) pkz = compute_pkz_func(delp, delz, pt, cappa) +def moist_te( + qvapor: FloatField, + qliquid: FloatField, + qrain: FloatField, + qsnow: FloatField, + qice: FloatField, + qgraupel: FloatField, + u: FloatField, + v: FloatField, + w: FloatField, + te: FloatFieldIJ, + pt: FloatField, + phis: FloatField, + delp: FloatField, + rsin2: FloatFieldIJ, + cosa_s: FloatFieldIJ, +): + """ + Args: + qvapor (in): + qliquid (in): + qrain (in): + qsnow (in): + qice (in): + qgraupel (in): + u (in): + v (in): + w (in): + te (out): + pt (in): + phis (in): + delp (in): + rsin2 (in): + cosa_s (in): + """ + with computation(FORWARD), interval(...): + cvm, gz = moist_cv_nwat6_fn( + qvapor, qliquid, qrain, qsnow, qice, qgraupel + ) + + te = te + delp * (cvm * pt + 0.5 * (phis + phis[0,0,1] + \ + w**2.0 + 0.5*rsin2 * (u**2.0 + u[0,1,0]**2.0 \ + + v**2.0 + v[1,0,0]**2.0 + - (u + u[0,1,0]) * (v + v[1,0,0]) * cosa_s))) + def fv_setup( qvapor: FloatField, diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index e956b33d..589a1939 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -99,4 +99,5 @@ from .translate_Pressures_mapV import TranslatePressures_mapV from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln from .translate_getMPIprop import TranslateGetMPIProp -from .translate_mpp_global_sum import TranslateMpp_global_sum \ No newline at end of file +from .translate_mpp_global_sum import TranslateMpp_global_sum +from .translate_moistcvpluste_2d import TranslateMoistCVPlusTe_2d \ No newline at end of file diff --git a/tests/savepoint/translate/translate_moistcvpluste_2d.py b/tests/savepoint/translate/translate_moistcvpluste_2d.py new file mode 100644 index 00000000..6847f836 --- /dev/null +++ b/tests/savepoint/translate/translate_moistcvpluste_2d.py @@ -0,0 +1,107 @@ +from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j +from pyFV3.stencils import moist_cv + +class TranslateMoistCVPlusTe_2d(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.in_vars["data_vars"] = { + "qvapor": {"serialname": "qvapor_js"}, + "qliquid": {"serialname": "qliquid_js"}, + "qice": {"serialname": "qice_js"}, + "qrain": {"serialname": "qrain_js"}, + "qsnow": {"serialname": "qsnow_js"}, + "qgraupel": {"serialname": "qgraupel_js"}, + "delp": {}, + "pt": {}, + "phis_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "te_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + "u": { + "istart": grid.isd, + "iend": grid.ied + 1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz, + }, + "v": { + "istart": grid.isd, + "iend": grid.ied + 1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz, + }, + "w":{ + "kend": grid.npz, + }, + "cosa_s": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, + "rsin2": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, + } + self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel"] + for k, v in self.in_vars["data_vars"].items(): + # if k not in self.write_vars: + if k in self.write_vars: + v["axis"] = 1 + + self.out_vars = { + "te_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + } + + self.compute_func = stencil_factory.from_origin_domain( + moist_cv.moist_te, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + + def compute_from_storage(self, inputs): + for name, value in inputs.items(): + if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: + inputs[name] = self.make_storage_data( + pad_field_in_j( + value, self.grid.njd, backend=self.stencil_factory.backend + ) + ) + + self.compute_func(inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + inputs["u"], + inputs["v"], + inputs["w"], + inputs["te_2d_"], + inputs["pt"], + inputs["phis_"], + inputs["delp"], + inputs["rsin2"], + inputs["cosa_s"], + ) + + return inputs From 54c5e377dbba864ce9f025d01f15fe4b2b3000e6 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 13 Jan 2025 10:14:20 -0800 Subject: [PATCH 103/252] Added fv_mapz serialization file for moist_te --- fv_mapz.F90.SER | 555 ++++-------------------------------------------- 1 file changed, 42 insertions(+), 513 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 91546df1..4c91154d 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -223,12 +223,14 @@ contains logical:: remap_t, remap_pt, remap_te real, dimension(is:ie, js:je, km+1) :: pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, peln_3d + real, dimension(is:ie, js:je, km+1) :: phis_3d real, dimension(is-1:ie+1,js-1:je+1,km+1) :: pe_3d real, dimension(is:ie+1, js:je+1, km+1) :: pe0_3d, pe3_3d real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d - real, dimension(is:ie, js:je) :: gz_2d + real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32 + real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s logical :: serial_flag @@ -250,6 +252,7 @@ contains !$ser verbatim pn2_3d = 0.0 !$ser verbatim pk2_3d = pk !$ser verbatim gz_2d = 0.0 + !$ser verbatim phis_3d = 0.0 remap_t = .false. remap_pt = .false. @@ -516,32 +519,6 @@ contains !$ser data gz1d=gz cvm=cvm pt=pt cappa=cappa q_con=q_con !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) - !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !!$ser savepoint Remapping_GEOS-Out - !!$ser data pe1_=pe1_3d - !!$ser data pe2_=pe2_3d - !!$ser data gz1d=gz - !!$ser data cvm=cvm - !!$ser data pt=pt - !!$ser data cappa=cappa - !!$ser data q_con=q_con - !!$ser data dp2_3d=dp2_3d - !$ser verbatim endif - endif ! hydro test elseif (remap_pt) then print*,'CODE NOT TESTED HERE 3' @@ -654,32 +631,6 @@ contains pk2(i,km+1) = pk1(i,km+1) enddo -!!$ser verbatim if(j == js) then -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie -!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km -!!$ser verbatim do i = is, ie -!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe1_=pe1_3d pe2_=pe2_3d -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data dp2_3d=dp2_3d -!!$ser data pn2_3d=pn2_3d -!!$ser verbatim endif - do k=2,km do i=is,ie pn2(i,k) = log(pe2(i,k)) @@ -687,35 +638,6 @@ contains enddo enddo -!!$ser verbatim if(j == js) then -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie -!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km -!!$ser verbatim do i = is, ie -!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe1_=pe1_3d pe2_=pe2_3d -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data dp2_3d=dp2_3d -!!$ser data pn1_3d=pn1_3d -!!$ser data pn2_3d=pn2_3d -!!$ser data pk2_3d=pk2_3d -!!$ser verbatim endif - if (remap_te) then print*,'CODE NOT TESTED HERE 5' !---------------------------------- @@ -773,35 +695,6 @@ contains !$ser savepoint Map_Scalar-Out !$ser data q1=pt !$ser verbatim endif - - !!$ser verbatim if(j == js) then - !!$ser verbatim do k = 1, km+1 - !!$ser verbatim do i = is, ie - !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser verbatim do k = 1, km - !!$ser verbatim do i = is, ie - !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser savepoint Remapping_GEOS-Out - !!$ser data pe1_=pe1_3d pe2_=pe2_3d - !!$ser data q_con=q_con - !!$ser data pt=pt - !!$ser data cappa=cappa - !!$ser data delp=delp - !!$ser data delz=delz - !!$ser data ps=ps - !!$ser data dp2_3d=dp2_3d - !!$ser data pn1_3d=pn1_3d - !!$ser data pn2_3d=pn2_3d - !!$ser data pk2_3d=pk2_3d - !!$ser verbatim endif endif endif @@ -845,44 +738,6 @@ contains !$ser data qsgs_tke=q(:,:,:,9) !$ser verbatim endif - !!$ser verbatim if(j == js) then - !!$ser verbatim do k = 1, km+1 - !!$ser verbatim do i = is, ie - !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser verbatim do k = 1, km - !!$ser verbatim do i = is, ie - !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser savepoint Remapping_GEOS-Out - !!$ser data pe1_=pe1_3d pe2_=pe2_3d - !!$ser data q_con=q_con - !!$ser data pt=pt - !!$ser data cappa=cappa - !!$ser data delp=delp - !!$ser data delz=delz - !!$ser data ps=ps - !!$ser data dp2_3d=dp2_3d - !!$ser data pn1_3d=pn1_3d - !!$ser data pn2_3d=pn2_3d - !!$ser data pk2_3d=pk2_3d - - !!$ser data qvapor=q(:,:,:,1) - !!$ser data qliquid=q(:,:,:,2) - !!$ser data qice=q(:,:,:,3) - !!$ser data qrain=q(:,:,:,4) - !!$ser data qsnow=q(:,:,:,5) - !!$ser data qgraupel=q(:,:,:,6) - !!$ser data qcld=q(:,:,:,7) - !!$ser data qo3mr=q(:,:,:,8) - !!$ser data qsgs_tke=q(:,:,:,9) - !!$ser verbatim endif elseif ( nq > 0 ) then print*,'CODE NOT TESTED HERE 7' ! Remap one tracer at a time @@ -923,47 +778,6 @@ contains !$ser data w_=w !$ser verbatim endif - !!$ser verbatim if(j == js) then - !!$ser verbatim do k = 1, km+1 - !!$ser verbatim do i = is, ie - !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser verbatim do k = 1, km - !!$ser verbatim do i = is, ie - !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser savepoint Remapping_GEOS-Out - !!$ser data pe1_=pe1_3d pe2_=pe2_3d - !!$ser data q_con=q_con - !!$ser data pt=pt - !!$ser data cappa=cappa - !!$ser data delp=delp - !!$ser data delz=delz - !!$ser data ps=ps - !!$ser data dp2_3d=dp2_3d - !!$ser data pn1_3d=pn1_3d - !!$ser data pn2_3d=pn2_3d - !!$ser data pk2_3d=pk2_3d - - !!$ser data qvapor=q(:,:,:,1) - !!$ser data qliquid=q(:,:,:,2) - !!$ser data qice=q(:,:,:,3) - !!$ser data qrain=q(:,:,:,4) - !!$ser data qsnow=q(:,:,:,5) - !!$ser data qgraupel=q(:,:,:,6) - !!$ser data qcld=q(:,:,:,7) - !!$ser data qo3mr=q(:,:,:,8) - !!$ser data qsgs_tke=q(:,:,:,9) - - !!$ser data w=w - !!$ser verbatim endif - !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie @@ -1009,47 +823,6 @@ contains !$ser data delz_=delz !$ser verbatim endif - !!$ser verbatim if(j == js) then - !!$ser verbatim do k = 1, km+1 - !!$ser verbatim do i = is, ie - !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser verbatim do k = 1, km - !!$ser verbatim do i = is, ie - !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser savepoint Remapping_GEOS-Out - !!$ser data pe1_=pe1_3d pe2_=pe2_3d - !!$ser data q_con=q_con - !!$ser data pt=pt - !!$ser data cappa=cappa - !!$ser data delp=delp - !!$ser data delz=delz - !!$ser data ps=ps - !!$ser data dp2_3d=dp2_3d - !!$ser data pn1_3d=pn1_3d - !!$ser data pn2_3d=pn2_3d - !!$ser data pk2_3d=pk2_3d - - !!$ser data qvapor=q(:,:,:,1) - !!$ser data qliquid=q(:,:,:,2) - !!$ser data qice=q(:,:,:,3) - !!$ser data qrain=q(:,:,:,4) - !!$ser data qsnow=q(:,:,:,5) - !!$ser data qgraupel=q(:,:,:,6) - !!$ser data qcld=q(:,:,:,7) - !!$ser data qo3mr=q(:,:,:,8) - !!$ser data qsgs_tke=q(:,:,:,9) - - !!$ser data w=w - !!$ser verbatim endif - !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array if ( w_limiter ) then @@ -1136,47 +909,6 @@ contains !$ser savepoint W_fix_consrv_moment-Out !$ser data w=w !$ser verbatim endif - - !!$ser verbatim if(j == js) then - !!$ser verbatim do k = 1, km+1 - !!$ser verbatim do i = is, ie - !!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser verbatim do k = 1, km - !!$ser verbatim do i = is, ie - !!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !!$ser verbatim enddo - !!$ser verbatim enddo - !!$ser savepoint Remapping_GEOS-Out - !!$ser data pe1_=pe1_3d pe2_=pe2_3d - !!$ser data q_con=q_con - !!$ser data pt=pt - !!$ser data cappa=cappa - !!$ser data delp=delp - !!$ser data delz=delz - !!$ser data ps=ps - !!$ser data dp2_3d=dp2_3d - !!$ser data pn1_3d=pn1_3d - !!$ser data pn2_3d=pn2_3d - !!$ser data pk2_3d=pk2_3d - - !!$ser data qvapor=q(:,:,:,1) - !!$ser data qliquid=q(:,:,:,2) - !!$ser data qice=q(:,:,:,3) - !!$ser data qrain=q(:,:,:,4) - !!$ser data qsnow=q(:,:,:,5) - !!$ser data qgraupel=q(:,:,:,6) - !!$ser data qcld=q(:,:,:,7) - !!$ser data qo3mr=q(:,:,:,8) - !!$ser data qsgs_tke=q(:,:,:,9) - - !!$ser data w=w - !!$ser verbatim endif endif endif @@ -1231,56 +963,6 @@ contains enddo enddo -!!$ser verbatim if(j == js) then -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie -!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km -!!$ser verbatim do i = is, ie -!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie+1 -!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe1_=pe1_3d pe2_=pe2_3d -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data dp2_3d=dp2_3d -!!$ser data pn1_3d=pn1_3d -!!$ser data pn2_3d=pn2_3d -!!$ser data pk2_3d=pk2_3d - -!!$ser data qvapor=q(:,:,:,1) -!!$ser data qliquid=q(:,:,:,2) -!!$ser data qice=q(:,:,:,3) -!!$ser data qrain=q(:,:,:,4) -!!$ser data qsnow=q(:,:,:,5) -!!$ser data qgraupel=q(:,:,:,6) -!!$ser data qcld=q(:,:,:,7) -!!$ser data qo3mr=q(:,:,:,8) -!!$ser data qsgs_tke=q(:,:,:,9) - -!!$ser data w=w - -!!$ser data pe0_=pe0_3d -!!$ser data pe3_=pe3_3d -!!$ser verbatim endif - call map1_ppm( km, pe0(is:ie,:), u, gz, & km, pe3(is:ie,:), u, & is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) @@ -1316,63 +998,6 @@ contains !$ser verbatim endif !$ser verbatim endif -!!$ser verbatim if(j == js) then -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie -!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km -!!$ser verbatim do i = is, ie -!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie+1 -!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe1_=pe1_3d pe2_=pe2_3d -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data dp2_3d=dp2_3d -!!$ser data pn1_3d=pn1_3d -!!$ser data pn2_3d=pn2_3d -!!$ser data pk2_3d=pk2_3d - -!!$ser data qvapor=q(:,:,:,1) -!!$ser data qliquid=q(:,:,:,2) -!!$ser data qice=q(:,:,:,3) -!!$ser data qrain=q(:,:,:,4) -!!$ser data qsnow=q(:,:,:,5) -!!$ser data qgraupel=q(:,:,:,6) -!!$ser data qcld=q(:,:,:,7) -!!$ser data qo3mr=q(:,:,:,8) -!!$ser data qsgs_tke=q(:,:,:,9) - -!!$ser data w=w - -!!$ser data pe0_=pe0_3d -!!$ser data pe3_=pe3_3d -!!$ser data u=u -!!$ser verbatim if (present(mfy)) then -!!$ser data mfy=mfy -!!$ser verbatim endif -!!$ser verbatim if (present(cy)) then -!!$ser data cy=cy -!!$ser verbatim endif -!!$ser verbatim endif - ! Note : This serialization portion will test the update of pe0 ! and pe3 !$ser verbatim if(j == js2d) then @@ -1456,69 +1081,6 @@ contains !$ser verbatim endif !$ser verbatim endif -!!$ser verbatim if(j == js) then -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie -!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km -!!$ser verbatim do i = is, ie -!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie+1 -!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe1_=pe1_3d pe2_=pe2_3d -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data dp2_3d=dp2_3d -!!$ser data pn1_3d=pn1_3d -!!$ser data pn2_3d=pn2_3d -!!$ser data pk2_3d=pk2_3d - -!!$ser data qvapor=q(:,:,:,1) -!!$ser data qliquid=q(:,:,:,2) -!!$ser data qice=q(:,:,:,3) -!!$ser data qrain=q(:,:,:,4) -!!$ser data qsnow=q(:,:,:,5) -!!$ser data qgraupel=q(:,:,:,6) -!!$ser data qcld=q(:,:,:,7) -!!$ser data qo3mr=q(:,:,:,8) -!!$ser data qsgs_tke=q(:,:,:,9) - -!!$ser data w=w - -!!$ser data pe0_=pe0_3d -!!$ser data pe3_=pe3_3d -!!$ser data u=u -!!$ser verbatim if (present(mfy)) then -!!$ser data mfy=mfy -!!$ser verbatim endif -!!$ser verbatim if (present(cy)) then -!!$ser data cy=cy -!!$ser verbatim endif -!!$ser data v=v -!!$ser verbatim if (present(mfx)) then -!!$ser data mfx_=mfx -!!$ser verbatim endif -!!$ser verbatim if (present(cx)) then -!!$ser data cx_=cx -!!$ser verbatim endif -!!$ser verbatim endif 1000 continue ! Update pressure variables and get new pkz, T_v, and omega @@ -1639,75 +1201,6 @@ contains !$ser data pk=pk !$ser verbatim endif -!!$ser verbatim if(j == js) then -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie -!!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!!$ser verbatim pn1_3d(i,j,k) = pn1(i,k) -!!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km -!!$ser verbatim do i = is, ie -!!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do i = is, ie+1 -!!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe1_=pe1_3d pe2_=pe2_3d -!!$ser data pe_=pe_3d -!!$ser data peln_3d=peln_3d -!!$ser data pk=pk -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data dp2_3d=dp2_3d -!!$ser data pn1_3d=pn1_3d -!!$ser data pn2_3d=pn2_3d -!!$ser data pk2_3d=pk2_3d - -!!$ser data qvapor=q(:,:,:,1) -!!$ser data qliquid=q(:,:,:,2) -!!$ser data qice=q(:,:,:,3) -!!$ser data qrain=q(:,:,:,4) -!!$ser data qsnow=q(:,:,:,5) -!!$ser data qgraupel=q(:,:,:,6) -!!$ser data qcld=q(:,:,:,7) -!!$ser data qo3mr=q(:,:,:,8) -!!$ser data qsgs_tke=q(:,:,:,9) - -!!$ser data w=w - -!!$ser data pe0_=pe0_3d -!!$ser data pe3_=pe3_3d -!!$ser data u=u -!!$ser verbatim if (present(mfy)) then -!!$ser data mfy=mfy -!!$ser verbatim endif -!!$ser verbatim if (present(cy)) then -!!$ser data cy=cy -!!$ser verbatim endif -!!$ser data v=v -!!$ser verbatim if (present(mfx)) then -!!$ser data mfx_=mfx -!!$ser verbatim endif -!!$ser verbatim if (present(cx)) then -!!$ser data cx_=cx -!!$ser verbatim endif -!!$ser verbatim endif - !--------------------- ! Compute pkz and T_v !--------------------- @@ -1958,12 +1451,12 @@ contains !$OMP do_adiabatic_init,zsum1,zsum0,te0_2d,domain, & !$OMP ng,gridstruct,E_Flux,pdt,dtmp,reproduce_sum,q, & !$OMP mdt,cld_amt,cappa,dtdt,out_dt,rrg,akap,do_sat_adj, & -!$OMP fast_mp_consv,kord_tm, & +!$OMP fast_mp_consv,kord_tm, phis_3d, te_2d_f32, & #ifdef SERIALIZE !$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, serial_flag & #endif !$OMP ) & -!$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp) +!$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp, rsin2, cosa_s) dtmp = 0. if( last_step .and. (.not.do_adiabatic_init) ) then @@ -2003,6 +1496,35 @@ contains phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) enddo enddo + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim phis_3d(i,j,k) = phis(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim rsin2=gridstruct%rsin2 + !$ser verbatim cosa_s=gridstruct%cosa_s + !$ser verbatim do i = is, ie + !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) + !$ser verbatim enddo + !!$ser verbatim print*,'MoistCVPlusTe_2D serialization' + !$ser savepoint MoistCVPlusTe_2d-In + !$ser data qvapor_js=q(:,j,:,sphum) + !$ser data qliquid_js=q(:,j,:,liq_wat) + !$ser data qice_js=q(:,j,:,ice_wat) + !$ser data qrain_js=q(:,j,:,rainwat) + !$ser data qsnow_js=q(:,j,:,snowwat) + !$ser data qgraupel_js=q(:,j,:,graupel) + !$ser data delp=delp + !$ser data pt=pt + !$ser data phis_=phis_3d + !$ser data te_2d_=te_2d_f32 + !$ser data u=u + !$ser data v=v + !$ser data w=w + !$ser data cosa_s=cosa_s + !$ser data rsin2=rsin2 + !$ser verbatim endif do k=1,km #ifdef MOIST_CAPPA ! if(j==js) print*,"MOIST_CAPPA EXECUTED" @@ -2023,6 +1545,13 @@ contains enddo #endif enddo ! k-loop + !$ser verbatim if(j == js) then + !$ser verbatim do i = is, ie + !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) + !$ser verbatim enddo + !$ser savepoint MoistCVPlusTe_2d-Out + !$ser data te_2d_=te_2d_f32 + !$ser verbatim endif endif ! end non-hydro do i=is,ie From e930fae7c60687fd3962e17e28fce642a900b785 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 13 Jan 2025 11:03:41 -0800 Subject: [PATCH 104/252] Added computations before moist_cv_te calcuation --- pyFV3/stencils/moist_cv.py | 12 +++++++++++- .../translate/translate_moistcvpluste_2d.py | 14 ++++++++++++-- 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index f318d603..468b324c 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -3,6 +3,7 @@ __INLINED, PARALLEL, FORWARD, + BACKWARD, computation, exp, interval, @@ -184,6 +185,9 @@ def moist_te( delp: FloatField, rsin2: FloatFieldIJ, cosa_s: FloatFieldIJ, + hs: FloatFieldIJ, + delz: FloatField, + grav: Float ): """ Args: @@ -202,8 +206,14 @@ def moist_te( delp (in): rsin2 (in): cosa_s (in): + hs (in): """ - with computation(FORWARD), interval(...): + with computation(FORWARD), interval(-1,None): + te = 0.0 + phis = hs + with computation(BACKWARD), interval(0,-1): + phis = phis[0,0,1] - grav*delz + with computation(FORWARD), interval(0,-1): cvm, gz = moist_cv_nwat6_fn( qvapor, qliquid, qrain, qsnow, qice, qgraupel ) diff --git a/tests/savepoint/translate/translate_moistcvpluste_2d.py b/tests/savepoint/translate/translate_moistcvpluste_2d.py index 6847f836..bd6ed568 100644 --- a/tests/savepoint/translate/translate_moistcvpluste_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluste_2d.py @@ -56,13 +56,20 @@ def __init__(self, grid, namelist, stencil_factory): "jstart": grid.jsd, "jend": grid.jed, }, + "hs": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, + "delz": {} } self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel"] for k, v in self.in_vars["data_vars"].items(): # if k not in self.write_vars: if k in self.write_vars: v["axis"] = 1 - + self.in_vars["parameters"] = ["grav"] self.out_vars = { "te_2d_": { "istart": grid.is_, @@ -75,7 +82,7 @@ def __init__(self, grid, namelist, stencil_factory): self.compute_func = stencil_factory.from_origin_domain( moist_cv.moist_te, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz), + domain=(grid.nic, 1, grid.npz+1), ) def compute_from_storage(self, inputs): @@ -102,6 +109,9 @@ def compute_from_storage(self, inputs): inputs["delp"], inputs["rsin2"], inputs["cosa_s"], + inputs["hs"], + inputs["delz"], + inputs["grav"], ) return inputs From 8d19bee90f0f7481c411d4c99adbc012127c9bf8 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 13 Jan 2025 13:31:20 -0800 Subject: [PATCH 105/252] Updated fv_mapz serialization file --- fv_mapz.F90.SER | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 4c91154d..64a936df 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -234,6 +234,8 @@ contains logical :: serial_flag + real :: grav_ + !$ser verbatim real :: w_max_, w_min_ !$ser verbatim integer:: mode, abskord, iep1, iedp1, jedp1, js2d, o3mr, sgs_tke @@ -1456,7 +1458,7 @@ contains !$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, serial_flag & #endif !$OMP ) & -!$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp, rsin2, cosa_s) +!$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp, rsin2, cosa_s, grav_) dtmp = 0. if( last_step .and. (.not.do_adiabatic_init) ) then @@ -1487,15 +1489,6 @@ contains else ! TE using 3D winds (pt is virtual temperature): ! if(j==js) print*,"CODE EXECUTED 1" - do i=is,ie - te_2d(i,j) = 0. - phis(i,km+1) = hs(i,j) - enddo - do k=km,1,-1 - do i=is,ie - phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) - enddo - enddo !$ser verbatim if(j == js) then !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie @@ -1507,6 +1500,7 @@ contains !$ser verbatim do i = is, ie !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) !$ser verbatim enddo + !$ser verbatim grav_=grav !!$ser verbatim print*,'MoistCVPlusTe_2D serialization' !$ser savepoint MoistCVPlusTe_2d-In !$ser data qvapor_js=q(:,j,:,sphum) @@ -1524,7 +1518,19 @@ contains !$ser data w=w !$ser data cosa_s=cosa_s !$ser data rsin2=rsin2 + !$ser data hs=hs + !$ser data grav=grav_ + !$ser data delz=delz !$ser verbatim endif + do i=is,ie + te_2d(i,j) = 0. + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) + enddo + enddo do k=1,km #ifdef MOIST_CAPPA ! if(j==js) print*,"MOIST_CAPPA EXECUTED" From 3143a07726f1c24cb91c10fa05cb2a1df87f535c Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 13 Jan 2025 21:24:20 -0800 Subject: [PATCH 106/252] Correction to moist_cv_te translate test --- tests/savepoint/translate/translate_moistcvpluste_2d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/savepoint/translate/translate_moistcvpluste_2d.py b/tests/savepoint/translate/translate_moistcvpluste_2d.py index bd6ed568..57c39805 100644 --- a/tests/savepoint/translate/translate_moistcvpluste_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluste_2d.py @@ -29,9 +29,9 @@ def __init__(self, grid, namelist, stencil_factory): }, "u": { "istart": grid.isd, - "iend": grid.ied + 1, + "iend": grid.ied , "jstart": grid.jsd, - "jend": grid.jed, + "jend": grid.jed+1, "kend": grid.npz, }, "v": { From a898a1ebc6690aa543bb716885ba59cb285d6644 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 14 Jan 2025 09:56:56 -0800 Subject: [PATCH 107/252] Added te and zsum1 computation stencil and translate test --- pyFV3/stencils/moist_cv.py | 13 ++++ tests/savepoint/translate/__init__.py | 3 +- .../savepoint/translate/translate_te_zsum.py | 68 +++++++++++++++++++ 3 files changed, 83 insertions(+), 1 deletion(-) create mode 100644 tests/savepoint/translate/translate_te_zsum.py diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index 468b324c..584a03c8 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -223,6 +223,19 @@ def moist_te( + v**2.0 + v[1,0,0]**2.0 - (u + u[0,1,0]) * (v + v[1,0,0]) * cosa_s))) +def te_zsum(te_2d: FloatFieldIJ, + te0_2d: FloatFieldIJ, + delp: FloatField, + pkz: FloatField, + zsum1: FloatFieldIJ, + ): + with computation(FORWARD): + with interval(0,1): + te_2d = te0_2d - te_2d + zsum1 = pkz * delp + + with interval(1,None): + zsum1 = zsum1 + pkz * delp def fv_setup( qvapor: FloatField, diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 589a1939..cb68f718 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -100,4 +100,5 @@ from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln from .translate_getMPIprop import TranslateGetMPIProp from .translate_mpp_global_sum import TranslateMpp_global_sum -from .translate_moistcvpluste_2d import TranslateMoistCVPlusTe_2d \ No newline at end of file +from .translate_moistcvpluste_2d import TranslateMoistCVPlusTe_2d +from .translate_te_zsum import TranslateTe_Zsum \ No newline at end of file diff --git a/tests/savepoint/translate/translate_te_zsum.py b/tests/savepoint/translate/translate_te_zsum.py new file mode 100644 index 00000000..6ff0e028 --- /dev/null +++ b/tests/savepoint/translate/translate_te_zsum.py @@ -0,0 +1,68 @@ +from ndsl.stencils.testing import TranslateFortranData2Py +from pyFV3.stencils import moist_cv + +class TranslateTe_Zsum(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.in_vars["data_vars"] = { + "delp": { + "kend": grid.npz, + }, + "te_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + "te0_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + "zsum1": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + "pkz": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz, + }, + } + self.out_vars = { + "te_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + "zsum1": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + } + + self.compute_func = stencil_factory.from_origin_domain( + moist_cv.te_zsum, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + + def compute_from_storage(self, inputs): + + self.compute_func(inputs["te_2d_"], + inputs["te0_2d_"], + inputs["delp"], + inputs["pkz"], + inputs["zsum1"], + ) + + return inputs From 036ec4c2742d88af3b2f45ec97b098469ea31b61 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 14 Jan 2025 13:26:06 -0800 Subject: [PATCH 108/252] Updated mpp_global_sum to use updated all_reduce function in NDSL --- pyFV3/stencils/mpp_global_sum.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py index 4688f823..01b8d3d9 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/stencils/mpp_global_sum.py @@ -1,4 +1,5 @@ from ndsl.quantity import Quantity +from ndsl.comm.comm_abc import ReductionOperator import numpy as np def mpp_global_sum(inputArray, communicator, stencil_factory=None): @@ -36,7 +37,7 @@ def mpp_global_sum(inputArray, communicator, stencil_factory=None): carry_overflow(ints_sum.data, prec, I_prec, prec_error) # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' after carry_over') - communicator.all_reduce_sum(ints_sum, ints_sum_reduce) + communicator.all_reduce(ints_sum, ReductionOperator.SUM, ints_sum_reduce) # print("rank ", communicator.rank, "sum(ints_sum_reduce) = ", sum(ints_sum_reduce.data), ' after all_reduce') regularize_ints(ints_sum_reduce.data, prec, I_prec) From 33b25207bc88011468b2947846900e55298fcd5f Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 14 Jan 2025 13:27:10 -0800 Subject: [PATCH 109/252] Added moist_cv_pt_last_step stencil and translate test --- tests/savepoint/translate/__init__.py | 3 +- .../translate_moistcvpluspt_2d_last_step.py | 161 ++++++++++++++++++ 2 files changed, 163 insertions(+), 1 deletion(-) create mode 100644 tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index cb68f718..9982b250 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -101,4 +101,5 @@ from .translate_getMPIprop import TranslateGetMPIProp from .translate_mpp_global_sum import TranslateMpp_global_sum from .translate_moistcvpluste_2d import TranslateMoistCVPlusTe_2d -from .translate_te_zsum import TranslateTe_Zsum \ No newline at end of file +from .translate_te_zsum import TranslateTe_Zsum +from .translate_moistcvpluspt_2d_last_step import TranslateMoistCVPlusPt_2d_last_step \ No newline at end of file diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py new file mode 100644 index 00000000..851d9e08 --- /dev/null +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py @@ -0,0 +1,161 @@ +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + +from ndsl import StencilFactory +from ndsl.dsl.typing import FloatField, Float +from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j +from pyFV3.stencils import moist_cv + + +def moist_pt( + qvapor: FloatField, + qliquid: FloatField, + qrain: FloatField, + qsnow: FloatField, + qice: FloatField, + qgraupel: FloatField, + q_con: FloatField, + pt: FloatField, + cappa: FloatField, + delp: FloatField, + delz: FloatField, + r_vir: Float, +): + with computation(PARALLEL), interval(...): + cvm, gz, q_con, cappa, pt = moist_cv.moist_pt_func( + qvapor, + qliquid, + qrain, + qsnow, + qice, + qgraupel, + q_con, + pt, + cappa, + delp, + delz, + r_vir, + ) + + +class MoistPT: + """ + Class to test with DaCe orchestration. test class is MoistCVPlusPt_2d + """ + + def __init__( + self, + stencil_factory: StencilFactory, + grid, + ): + self._moist_cv_pt = stencil_factory.from_origin_domain( + moist_pt, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + + def __call__( + self, + qvapor: FloatField, + qliquid: FloatField, + qrain: FloatField, + qsnow: FloatField, + qice: FloatField, + qgraupel: FloatField, + q_con: FloatField, + pt: FloatField, + cappa: FloatField, + delp: FloatField, + delz: FloatField, + r_vir: float, + ): + self._moist_cv_pt( + qvapor, + qliquid, + qrain, + qsnow, + qice, + qgraupel, + q_con, + pt, + cappa, + delp, + delz, + r_vir, + ) + + +class TranslateMoistCVPlusPt_2d_last_step(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.compute_func = MoistPT(stencil_factory, self.grid) # type: ignore + self.in_vars["data_vars"] = { + "qvapor": { + "kend": grid.npz-1, + }, + "qliquid": { + "kend": grid.npz-1, + }, + "qice": { + "kend": grid.npz-1, + }, + "qrain": { + "kend": grid.npz-1, + }, + "qsnow": { + "kend": grid.npz-1, + }, + "qgraupel": { + "kend": grid.npz-1, + }, + "pt": {}, + "pkz": {}, + } + # self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] + # for k, v in self.in_vars["data_vars"].items(): + # # if k not in self.write_vars: + # if k in self.write_vars: + # v["axis"] = 1 + + self.in_vars["parameters"] = ["r_vir", "dtmp"] + self.out_vars = { + "pt": {}, + } + + self.compute_func = stencil_factory.from_origin_domain( + moist_cv.moist_pt_last_step, + origin=grid.compute_origin(), + domain=(grid.nic, grid.njc, grid.npz), + ) + + self.quantity_factory = grid.quantity_factory + + self._gz = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz, + ), dtype=Float, + ) + + def compute_from_storage(self, inputs): + # for name, value in inputs.items(): + # if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: + # inputs[name] = self.make_storage_data( + # pad_field_in_j( + # value, self.grid.njd, backend=self.stencil_factory.backend + # ) + # ) + self.compute_func(inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + self._gz, + inputs["pt"], + inputs["pkz"], + inputs["dtmp"], + inputs["r_vir"], + ) + return inputs From 3c4a5643a5c7a24019428cea9680680dd4af805d Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 14 Jan 2025 13:32:44 -0800 Subject: [PATCH 110/252] Remove extraneous code from moist_cv_pt_last_step translate test --- .../translate_moistcvpluspt_2d_last_step.py | 93 +------------------ 1 file changed, 1 insertion(+), 92 deletions(-) diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py index 851d9e08..ec1728ee 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py @@ -5,90 +5,10 @@ from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j from pyFV3.stencils import moist_cv - -def moist_pt( - qvapor: FloatField, - qliquid: FloatField, - qrain: FloatField, - qsnow: FloatField, - qice: FloatField, - qgraupel: FloatField, - q_con: FloatField, - pt: FloatField, - cappa: FloatField, - delp: FloatField, - delz: FloatField, - r_vir: Float, -): - with computation(PARALLEL), interval(...): - cvm, gz, q_con, cappa, pt = moist_cv.moist_pt_func( - qvapor, - qliquid, - qrain, - qsnow, - qice, - qgraupel, - q_con, - pt, - cappa, - delp, - delz, - r_vir, - ) - - -class MoistPT: - """ - Class to test with DaCe orchestration. test class is MoistCVPlusPt_2d - """ - - def __init__( - self, - stencil_factory: StencilFactory, - grid, - ): - self._moist_cv_pt = stencil_factory.from_origin_domain( - moist_pt, - origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz), - ) - - def __call__( - self, - qvapor: FloatField, - qliquid: FloatField, - qrain: FloatField, - qsnow: FloatField, - qice: FloatField, - qgraupel: FloatField, - q_con: FloatField, - pt: FloatField, - cappa: FloatField, - delp: FloatField, - delz: FloatField, - r_vir: float, - ): - self._moist_cv_pt( - qvapor, - qliquid, - qrain, - qsnow, - qice, - qgraupel, - q_con, - pt, - cappa, - delp, - delz, - r_vir, - ) - - class TranslateMoistCVPlusPt_2d_last_step(TranslateFortranData2Py): def __init__(self, grid, namelist, stencil_factory): super().__init__(grid, namelist, stencil_factory) self.stencil_factory = stencil_factory - self.compute_func = MoistPT(stencil_factory, self.grid) # type: ignore self.in_vars["data_vars"] = { "qvapor": { "kend": grid.npz-1, @@ -111,11 +31,6 @@ def __init__(self, grid, namelist, stencil_factory): "pt": {}, "pkz": {}, } - # self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] - # for k, v in self.in_vars["data_vars"].items(): - # # if k not in self.write_vars: - # if k in self.write_vars: - # v["axis"] = 1 self.in_vars["parameters"] = ["r_vir", "dtmp"] self.out_vars = { @@ -139,13 +54,7 @@ def __init__(self, grid, namelist, stencil_factory): ) def compute_from_storage(self, inputs): - # for name, value in inputs.items(): - # if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: - # inputs[name] = self.make_storage_data( - # pad_field_in_j( - # value, self.grid.njd, backend=self.stencil_factory.backend - # ) - # ) + self.compute_func(inputs["qvapor"], inputs["qliquid"], inputs["qrain"], From 52a1b1dcd88b73a14fae05e52ae33eeb97264493 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 14 Jan 2025 14:00:16 -0800 Subject: [PATCH 111/252] Added cond_output stencil and corresponding translate test --- pyFV3/stencils/moist_cv.py | 11 ++++ tests/savepoint/translate/__init__.py | 3 +- .../translate/translate_cond_output.py | 54 +++++++++++++++++++ 3 files changed, 67 insertions(+), 1 deletion(-) create mode 100644 tests/savepoint/translate/translate_cond_output.py diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index 584a03c8..f975009d 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -237,6 +237,17 @@ def te_zsum(te_2d: FloatFieldIJ, with interval(1,None): zsum1 = zsum1 + pkz * delp +def cond_output(q_con: FloatField, + qliquid: FloatField, + qrain: FloatField, + qsnow: FloatField, + qice: FloatField, + qgraupel: FloatField, + ): + with computation(PARALLEL), interval(...): + q_con = 0.0 + q_con = qliquid + qice + qrain + qsnow + qgraupel + def fv_setup( qvapor: FloatField, qliquid: FloatField, diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 9982b250..e3a4fc6c 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -102,4 +102,5 @@ from .translate_mpp_global_sum import TranslateMpp_global_sum from .translate_moistcvpluste_2d import TranslateMoistCVPlusTe_2d from .translate_te_zsum import TranslateTe_Zsum -from .translate_moistcvpluspt_2d_last_step import TranslateMoistCVPlusPt_2d_last_step \ No newline at end of file +from .translate_moistcvpluspt_2d_last_step import TranslateMoistCVPlusPt_2d_last_step +from .translate_cond_output import TranslateCond_output \ No newline at end of file diff --git a/tests/savepoint/translate/translate_cond_output.py b/tests/savepoint/translate/translate_cond_output.py new file mode 100644 index 00000000..0dbbe71b --- /dev/null +++ b/tests/savepoint/translate/translate_cond_output.py @@ -0,0 +1,54 @@ +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + +from ndsl import StencilFactory +from ndsl.dsl.typing import FloatField, Float +from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j +from pyFV3.stencils import moist_cv + +class TranslateCond_output(TranslateFortranData2Py): + def __init__(self, grid, namelist, stencil_factory): + super().__init__(grid, namelist, stencil_factory) + self.stencil_factory = stencil_factory + self.in_vars["data_vars"] = { + "qliquid": { + "kend": grid.npz-1, + }, + "qice": { + "kend": grid.npz-1, + }, + "qrain": { + "kend": grid.npz-1, + }, + "qsnow": { + "kend": grid.npz-1, + }, + "qgraupel": { + "kend": grid.npz-1, + }, + "q_con": { + "kend": grid.npz-1, + } + } + + self.out_vars = { + "q_con": { + "kend": grid.npz-1, + } + } + + self.compute_func = stencil_factory.from_origin_domain( + moist_cv.cond_output, + origin=grid.compute_origin(), + domain=(grid.nic, grid.njc, grid.npz), + ) + + def compute_from_storage(self, inputs): + + self.compute_func(inputs["q_con"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + ) + return inputs From 5b27f76e98273aca04d044b978581b9d07a983e4 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 15 Jan 2025 09:01:08 -0800 Subject: [PATCH 112/252] Updated remapping translate test with stencils needed for fv_mapz computation. Note that since this isn't a parallel translate test, the mpp_global_sum function won't work, so it's commented out --- fv_mapz.F90.SER | 225 ++++++++++++------ .../translate/translate_remapping_GEOS.py | 137 ++++++++++- 2 files changed, 275 insertions(+), 87 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 64a936df..ddd4b1c7 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -229,8 +229,8 @@ contains real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d - real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32 - real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s + real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32, te0_2d_f32, zsum1_f32 + real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64 logical :: serial_flag @@ -376,9 +376,9 @@ contains !$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & !$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,& !$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & -!$OMP peln_3d, gz_2d, pe0_3d, pe3_3d) & +!$OMP peln_3d, gz_2d, pe0_3d, pe3_3d, phis_3d, te_2d, te_2d_f32, te0_2d, te0_2d_f32) & !$OMP private(gz,cvm,bkh,dp2, & -!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_) +!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64) do 1000 j=js,je+1 !$ser verbatim if(j == js) then @@ -410,6 +410,21 @@ contains !$ser verbatim qmin=t_min !$ser verbatim w_max_ = w_max !$ser verbatim w_min_ = w_min +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim phis_3d(i,j,k) = phis(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim rsin2=gridstruct%rsin2 +!$ser verbatim cosa_s=gridstruct%cosa_s +!!$ser verbatim do i = is, ie +!!$ser verbatim te_2d_f32(i,:) = te_2d(i,:) +!!$ser verbatim area_64(i,:) = gridstruct%area_64(i,:) +!!$ser verbatim enddo +!$ser verbatim grav_=grav +!$ser verbatim te_2d_f32 = real(te_2d) +!$ser verbatim area_64 = real(gridstruct%area_64) +!$ser verbatim te0_2d_f32 = real(te0_2d) !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d !$ser data pe1_=pe1_3d @@ -465,6 +480,14 @@ contains !$ser data w_max=w_max_ !$ser data w_min=w_min_ !$ser data kord_mt=kord_mt +!$ser data cosa_s=cosa_s +!$ser data rsin2=rsin2 +!$ser data hs=hs +!$ser data grav=grav_ +!$ser data te_2d_=te_2d_f32 +!$ser data te0_2d_=te0_2d_f32 +!$ser data phis_=phis_3d +!$ser data area_64=area_64 !!$ser data gz=gz !!$ser data k1k=k1k !!$ser data remap_t=remap_t @@ -1323,76 +1346,6 @@ contains !$ser data pkz=pkz !$ser data cappa=cappa !$ser verbatim endif - - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) - !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie+1 - !$ser verbatim pe0_3d(i,j,k) = pe0(i,k) - !$ser verbatim pe3_3d(i,j,k) = pe3(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data pe_=pe_3d - !$ser data peln_3d=peln_3d - !$ser data pk=pk - !$ser data pkz=pkz - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d - - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) - - !$ser data w=w - - !$ser data pe0_=pe0_3d - !$ser data pe3_=pe3_3d - !$ser data u=u - !$ser verbatim if (present(mfy)) then - !$ser data mfy=mfy - !$ser verbatim endif - !$ser verbatim if (present(cy)) then - !$ser data cy=cy - !$ser verbatim endif - !$ser data v=v - !$ser verbatim if (present(mfx)) then - !$ser data mfx_=mfx - !$ser verbatim endif - !$ser verbatim if (present(cx)) then - !$ser data cx_=cx - !$ser verbatim endif - !$ser verbatim endif endif if (remap_pt) then print*,'CODE NOT TESTED HERE 10' @@ -1453,11 +1406,12 @@ contains !$OMP do_adiabatic_init,zsum1,zsum0,te0_2d,domain, & !$OMP ng,gridstruct,E_Flux,pdt,dtmp,reproduce_sum,q, & !$OMP mdt,cld_amt,cappa,dtdt,out_dt,rrg,akap,do_sat_adj, & -!$OMP fast_mp_consv,kord_tm, phis_3d, te_2d_f32, & +!$OMP fast_mp_consv,kord_tm, phis_3d, te_2d_f32, te0_2d_f32, zsum1_f32, & #ifdef SERIALIZE -!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, serial_flag & +!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, serial_flag, & #endif -!$OMP ) & +!$OMP js2d, pe1_3d, pe2_3d, pn1, pn1_3d, pn2, pn2_3d, pk2, pk2_3d, pe_3d, peln_3d, dp2, dp2_3d, & +!$OMP pe0_3d, pe3_3d, ps, mfy, cy, mfx, cx) & !$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp, rsin2, cosa_s, grav_) dtmp = 0. @@ -1558,8 +1512,94 @@ contains !$ser savepoint MoistCVPlusTe_2d-Out !$ser data te_2d_=te_2d_f32 !$ser verbatim endif + + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) + !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) + !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie+1 + !$ser verbatim pe0_3d(i,j,k) = pe0(i,k) + !$ser verbatim pe3_3d(i,j,k) = pe3(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do i = is, ie + !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) + !$ser verbatim enddo + !$ser savepoint Remapping_GEOS-Out + !$ser data pe1_=pe1_3d pe2_=pe2_3d + !$ser data pe_=pe_3d + !$ser data peln_3d=peln_3d + !$ser data pk=pk + !$ser data pkz=pkz + !$ser data q_con=q_con + !$ser data pt=pt + !$ser data cappa=cappa + !$ser data delp=delp + !$ser data delz=delz + !$ser data ps=ps + !$ser data dp2_3d=dp2_3d + !$ser data pn1_3d=pn1_3d + !$ser data pn2_3d=pn2_3d + !$ser data pk2_3d=pk2_3d + + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + + !$ser data w=w + + !$ser data pe0_=pe0_3d + !$ser data pe3_=pe3_3d + !$ser data u=u + !$ser verbatim if (present(mfy)) then + !$ser data mfy=mfy + !$ser verbatim endif + !$ser verbatim if (present(cy)) then + !$ser data cy=cy + !$ser verbatim endif + !$ser data v=v + !$ser verbatim if (present(mfx)) then + !$ser data mfx_=mfx + !$ser verbatim endif + !$ser verbatim if (present(cx)) then + !$ser data cx_=cx + !$ser verbatim endif + !$ser data te_2d_=te_2d_f32 + !$ser verbatim endif + endif ! end non-hydro + !$ser verbatim if(j == js2d) then + !$ser verbatim te_2d_f32 = real(te_2d) + !$ser verbatim te0_2d_f32 = real(te0_2d) + !$ser verbatim zsum1_f32 = real(zsum1) + !$ser savepoint Te_Zsum-In + !$ser data te_2d_=te_2d_f32 + !$ser data te0_2d_=te0_2d_f32 + !$ser data zsum1=zsum1_f32 + !$ser data delp=delp + !$ser data pkz=pkz + !$ser verbatim endif do i=is,ie te_2d(i,j) = te0_2d(i,j) - te_2d(i,j) zsum1(i,j) = pkz(i,j,1)*delp(i,j,1) @@ -1574,6 +1614,14 @@ contains zsum0(i,j) = ptop*(pk(i,j,1)-pk(i,j,km+1)) + zsum1(i,j) enddo endif + !$ser verbatim if(j == js2d) then + !$ser verbatim te_2d_f32 = real(te_2d) + !$ser verbatim zsum1_f32 = real(zsum1) + !!$ser verbatim print*, 'sum te_2d: ', sum(te_2d_f32(:,1)), sum(te_2d(:,1)) + !$ser savepoint Te_Zsum-Out + !$ser data te_2d_=te_2d_f32 + !$ser data zsum1=zsum1_f32 + !$ser verbatim endif enddo ! j-loop @@ -1701,6 +1749,17 @@ contains ! Output temperature if last_step if ( .not. hydrostatic ) then ! print*,'CODE EXECUTED' + !$ser savepoint MoistCVPlusPt_2d_last_step-In + !$ser data qvapor=q(:,:,:,sphum) + !$ser data qliquid=q(:,:,:,liq_wat) + !$ser data qice=q(:,:,:,ice_wat) + !$ser data qrain=q(:,:,:,rainwat) + !$ser data qsnow=q(:,:,:,snowwat) + !$ser data qgraupel=q(:,:,:,graupel) + !$ser data r_vir=r_vir + !$ser data dtmp=dtmp + !$ser data pt=pt + !$ser data pkz=pkz !$OMP do do k=1,km do j=js,je @@ -1717,8 +1776,17 @@ contains #endif enddo ! j-loop enddo ! k-loop + !$ser savepoint MoistCVPlusPt_2d_last_step-Out + !$ser data pt=pt #ifdef USE_COND ! print*, "USE_COND active" + !$ser savepoint Cond_output-In + !$ser data q_con=q_con + !$ser data qliquid=q(:,:,:,liq_wat) + !$ser data qice=q(:,:,:,ice_wat) + !$ser data qrain=q(:,:,:,rainwat) + !$ser data qsnow=q(:,:,:,snowwat) + !$ser data qgraupel=q(:,:,:,graupel) ! Fill condensate output !$OMP do do k=1,km @@ -1734,6 +1802,9 @@ contains enddo ! j-loop enddo ! k-loop #endif + !$ser savepoint Cond_output-Out + !$ser data q_con=q_con + else print*,'CODE NOT TESTED HERE 16' !$OMP do diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 31776a2e..40b34799 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -7,6 +7,7 @@ from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax, pressures_mapv, pe_pk_delp_peln +from pyFV3.stencils.mpp_global_sum import mpp_global_sum from ndsl.stencils.testing import pad_field_in_j, Grid from pyFV3.testing import TranslateDycoreFortranData2Py from ndsl.constants import ( @@ -200,14 +201,14 @@ def __init__( "iend": grid.ied, "jstart": grid.jsd, "jend": grid.jed+1, - "kend": grid.npz-1, + "kend": grid.npz, }, "v": { "istart": grid.isd, "iend": grid.ied+1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz, }, "mfy": { "istart": grid.is_, @@ -239,6 +240,49 @@ def __init__( "jend": grid.jed, "kend": grid.npz-1, }, + "phis_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "cosa_s": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, + "rsin2": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, + "te_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + "hs": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, + "te0_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + "area_64": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, } # self.write_vars = ["gz", "cvm"] self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] @@ -256,6 +300,7 @@ def __init__( "w_max", "w_min", "kord_mt", + "grav", # "zvir", # "last_step", # "consv_te", @@ -431,6 +476,12 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, + "te_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, } self.stencil_factory = stencil_factory @@ -479,6 +530,13 @@ def __init__( ), dtype=Float, ) + self._zsum1 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + ), dtype=Float, + ) + self._compute_performed = self.quantity_factory._numpy.zeros( ( grid.nid, @@ -568,6 +626,24 @@ def __init__( domain=(grid.nic, 1, grid.npz), ) + self._moist_cv_te = stencil_factory.from_origin_domain( + moist_cv.moist_te, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz+1), + ) + + self._te_zsum = stencil_factory.from_origin_domain( + moist_cv.te_zsum, + origin=grid.compute_origin(), + domain=(grid.nic, 1, grid.npz), + ) + + self._most_cv_pt_last_step = stencil_factory.from_origin_domain( + moist_cv.moist_pt_last_step, + origin=grid.compute_origin(), + domain=(grid.nic, grid.njc, grid.npz), + ) + def compute_from_storage(self, inputs): # Replicates tracer values in I along the J direction @@ -795,16 +871,57 @@ def compute_from_storage(self, inputs): Float(inputs["r_vir"]), ) - # If loop based on if( last_step .and. (.not.do_adiabatic_init) ) then - # PHIS computation - # Some variation of moist_cv_pt - # zsum1 computation + # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then + + self._moist_cv_te(inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + inputs["u"], + inputs["v"], + inputs["w"], + inputs["te_2d_"], + inputs["pt"], + inputs["phis_"], + inputs["delp"], + inputs["rsin2"], + inputs["cosa_s"], + inputs["hs"], + inputs["delz"], + inputs["grav"], + ) + + self._te_zsum(inputs["te_2d_"], + inputs["te0_2d_"], + inputs["delp"], + inputs["pkz"], + self._zsum1, + ) - # MPP GLOBAL SUM - # E_FLUX calcuation + # Note, since this is a serial translate test, mpp_global_sum won't work without the communicator. + # Also, mpp_global_sum is currently set up for the C24 TBC setup + + # tesum = mpp_global_sum(inputs["te_2d_"]*inputs["area_64"], communicator, self.stencil_factory) + + # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed + + # zsum = mpp_global_sum(self._zsum1*inputs["area_64"], communicator, self.stencil_factory) + # dtmp = tesum / (cv_air*zsum) # If loop based on if ( last_step .and. (.not. adiabatic) ) then - # Some variation of moist_cv_pt - # Condensation update + # self._most_cv_pt_last_step(inputs["qvapor"], + # inputs["qliquid"], + # inputs["qrain"], + # inputs["qsnow"], + # inputs["qice"], + # inputs["qgraupel"], + # self._gz, + # inputs["pt"], + # inputs["pkz"], + # dtmp, + # inputs["r_vir], + # ) return inputs From 41435609df8e649149d237065dc09b2fd86d523f Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 15 Jan 2025 16:48:45 -0500 Subject: [PATCH 113/252] Specify correct dtype for `dpx` --- pyFV3/stencils/dyn_core.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 29d3591f..8a693e36 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -239,6 +239,13 @@ def dyncore_temporaries( units="unknown", dtype=Float, ) + for name in ["dpx"]: + temporaries[name] = quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=np.float64, + allow_mismatch_float_precision=True, + ) return temporaries @@ -448,8 +455,7 @@ def __init__( ) if config.beta < -0.1: raise RuntimeError( - "Acoustics (dyn_core): beta < 0.1 is not implemented" - " (one_grad_p, etc.)" + "Acoustics (dyn_core): beta < 0.1 is not implemented (one_grad_p, etc.)" ) if config.use_logp: raise RuntimeError("Acoustics (dyn_core): use_logp=True is not implemented") From 97cf917e15232e3dd4c479f577035aa2166b18a4 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 16 Jan 2025 08:28:11 -0800 Subject: [PATCH 114/252] Starting translate test with full j-axis being computed --- fv_mapz.F90.SER.v2 | 4828 +++++++++++++++++ .../translate/translate_remapping_GEOS.py | 976 ++-- 2 files changed, 5327 insertions(+), 477 deletions(-) create mode 100644 fv_mapz.F90.SER.v2 diff --git a/fv_mapz.F90.SER.v2 b/fv_mapz.F90.SER.v2 new file mode 100644 index 00000000..85f57ff5 --- /dev/null +++ b/fv_mapz.F90.SER.v2 @@ -0,0 +1,4828 @@ +!*********************************************************************** +!* GNU Lesser General Public License +!* +!* This file is part of the FV3 dynamical core. +!* +!* The FV3 dynamical core is free software: you can redistribute it +!* and/or modify it under the terms of the +!* GNU Lesser General Public License as published by the +!* Free Software Foundation, either version 3 of the License, or +!* (at your option) any later version. +!* +!* The FV3 dynamical core is distributed in the hope that it will be +!* useful, but WITHOUT ANYWARRANTY; without even the implied warranty +!* of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. +!* See the GNU General Public License for more details. +!* +!* You should have received a copy of the GNU Lesser General Public +!* License along with the FV3 dynamical core. +!* If not, see . +!*********************************************************************** + +!>@brief The module 'fv_mapz' contains the vertical mapping routines \cite lin2004vertically +!>@note April 12, 2012 -SJL: This revision may actually produce rounding level differences +!! due to the elimination of KS to compute pressure level for remapping. + +module fv_mapz_mod + +! Modules Included: +! +! +! +! +! +!
Module NameFunctions Included
+! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +!
constants_modradius, pi=>pi_8, rvgas, rdgas, grav, hlv, hlf, cp_air, cp_vapor
field_manager_modMODEL_ATMOS
fv_arrays_modfv_grid_type
fv_cmp_modqs_init, fv_sat_adj
fv_fill_modfillz
fv_grid_utils_modg_sum, ptop_min
fv_mp_modis_master
fv_timing_modtiming_on, timing_off
fv_tracer2d_modtracer_2d, tracer_2d_1L, tracer_2d_nested
mpp_mod/td> +! NOTE, mpp_error, get_unit, mpp_root_pe, mpp_pe
mpp_domains_mod/td> +! mpp_update_domains, domain2d
tracer_manager_modget_tracer_index
+ !$ser verbatim USE m_serialize, ONLY: fs_is_serialization_on + use constants_mod, only: radius, pi=>pi_8, rvgas, rdgas, grav, hlv, hlf, hls, cp_air, cp_vapor + use tracer_manager_mod,only: get_tracer_index + use field_manager_mod, only: MODEL_ATMOS + use fv_grid_utils_mod, only: g_sum, ptop_min + use fv_fill_mod, only: fillz + use mpp_domains_mod, only: mpp_update_domains, domain2d, mpp_global_sum, BITWISE_EFP_SUM, BITWISE_EXACT_SUM + use mpp_mod, only: NOTE, mpp_error, get_unit, mpp_root_pe, mpp_pe + use fv_arrays_mod, only: fv_grid_type, fv_flags_type + use fv_timing_mod, only: timing_on, timing_off + use fv_mp_mod, only: is_master + use fv_cmp_mod, only: qs_init, fv_sat_adj + + implicit none + real, parameter:: consv_min= 0.001 !< below which no correction applies + real, parameter:: te_min= -1.e25 + real, parameter:: t_min= 184. !< below which applies stricter constraint + real, parameter:: r2=1./2., r0=0.0 + real, parameter:: r3 = 1./3., r23 = 2./3., r12 = 1./12. + real, parameter:: cv_vap = 3.*rvgas !< 1384.5 + real, parameter:: cv_air = cp_air - rdgas !< = rdgas * (7/2-1) = 2.5*rdgas=717.68 +! real, parameter:: c_ice = 2106. !< heat capacity of ice at 0.C + real, parameter:: c_ice = 1972. !< heat capacity of ice at -15.C + real, parameter:: c_liq = 4.1855e+3 !< GFS: heat capacity of water at 0C +! real, parameter:: c_liq = 4218. !< ECMWF-IFS + real, parameter:: cp_vap = cp_vapor !< 1846. + real, parameter:: tice = 273.16 + + logical, parameter :: w_limiter = .true. + real, parameter :: w_max = 90. + real, parameter :: w_min = -60. + + real(kind=8) :: E_Flux = 0. + private + + public compute_total_energy, Lagrangian_to_Eulerian, moist_cv, moist_cp, & + rst_remap, mappm, E_Flux, mapn_tracer, map1_q2 + +!---- version number ----- + character(len=128) :: version = '$Id$' + character(len=128) :: tagname = '$Name$' + +contains + +!>@brief The subroutine 'Lagrangian_to_Eulerian' remaps deformed Lagrangian layers back to the reference Eulerian coordinate. +!>@details It also includes the entry point for calling fast microphysical processes. This is typically calle on the k_split loop. + subroutine Lagrangian_to_Eulerian(last_step, consv, ps, pe, delp, pkz, pk, & + mdt, pdt, km, is,ie,js,je, isd,ied,jsd,jed, & + nq, nwat, sphum, q_con, u, v, w, delz, pt, q, hs, r_vir, cp, & + akap, cappa, kord_mt, kord_wz, kord_tr, kord_tm, peln, te0_2d, & + ng, ua, va, omga, te, ws, fill, reproduce_sum, out_dt, dtdt, & + ptop, ak, bk, pfull, flagstruct, gridstruct, domain, do_sat_adj, & + hydrostatic, hybrid_z, do_omega, adiabatic, do_adiabatic_init, & + mfx, mfy, cx, cy, remap_option, gmao_remap) + logical, intent(in):: last_step + real, intent(in):: mdt !< remap time step + real, intent(in):: pdt !< phys time step + integer, intent(in):: km + integer, intent(in):: nq !< number of tracers (including h2o) + integer, intent(in):: nwat + integer, intent(in):: sphum !< index for water vapor (specific humidity) + integer, intent(in):: ng + integer, intent(in):: is,ie,isd,ied !< starting & ending X-Dir index + integer, intent(in):: js,je,jsd,jed !< starting & ending Y-Dir index + integer, intent(in):: kord_mt !< Mapping order for the vector winds + integer, intent(in):: kord_wz !< Mapping order/option for w + integer, intent(in):: kord_tr(nq) !< Mapping order for tracers + integer, intent(in):: kord_tm !< Mapping order for thermodynamics + + real, intent(in):: consv !< factor for TE conservation + real, intent(in):: r_vir + real, intent(in):: cp + real, intent(in):: akap + real, intent(in):: hs(isd:ied,jsd:jed) !< surface geopotential + real(kind=8), intent(inout):: te0_2d(is:ie,js:je) + real, intent(in):: ws(is:ie,js:je) + + logical, intent(in):: do_sat_adj + logical, intent(in):: fill !< fill negative tracers + logical, intent(in):: reproduce_sum + logical, intent(in):: do_omega, adiabatic, do_adiabatic_init + real, intent(in) :: ptop + real, intent(in) :: ak(km+1) + real, intent(in) :: bk(km+1) + real, intent(in):: pfull(km) + type(fv_grid_type), intent(IN), target :: gridstruct + type(fv_flags_type), intent(INOUT) :: flagstruct + type(domain2d), intent(INOUT) :: domain + + ! INPUT/OUTPUT + real, intent(inout):: pk(is:ie,js:je,km+1) !< pe to the kappa + real, intent(inout):: q(isd:ied,jsd:jed,km,*) + real, intent(inout):: delp(isd:ied,jsd:jed,km) !< pressure thickness + real, intent(inout):: pe(is-1:ie+1,km+1,js-1:je+1) !< pressure at layer edges + real, intent(inout):: ps(isd:ied,jsd:jed) !< surface pressure + + ! u-wind will be ghosted one latitude to the north upon exit + real, intent(inout):: u(isd:ied ,jsd:jed+1,km) !< u-wind (m/s) + real, intent(inout):: v(isd:ied+1,jsd:jed ,km) !< v-wind (m/s) + real, intent(inout):: w(isd: ,jsd: ,1:) !< vertical velocity (m/s) + real, intent(inout):: pt(isd:ied ,jsd:jed ,km) !< cp*virtual potential temperature + !< as input; output: temperature + real, intent(inout), dimension(isd:,jsd:,1:)::delz, q_con, cappa + logical, intent(in):: hydrostatic + logical, intent(in):: hybrid_z + logical, intent(in):: out_dt + + real, intent(inout):: ua(isd:ied,jsd:jed,km) !< u-wind (m/s) on physics grid + real, intent(inout):: va(isd:ied,jsd:jed,km) !< v-wind (m/s) on physics grid + real, intent(inout):: omga(isd:ied,jsd:jed,km) !< vertical press. velocity (pascal/sec) + real, intent(inout):: peln(is:ie,km+1,js:je) !< log(pe) + real, intent(inout):: dtdt(is:ie,js:je,km) + real, intent(out):: pkz(is:ie,js:je,km) !< layer-mean pk for converting t to pt + real, intent(out):: te(isd:ied,jsd:jed,km) + ! Mass fluxes + real, optional, intent(inout):: mfx(is:ie+1,js:je ,km) ! X-dir Mass Flux + real, optional, intent(inout):: mfy(is:ie ,js:je+1,km) ! Y-dir Mass Flux + ! Courant numbers + real, optional, intent(inout):: cx(is:ie+1, jsd:jed,km) + real, optional, intent(inout):: cy(isd:ied ,js:je+1,km) + + integer, intent(in):: remap_option, gmao_remap + + ! !DESCRIPTION: + ! + ! !REVISION HISTORY: + ! SJL 03.11.04: Initial version for partial remapping + ! + !----------------------------------------------------------------------- + real(kind=8), dimension(is:ie,js:je):: te_2d, zsum0, zsum1 + real, dimension(is:ie,js:je):: dpln + real, dimension(is:ie,km) :: q2, dp2, w2 + real, dimension(is:ie,km+1):: pe1, pe2, pk1, pk2, pn1, pn2, phis + real, dimension(is:ie+1,km+1):: pe0, pe3 + real, dimension(is:ie):: gz, cvm + real(kind=8):: tesum, zsum, dtmp + real :: rcp, rg, tmp, tpe, rrg, bkh, k1k, dlnp + logical:: fast_mp_consv + integer:: i,j,k + integer:: nt, liq_wat, ice_wat, rainwat, snowwat, cld_amt, graupel, iq, n, kmp, kp, k_next + logical:: remap_t, remap_pt, remap_te + + real, dimension(is:ie, js:je, km+1) :: pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, peln_3d + real, dimension(is:ie, js:je, km+1) :: phis_3d + real, dimension(is-1:ie+1,js-1:je+1,km+1) :: pe_3d + real, dimension(is:ie+1, js:je+1, km+1) :: pe0_3d, pe3_3d + + real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d + + real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32, te0_2d_f32, zsum1_f32 + real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64 + + logical :: serial_flag + + real :: grav_ + +!$ser verbatim real :: w_max_, w_min_ + + !$ser verbatim integer:: mode, abskord, iep1, iedp1, jedp1, js2d, o3mr, sgs_tke + !$ser verbatim real :: qmin + !$ser verbatim qmin = 0.0 + !$ser verbatim iep1=ie+1 + !$ser verbatim iedp1=ied+1 + !$ser verbatim jedp1=jed+1 + !$ser verbatim js2d=js + !$ser verbatim pe_3d = 0.0 + !$ser verbatim pe1_3d = 0.0 + !$ser verbatim pe2_3d = 0.0 + !$ser verbatim peln_3d = 0.0 + !$ser verbatim dp2_3d = 0.0 + !$ser verbatim pn1_3d = 0.0 + !$ser verbatim pn2_3d = 0.0 + !$ser verbatim pk2_3d = pk + !$ser verbatim gz_2d = 0.0 + !$ser verbatim phis_3d = 0.0 + + remap_t = .false. + remap_pt = .false. + remap_te = .false. + select case (remap_option) + case(0) + remap_t = .true. + case(1) + remap_pt = .true. + case(2) + remap_te = .true. + case default + print*, ' INVALID REMAP_OPTION ' + stop + end select + + select case (gmao_remap) + case(0) + ! use GFDL schemes + case(1) + ! GMAO linear remap + case(2) + ! GMAO quadratic remap + case(3) + ! GMAO cubic remap + case default + print*, ' INVALID GMAO_REMAP' + stop + end select + + if (is_master() .and. flagstruct%fv_debug) then + print*, '' + select case (remap_option) + case(0) + print*, ' REMAPPING T in logP' + case(1) + print*, ' REMAPPING PT in P' + case(2) + print*, ' REMAPPING TE in logP' + end select + + print*, '' + select case (gmao_remap) + case(0) + print*, ' Using GFDL schemes' + case(1) + print*, ' Using GMAO linear scheme' + case(2) + print*, ' Using GMAO quadratic scheme' + case(3) + print*, ' Using GMAO cubic scheme' + end select + + ! Total eergy conservation + print*, '' + print*, ' REMAPPING CONSV: ', consv + print*, ' REMAPPING CONSV_MIN: ', consv_min + print*, '' + endif + + k1k = rdgas/cv_air ! akap / (1.-akap) = rg/Cv=0.4 + rg = rdgas + rcp = 1./ cp + rrg = -rdgas/grav + +#ifdef MAPL_MODE + select case(nwat) + case(1) + liq_wat = -1 + ice_wat = -1 + rainwat = -1 + snowwat = -1 + graupel = -1 + cld_amt = -1 + case(3) + liq_wat = 2 + ice_wat = 3 + rainwat = -1 + snowwat = -1 + graupel = -1 + cld_amt = -1 + case(6:7) + liq_wat = 2 + ice_wat = 3 + rainwat = 4 + snowwat = 5 + graupel = 6 + cld_amt = 7 + end select +#else + liq_wat = get_tracer_index (MODEL_ATMOS, 'liq_wat') + ice_wat = get_tracer_index (MODEL_ATMOS, 'ice_wat') + rainwat = get_tracer_index (MODEL_ATMOS, 'rainwat') + snowwat = get_tracer_index (MODEL_ATMOS, 'snowwat') + graupel = get_tracer_index (MODEL_ATMOS, 'graupel') + cld_amt = get_tracer_index (MODEL_ATMOS, 'cld_amt') + !$ser verbatim o3mr = get_tracer_index (MODEL_ATMOS, 'o3mr') + !$ser verbatim sgs_tke = get_tracer_index (MODEL_ATMOS, 'sgs_tke') +#endif + + if ( do_sat_adj .and. nwat>=6 ) then + print*,'CODE NOT TESTED HERE 1' + fast_mp_consv = (.not.do_adiabatic_init) .and. consv>consv_min + do k=1,km + kmp = k + if ( pfull(k) > 10.E2 ) exit + enddo + call qs_init(kmp) + endif + +!$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & +!$OMP remap_pt,remap_te,mfy,mfx,cx,cy,hydrostatic, & +!$OMP pt,pk,rg,peln,q,nwat,liq_wat,rainwat,ice_wat,snowwat, & +#ifdef SERIALIZE +!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, cld_amt, mode,qmin, abskord,iep1, iedp1, jedp1, js2d, o3mr, sgs_tke, & +#endif +!$OMP graupel,q_con, sphum,cappa,r_vir,rcp,cp,k1k,delp, & +!$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & +!$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & +!$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,& +!$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & +!$OMP peln_3d, gz_2d, pe0_3d, pe3_3d, phis_3d, te_2d, te_2d_f32, te0_2d, te0_2d_f32) & +!$OMP private(gz,cvm,bkh,dp2, & +!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64) + +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js, je +!$ser verbatim do i = is, ie +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js-1, je+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim qmin=t_min +!$ser verbatim w_max_ = w_max +!$ser verbatim w_min_ = w_min +!$ser verbatim rsin2=gridstruct%rsin2 +!$ser verbatim cosa_s=gridstruct%cosa_s +!$ser verbatim grav_=grav +!$ser verbatim te_2d_f32 = real(te_2d) +!$ser verbatim area_64 = real(gridstruct%area_64) +!$ser verbatim te0_2d_f32 = real(te0_2d) +!$ser savepoint Remapping_GEOS-In +!$ser data pe_=pe_3d +!$ser data ptop=ptop +!$ser data qvapor=q(:,:,:,sphum) +!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qice=q(:,:,:,ice_wat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qgraupel=q(:,:,:,graupel) +!$ser data qcld=q(:,:,:,cld_amt) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ak=ak +!$ser data bk=bk +!$ser data ps=ps +!$ser data peln_3d=peln_3d +!$ser data r_vir=r_vir +!$ser data cvm=cvm +!$ser data pk=pk +!$ser data pkz=pkz +!$ser data akap=akap +!$ser data t_min=qmin +!$ser data kord_wz=kord_wz +!$ser data ws_=ws +!$ser data w=w +!$ser data u=u +!$ser data v=v +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser data w_max=w_max_ +!$ser data w_min=w_min_ +!$ser data kord_mt=kord_mt +!$ser data cosa_s=cosa_s +!$ser data rsin2=rsin2 +!$ser data hs=hs +!$ser data grav=grav_ +!$ser data te_2d_=te_2d_f32 +!$ser data te0_2d_=te0_2d_f32 +!$ser data area_64=area_64 +!$ser data te=te + + do 1000 j=js,je+1 + + do k=1,km+1 + do i=is,ie + pe1(i,k) = pe(i,k,j) + enddo + enddo + + do i=is,ie + pe2(i, 1) = ptop + pe2(i,km+1) = pe(i,km+1,j) + enddo + + if ( j /= (je+1) ) then + + if (remap_t) then + ! Remap T in logP + ! Note: pt at this stage is Theta_v + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 2' + ! Transform virtual pt to virtual Temp + do k=1,km + do i=is,ie + pt(i,j,k) = pt(i,j,k)*(pk(i,j,k+1)-pk(i,j,k))/(akap*(peln(i,k+1,j)-peln(i,k,j))) + enddo + enddo + else + ! Transform "density pt" to "density temp" + !$ser verbatim if(j == js2d) then + !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' + !$ser savepoint MoistCVPlusPt_2d-In + !$ser data qvapor_js=q(:,j,:,sphum) qliquid_js=q(:,j,:,liq_wat) qice_js=q(:,j,:,ice_wat) qrain_js=q(:,j,:,rainwat) qsnow_js=q(:,j,:,snowwat) qgraupel_js=q(:,j,:,graupel) qcld_js=q(:,j,:,cld_amt) gz1d=gz cvm=cvm r_vir=r_vir cappa=cappa rrg=rrg delp=delp delz=delz pt=pt k1k=k1k j_2d=js2d q_con=q_con + !$ser verbatim endif + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + !$ser verbatim q_con(i,j,k) = gz(i) + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pt(i,j,k) = pt(i,j,k)*exp(cappa(i,j,k)/(1.-cappa(i,j,k))*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#else + do i=is,ie + pt(i,j,k) = pt(i,j,k)*exp(k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#endif + enddo + !$ser verbatim if(j == js2d) then + !$ser savepoint MoistCVPlusPt_2d-Out + !$ser data gz1d=gz cvm=cvm pt=pt cappa=cappa q_con=q_con + !$ser verbatim endif + + endif ! hydro test + elseif (remap_pt) then + print*,'CODE NOT TESTED HERE 3' + ! Remap PT in P + ! pt is already virtual PT + elseif (remap_te) then + print*,'CODE NOT TESTED HERE 4' + ! Remap TE in logP + ! Transform virtual pt to total energy + if ( hydrostatic ) then + call pkez(km, is, ie, js, je, j, pe, pk, akap, peln, pkz, ptop) + do i=is,ie + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) + cp_air*pt(i,j,k)*(pk(i,j,k+1)-pk(i,j,k)) + enddo + enddo + do k=1,km+1 + do i=is,ie + phis(i,k) = phis(i,k) * pe(i,k,j) + enddo + enddo + ! Compute cp*T + KE + do k=1,km + do i=is,ie + te(i,j,k) = 0.25*gridstruct%rsin2(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + cp_air*pt(i,j,k)*pkz(i,j,k) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) + enddo + enddo + else + ! TE using 3D winds (pt is virtual potential temperature): + do i=is,ie + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do k=1,km+1 + do i=is,ie + phis(i,k) = phis(i,k) * pe(i,k,j) + enddo + enddo + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pkz(i,j,k) = exp(cappa(i,j,k)/(1.-cappa(i,j,k))*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + ! TE = KE + Cv*T_v + PE + te(i,j,k) = 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + cvm(i)*pt(i,j,k)*pkz(i,j,k) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) + enddo +#else + do i=is,ie + pkz(i,j,k) = exp(k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + ! TE = KE + Cv*T_v + PE + te(i,j,k) = 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + cv_air*pt(i,j,k)*pkz(i,j,k) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) + enddo +#endif + enddo + endif ! hydro test + endif + + ! update ps + do i=is,ie + ps(i,j) = pe1(i,km+1) + enddo + ! + ! Hybrid sigma-P coordinate: + ! + do k=2,km + do i=is,ie + pe2(i,k) = ak(k) + bk(k)*pe(i,km+1,j) + enddo + enddo + do k=1,km + do i=is,ie + dp2(i,k) = pe2(i,k+1) - pe2(i,k) + enddo + enddo + !------------------ + ! Compute p**Kappa + !------------------ + do k=1,km+1 + do i=is,ie + pk1(i,k) = pk(i,j,k) + pn1(i,k) = peln(i,k,j) + enddo + enddo + + do i=is,ie + pn2(i, 1) = peln(i, 1,j) + pn2(i,km+1) = peln(i,km+1,j) + pk2(i, 1) = pk1(i, 1) + pk2(i,km+1) = pk1(i,km+1) + enddo + + do k=2,km + do i=is,ie + pn2(i,k) = log(pe2(i,k)) + pk2(i,k) = exp(akap*pn2(i,k)) + enddo + enddo + + if (remap_te) then + print*,'CODE NOT TESTED HERE 5' + !---------------------------------- + ! map TE in log P + !---------------------------------- + if ( gmao_remap > 0 ) then + call map1_gmao (km, pe1, te, & + km, pe2, te, & + is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.true.) + else + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint MapScalar_2d-In + !$ser verbatim mode=1 + !$ser data pe1_=pn1_3d pt=pt q1=te mode=mode j_2d=js2d q_min=-1.e25 + !$ser verbatim endif + call map_scalar(km, pn1, te, gz, & + km, pn2, te, & + is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), te_min) + !$ser verbatim if(j == js2d) then + !$ser savepoint MapScalar_2d-Out + !$ser data q1=te + !$ser verbatim endif + endif + else + !---------------------------------- + ! map T or PT in log P + !---------------------------------- + if ( gmao_remap > 0 ) then + print*,'CODE NOT TESTED HERE 6' + call map1_gmao (km, pe1, pt, & + km, pe2, pt, & + is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.false.) + else + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Map_Scalar-In + !$ser verbatim qmin=184.0 + !$ser data pe1_=pn1_3d pe2_=pn2_3d q1=pt j_2d=js2d q_min=qmin + !$ser verbatim endif + call map_scalar(km, pn1, pt, gz, & + km, pn2, pt, & + is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), t_min) + !$ser verbatim if(j == js2d) then + !$ser savepoint Map_Scalar-Out + !$ser data q1=pt + !$ser verbatim endif + endif + endif + + !---------------- + ! Map constituents + !---------------- + if( nq > 5 ) then + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) + !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) + !$ser verbatim + !$ser verbatim if (k < km+1) dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint MapN_Tracer_2d-In + !$ser data j_2d=js2d nq=nq pe1_=pe1_3d pe2_=pe2_3d dp2_=dp2_3d + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + !$ser verbatim endif + call mapn_tracer(nq, km, pe1, pe2, q, dp2, kord_tr, j, & + is, ie, isd, ied, jsd, jed, 0., fill) + !$ser verbatim if(j == js2d) then + !$ser savepoint MapN_Tracer_2d-Out + !$ser data qvapor=q(:,:,:,1) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + !$ser verbatim endif + + elseif ( nq > 0 ) then + print*,'CODE NOT TESTED HERE 7' + ! Remap one tracer at a time + do iq=1,nq + call map1_q2(km, pe1, q(isd,jsd,1,iq), & + km, pe2, q2, dp2, & + is, ie, 0, kord_tr(iq), j, isd, ied, jsd, jed, 0.) + if (fill) call fillz(ie-is+1, km, 1, q2, dp2) + do k=1,km + do i=is,ie + q(i,j,k,iq) = q2(i,k) + enddo + enddo + enddo + endif + + if ( .not. hydrostatic ) then + ! Remap vertical wind: + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Map1_PPM_W-In + !$ser data pe1_=pe1_3d + !$ser data pe2_=pe2_3d + !$ser data ws_=ws + !$ser data w_=w + !$ser data kord_wz=kord_wz + !$ser verbatim endif + call map1_ppm (km, pe1, w, ws(is,j), & + km, pe2, w, & + is, ie, j, isd, ied, jsd, jed, -2, kord_wz) + !$ser verbatim if(j == js2d) then + !$ser savepoint Map1_PPM_W-Out + !$ser data w_=w + !$ser verbatim endif + + !$ser verbatim if(j == js2d) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) + !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim do i = is, ie + !$ser verbatim gz_2d(i,j) = gz(i) + !$ser verbatim enddo + !$ser verbatim do k = 1, km + !$ser verbatim do i = is, ie + !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser savepoint Map1_PPM_delz-In + !$ser data pe1_=pe1_3d + !$ser data pe2_=pe2_3d + !$ser data dp2_3d=dp2_3d + !$ser data gz_=gz_2d + !$ser data delz_=delz + !$ser data delp=delp + !$ser data kord_wz=kord_wz + !$ser verbatim endif + ! Remap delz for hybrid sigma-p coordinate + do k=1,km + do i=is,ie + delz(i,j,k) = -delz(i,j,k) / delp(i,j,k) ! ="specific volume"/grav + enddo + enddo + + call map1_ppm (km, pe1, delz, gz, & + km, pe2, delz, & + is, ie, j, isd, ied, jsd, jed, 1, abs(kord_wz)) + + do k=1,km + do i=is,ie + delz(i,j,k) = -delz(i,j,k)*dp2(i,k) + enddo + enddo + !$ser verbatim if(j == js2d) then + !$ser savepoint Map1_PPM_delz-Out + !$ser data delz_=delz + !$ser verbatim endif + + !Fix excessive w - momentum conserving --- sjl + ! gz(:) used here as a temporary array + if ( w_limiter ) then +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km +!$ser verbatim do i = is, ie +!$ser verbatim dp2_3d(i,:,k) = dp2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim w_max_ = w_max +!$ser verbatim w_min_ = w_min + +!$ser savepoint W_fix_consrv_moment-In +!$ser data w=w dp2_W=dp2_3d w_max=w_max_ w_min=w_min_ +!$ser verbatim endif + do k=1,km + do i=is,ie + w2(i,k) = w(i,j,k) + enddo + enddo + do k=1, km-1 + do i=is,ie + if ( w2(i,k) > w_max ) then + if(j == js2d) then + print*,"ENTERED LOOP 1" + endif + gz(i) = (w2(i,k)-w_max) * dp2(i,k) + w2(i,k ) = w_max + w2(i,k+1) = w2(i,k+1) + gz(i)/dp2(i,k+1) + !print*, ' W_LIMITER down: ', i,j,k, w2(i,k:k+1), w(i,j,k:k+1) + elseif ( w2(i,k) < w_min ) then + if(j == js2d) then + print*,"ENTERED LOOP 2" + endif + gz(i) = (w2(i,k)-w_min) * dp2(i,k) + w2(i,k ) = w_min + w2(i,k+1) = w2(i,k+1) + gz(i)/dp2(i,k+1) + !print*, ' W_LIMITER down: ', i,j,k, w2(i,k:k+1), w(i,j,k:k+1) + endif + enddo + enddo + do k=km, 2, -1 + do i=is,ie + if ( w2(i,k) > w_max ) then + if(j == js2d) then + print*,"ENTERED LOOP 3" + endif + gz(i) = (w2(i,k)-w_max) * dp2(i,k) + w2(i,k ) = w_max + w2(i,k-1) = w2(i,k-1) + gz(i)/dp2(i,k-1) + !print*, ' W_LIMITER up: ', i,j,k, w2(i,k-1:k), w(i,j,k-1:k) + elseif ( w2(i,k) < w_min ) then + if(j == js2d) then + print*,"ENTERED LOOP 4" + endif + gz(i) = (w2(i,k)-w_min) * dp2(i,k) + w2(i,k ) = w_min + w2(i,k-1) = w2(i,k-1) + gz(i)/dp2(i,k-1) + !print*, ' W_LIMITER up: ', i,j,k, w2(i,k-1:k), w(i,j,k-1:k) + endif + enddo + enddo + do i=is,ie + if (w2(i,1) > w_max*2. ) then + if(j == js2d) then + print*,'ENTERED LOOP 5' + endif + w2(i,1) = w_max*2 ! sink out of the top of the domain + !print*, ' W_LIMITER top limited: ', i,j,1, w2(i,1), w(i,j,1) + elseif (w2(i,1) < w_min*2. ) then + if(j == js2d) then + print*,'ENTERED LOOP 6' + endif + w2(i,1) = w_min*2. + !print*, ' W_LIMITER top limited: ', i,j,1, w2(i,1), w(i,j,1) + endif + enddo + do k=1,km + do i=is,ie + w(i,j,k) = w2(i,k) + enddo + enddo +!$ser verbatim if(j == js2d) then +!$ser savepoint W_fix_consrv_moment-Out +!$ser data w=w +!$ser verbatim endif + endif + endif + + endif !(j < je+1) + +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim pe_3d(i,j-1,k) = pe(i,k,j-1) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Pressures_mapU-In +!$ser data pe_=pe_3d +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser data ak=ak +!$ser data bk=bk +!$ser data ptop=ptop +!$ser data u_=u +!$ser data kord_mt=kord_mt +!$ser verbatim if (present(mfy)) then +!$ser data mfy_=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy_=cy +!$ser verbatim endif +!$ser verbatim endif + !------ + ! map u + !------ + do i=is,ie+1 + pe0(i,1) = ptop + enddo + + do k=2,km+1 + do i=is,ie + pe0(i,k) = 0.5*(pe(i,k,j-1)+pe(i,k,j)) + enddo + enddo + + do k=1,km+1 + bkh = 0.5*bk(k) + do i=is,ie + pe3(i,k) = ak(k) + bkh*(pe(i,km+1,j-1)+pe(i,km+1,j)) + enddo + enddo + + call map1_ppm( km, pe0(is:ie,:), u, gz, & + km, pe3(is:ie,:), u, & + is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) + if (present(mfy)) then + ! print*, 'mfy present' + call map1_ppm( km, pe0(is:ie,:), mfy, gz, & + km, pe3(is:ie,:), mfy, & + is, ie, j, is, ie, js, je+1, -1, kord_mt) + endif + if (present(cy)) then + ! print*, 'cy present' + call map1_ppm( km, pe0(is:ie,:), cy, gz, & + km, pe3(is:ie,:), cy, & + is, ie, j, isd, ied, js, je+1, -1, kord_mt) + endif + +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Pressures_mapU-Out +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser data u_=u +!$ser verbatim if (present(mfy)) then +!$ser data mfy_=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy_=cy +!$ser verbatim endif +!$ser verbatim endif + +! Note : This serialization portion will test the update of pe0 +! and pe3 +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim pe_3d(i-1,j,k) = pe(i-1,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Pressures_mapV-In +!$ser data pe_=pe_3d +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser data ak=ak +!$ser data bk=bk +!$ser data v_=v +!$ser data kord_mt=kord_mt +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser verbatim endif + !------ + ! map v + !------ + if (j < je+1) then + do k=2,km+1 + do i=is,ie+1 + pe0(i,k) = 0.5*(pe(i-1,k, j)+pe(i,k, j)) + enddo + enddo + + do k=1,km+1 + bkh = 0.5*bk(k) + do i=is,ie+1 + pe3(i,k) = ak(k) + bkh*(pe(i-1,km+1,j)+pe(i,km+1,j)) + enddo + enddo + + call map1_ppm (km, pe0, v, gz, & + km, pe3, v, is, ie+1, & + j, isd, ied+1, jsd, jed, -1, kord_mt) + if (present(mfx)) then + ! print*, 'mfx present' + call map1_ppm (km, pe0, mfx, gz, & + km, pe3, mfx, is, ie+1, & + j, is, ie+1, js, je, -1, kord_mt) + endif + if (present(cx)) then + ! print*, 'cx present' + call map1_ppm (km, pe0, cx, gz, & + km, pe3, cx, is, ie+1, & + j, is, ie+1, jsd, jed, -1, kord_mt) + endif + endif ! (j < je+1) + +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie+1 +!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) +!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Pressures_mapV-Out +!$ser data pe0_=pe0_3d +!$ser data pe3_=pe3_3d +!$ser data v_=v +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser verbatim endif + +1000 continue + +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js, je +!$ser verbatim do i = is, ie +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js-1, je+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe_=pe_3d +!$ser data qvapor=q(:,:,:,sphum) +!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qice=q(:,:,:,ice_wat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qgraupel=q(:,:,:,graupel) +!$ser data qcld=q(:,:,:,cld_amt) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data peln_3d=peln_3d +!$ser data r_vir=r_vir +!$ser data cvm=cvm +!$ser data pk=pk +!$ser data pkz=pkz +!$ser data akap=akap +!$ser data t_min=qmin +!$ser data kord_wz=kord_wz +!$ser data ws_=ws +!$ser data w=w +!$ser data u=u +!$ser data v=v +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser data w_max=w_max_ +!$ser data w_min=w_min_ +!$ser data kord_mt=kord_mt +!$ser data cosa_s=cosa_s +!$ser data rsin2=rsin2 +!$ser data hs=hs +!$ser data grav=grav_ +!$ser data te_2d_=te_2d_f32 +!$ser data te0_2d_=te0_2d_f32 +!$ser data area_64=area_64 +!$ser data te=te + +! Update pressure variables and get new pkz, T_v, and omega + +!$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & +!$OMP remap_pt,remap_te,mfy,mfx,cx,cy,hydrostatic, & +!$OMP pt,pk,rg,peln,q,nwat,liq_wat,rainwat,ice_wat,snowwat, & +!$OMP graupel,sphum,cappa,r_vir,rcp,cp,k1k,delp, cld_amt, & +!$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & +!$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, & +!$OMP hs,w,ws,kord_wz,do_omega,omga,rrg,kord_mt, & +#ifdef SERIALIZE +!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, js2d, & +#endif +!$OMP pe2_3d, pe_3d, peln_3d, pn2_3d, pk2_3d, pe0, pe0_3d, pe3_3d, & +!$OMP pe1_3d, pn1, dp2_3d, q_con, pe1, pn1_3d) & +!$OMP private(gz,cvm,kp,k_next,bkh,dp2, & +!$OMP pe2,pe3,pk2,pn2,phis,tpe,dlnp,tmp) + do 2000 j=js,je + +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint PE_pk_delp_peln-In +!$ser data pe2_=pe2_3d +!$ser data pe_=pe_3d +!$ser data peln_=peln_3d +!$ser data pn2_=pn2_3d +!$ser data pk2_=pk2_3d +!$ser data delp=delp +!$ser data pk=pk +!$ser data ak=ak +!$ser data bk=bk +!$ser data akap=akap +!$ser data ptop=ptop +!$ser verbatim endif + !---------- + ! Update pe + !---------- + do i=is,ie + pe2(i, 1) = ptop + pe2(i,km+1) = pe(i,km+1,j) + enddo + do k=2,km + do i=is,ie + pe2(i,k) = ak(k) + bk(k)*pe(i,km+1,j) + enddo + enddo + do k=1,km+1 + do i=is,ie + pe(i,k,j) = pe2(i,k) + enddo + enddo + + !---------- + ! Update pk + !---------- + do i=is,ie + pn2(i, 1) = peln(i, 1,j) + pn2(i,km+1) = peln(i,km+1,j) + pk2(i, 1) = pk(i,j, 1) + pk2(i,km+1) = pk(i,j,km+1) + enddo + do k=2,km + do i=is,ie + pn2(i,k) = log(pe2(i,k)) + pk2(i,k) = exp(akap*pn2(i,k)) + enddo + enddo + do k=1,km+1 + do i=is,ie + pk(i,j,k) = pk2(i,k) + enddo + enddo + + !------------ + ! update delp + !------------ + do k=1,km + do i=is,ie + delp(i,j,k) = pe2(i,k+1) - pe2(i,k) + enddo + enddo + + !------------ + ! update logP + !------------ + do k=1,km+1 + do i=is,ie + peln(i,k,j) = pn2(i,k) + enddo + enddo + +!$ser verbatim if(j == js2d) then +!$ser verbatim do k = 1, km+1 +!$ser verbatim do i = is, ie +!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) +!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint PE_pk_delp_peln-Out +!$ser data pe2_=pe2_3d +!$ser data pe_=pe_3d +!$ser data peln_=peln_3d +!$ser data pn2_=pn2_3d +!$ser data pk2_=pk2_3d +!$ser data delp=delp +!$ser data pk=pk +!$ser verbatim endif + + !--------------------- + ! Compute pkz and T_v + !--------------------- + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 8' + do k=1,km + do i=is,ie + pkz(i,j,k) = (pk(i,j,k+1)-pk(i,j,k))/(akap*(peln(i,k+1,j)-peln(i,k,j))) + enddo + enddo + if (.not.remap_t) then + if (remap_te) then + ! Get updated T_v (store in pt) + do i=is,ie + gz(i) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + tpe = te(i,j,k) - gz(i) - 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j) ) + dlnp = rg*(peln(i,k+1,j) - peln(i,k,j)) + tmp = tpe / (cp - pe(i,k,j)*dlnp/delp(i,j,k)) + pt(i,j,k) = tmp + gz(i) = gz(i) + dlnp*tmp + enddo + enddo ! end k-loop + else + ! Make pt T_v + do k=1,km + do i=is,ie + pt(i,j,k) = pt(i,j,k)*pkz(i,j,k) + enddo + enddo + endif + endif + else + if (remap_te) then + print*,'CODE NOT TESTED HERE 9' + ! Invert TE using 3D winds to get pt (virtual temperature) and pkz: + do i=is,ie + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do k=1,km+1 + do i=is,ie + phis(i,k) = phis(i,k) * pe(i,k,j) + enddo + enddo + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + tpe = te(i,j,k) - & + ( 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & + + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) ) + pt(i,j,k) = tpe/cvm(i) + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#else + do i=is,ie + tpe = te(i,j,k) - & + ( 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j))) ) + pt(i,j,k) = tpe/cv_air + pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#endif + enddo + endif + if (remap_t) then + ! print*,'CODE EXECUTED' +! Note: pt at this stage is T_v or T_m + !$ser verbatim if(j == js2d) then + !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' + !$ser savepoint MoistCVPlusPkz_2d-In + !$ser data qvapor_js=q(:,j,:,sphum) + !$ser data qliquid_js=q(:,j,:,liq_wat) + !$ser data qice_js=q(:,j,:,ice_wat) + !$ser data qrain_js=q(:,j,:,rainwat) + !$ser data qsnow_js=q(:,j,:,snowwat) + !$ser data qgraupel_js=q(:,j,:,graupel) + !$ser data qcld_js=q(:,j,:,cld_amt) + !$ser data r_vir=r_vir + !$ser data cappa=cappa + !$ser data rrg=rrg + !$ser data delp=delp + !$ser data delz=delz + !$ser data pkz=pkz + !$ser data pt=pt + !$ser data k1k=k1k + !$ser verbatim endif + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) + pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#else + do i=is,ie + pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) + enddo +#endif + enddo + !$ser verbatim if(j == js2d) then + !$ser savepoint MoistCVPlusPkz_2d-Out + !$ser data pkz=pkz + !$ser data cappa=cappa + !$ser verbatim endif + endif + if (remap_pt) then + print*,'CODE NOT TESTED HERE 10' +! Note: pt at this stage is Theta_v + do k=1,km + do i=is,ie + pkz(i,j,k) = exp( k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k)) ) + ! Make pt T_v + pt(i,j,k) = pt(i,j,k)*pkz(i,j,k) + enddo + enddo + endif + endif + +! Interpolate omega/pe3 (defined at peln) to remapped cell center (dp2) + if ( do_omega ) then +! Copy omega field to pe3 + print*,'CODE NOT TESTED HERE 11' + do i=is,ie + pe3(i,1) = 0. + enddo + do k=2,km+1 + do i=is,ie + pe3(i,k) = omga(i,j,k-1) + enddo + enddo + do k=1,km + do i=is,ie + dp2(i,k) = 0.5*(peln(i,k,j) + peln(i,k+1,j)) + enddo + enddo + do i=is,ie + k_next = 1 + do n=1,km + kp = k_next + do k=kp,km + if( dp2(i,n) <= peln(i,k+1,j) .and. dp2(i,n) >= peln(i,k,j) ) then + omga(i,j,n) = pe3(i,k) + (pe3(i,k+1) - pe3(i,k)) * & + (dp2(i,n)-peln(i,k,j)) / (peln(i,k+1,j)-peln(i,k,j) ) + k_next = k + exit + endif + enddo + enddo + enddo + endif ! end do_omega + +2000 continue + +! Do total energy conservation and fast saturation adjustment as requested +! and fill new PT (Theta_V) for next k_split step or export dry T + +!$OMP parallel default(none) shared(is,ie,js,je,km,kmp,ptop,u,v,pe,isd,ied,jsd,jed,kord_mt, & +!$OMP remap_t,remap_pt,remap_te, & +!$OMP te_2d,te,delp,hydrostatic,hs,rg,pt,peln, adiabatic, & +!$OMP cp,delz,nwat,rainwat,liq_wat,ice_wat,snowwat, & +!$OMP graupel,q_con,r_vir,sphum,w,pk,pkz,last_step,consv, & +!$OMP do_adiabatic_init,zsum1,zsum0,te0_2d,domain, & +!$OMP ng,gridstruct,E_Flux,pdt,dtmp,reproduce_sum,q, & +!$OMP mdt,cld_amt,cappa,dtdt,out_dt,rrg,akap,do_sat_adj, & +!$OMP fast_mp_consv,kord_tm, phis_3d, te_2d_f32, te0_2d_f32, zsum1_f32, & +#ifdef SERIALIZE +!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, serial_flag, & +#endif +!$OMP js2d, pe1_3d, pe2_3d, pn1, pn1_3d, pn2, pn2_3d, pk2, pk2_3d, pe_3d, peln_3d, dp2, dp2_3d, & +!$OMP pe0_3d, pe3_3d, ps, mfy, cy, mfx, cx) & +!$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp, rsin2, cosa_s, grav_) + + dtmp = 0. + if( last_step .and. (.not.do_adiabatic_init) ) then + ! NOTE : Code can enter here since do_adiabatic_init can be False + if ( consv > consv_min ) then + ! print*, "consv > consv_min = entered", consv, consv_min +!$OMP do + do j=js,je + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 12' + do i=is,ie + gz(i) = hs(i,j) + do k=1,km + gz(i) = gz(i) + rg*pt(i,j,k)*(peln(i,k+1,j)-peln(i,k,j)) + enddo + enddo + do i=is,ie + te_2d(i,j) = pe(i,km+1,j)*hs(i,j) - pe(i,1,j)*gz(i) + enddo + do k=1,km + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cp*pt(i,j,k) + & + 0.25*gridstruct%rsin2(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j))) + enddo + enddo + else +! TE using 3D winds (pt is virtual temperature): + ! if(j==js) print*,"CODE EXECUTED 1" + !$ser verbatim if(j == js) then + !$ser verbatim do k = 1, km+1 + !$ser verbatim do i = is, ie + !$ser verbatim phis_3d(i,j,k) = phis(i,k) + !$ser verbatim enddo + !$ser verbatim enddo + !$ser verbatim rsin2=gridstruct%rsin2 + !$ser verbatim cosa_s=gridstruct%cosa_s + !$ser verbatim do i = is, ie + !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) + !$ser verbatim enddo + !$ser verbatim grav_=grav + !!$ser verbatim print*,'MoistCVPlusTe_2D serialization' + !$ser savepoint MoistCVPlusTe_2d-In + !$ser data qvapor_js=q(:,j,:,sphum) + !$ser data qliquid_js=q(:,j,:,liq_wat) + !$ser data qice_js=q(:,j,:,ice_wat) + !$ser data qrain_js=q(:,j,:,rainwat) + !$ser data qsnow_js=q(:,j,:,snowwat) + !$ser data qgraupel_js=q(:,j,:,graupel) + !$ser data delp=delp + !$ser data pt=pt + !$ser data phis_=phis_3d + !$ser data te_2d_=te_2d_f32 + !$ser data u=u + !$ser data v=v + !$ser data w=w + !$ser data cosa_s=cosa_s + !$ser data rsin2=rsin2 + !$ser data hs=hs + !$ser data grav=grav_ + !$ser data delz=delz + !$ser verbatim endif + do i=is,ie + te_2d(i,j) = 0. + phis(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do k=1,km +#ifdef MOIST_CAPPA + ! if(j==js) print*,"MOIST_CAPPA EXECUTED" + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cvm(i)*pt(i,j,k) + & + 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)))) + enddo +#else + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cv_air*pt(i,j,k) + & + 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & + u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)))) + enddo +#endif + enddo ! k-loop + !$ser verbatim if(j == js) then + !$ser verbatim do i = is, ie + !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) + !$ser verbatim enddo + !$ser savepoint MoistCVPlusTe_2d-Out + !$ser data te_2d_=te_2d_f32 + !$ser verbatim endif + + endif ! end non-hydro + + !$ser verbatim if(j == js2d) then + !$ser verbatim te_2d_f32 = real(te_2d) + !$ser verbatim te0_2d_f32 = real(te0_2d) + !$ser verbatim zsum1_f32 = real(zsum1) + !$ser savepoint Te_Zsum-In + !$ser data te_2d_=te_2d_f32 + !$ser data te0_2d_=te0_2d_f32 + !$ser data zsum1=zsum1_f32 + !$ser data delp=delp + !$ser data pkz=pkz + !$ser verbatim endif + do i=is,ie + te_2d(i,j) = te0_2d(i,j) - te_2d(i,j) + zsum1(i,j) = pkz(i,j,1)*delp(i,j,1) + enddo + do k=2,km + do i=is,ie + zsum1(i,j) = zsum1(i,j) + pkz(i,j,k)*delp(i,j,k) + enddo + enddo + if ( hydrostatic ) then + do i=is,ie + zsum0(i,j) = ptop*(pk(i,j,1)-pk(i,j,km+1)) + zsum1(i,j) + enddo + endif + !$ser verbatim if(j == js2d) then + !$ser verbatim te_2d_f32 = real(te_2d) + !$ser verbatim zsum1_f32 = real(zsum1) + !!$ser verbatim print*, 'sum te_2d: ', sum(te_2d_f32(:,1)), sum(te_2d(:,1)) + !$ser savepoint Te_Zsum-Out + !$ser data te_2d_=te_2d_f32 + !$ser data zsum1=zsum1_f32 + !$ser verbatim endif + + enddo ! j-loop + +!$OMP single + !print*,"MPP GLOBAL SUM CODE EXECUTED 1" + !$ser savepoint Mpp_global_sum-In + !$ser verbatim serial_flag=.true. + !$ser verbatim j=24 + !$ser data x_compute_size=j + !$ser data y_compute_size=j + !$ser verbatim j=1 + !$ser data x_compute_begin=j + !$ser data y_compute_begin=j + !$ser data max_ntile_pe=j + !$ser data tile=j + !$ser verbatim j=0 + !$ser data ioff=j + !$ser data joff=j + !$ser data serial_flag=serial_flag + !$ser data tesum=tesum + tesum = mpp_global_sum(domain, te_2d*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + + !$ser verbatim te_2d=te_2d*gridstruct%area_64(is:ie,js:je) + !!$ser verbatim print*,'Sum of input into mpp_global_sum = ', sum(te_2d) + !$ser data inputArray=te_2d + + !$ser savepoint Mpp_global_sum-Out + !$ser data tesum=tesum + + !print*,'tesum = ', tesum + E_Flux = DBLE(consv)*tesum / DBLE(grav*pdt*4.*pi*radius**2) ! unit: W/m**2 + ! Note pdt is "phys" time step + if ( hydrostatic ) then + print*,'CODE NOT TESTED HERE 13' + zsum = mpp_global_sum(domain, zsum0*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = tesum / DBLE(cp*zsum) + else + zsum = mpp_global_sum(domain, zsum1*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = tesum / DBLE(cv_air*zsum) + endif +!$OMP end single + + elseif ( consv < -consv_min ) then + print*,'CODE NOT TESTED HERE 14' +!$OMP do + do j=js,je + do i=is,ie + zsum1(i,j) = pkz(i,j,1)*delp(i,j,1) + enddo + do k=2,km + do i=is,ie + zsum1(i,j) = zsum1(i,j) + pkz(i,j,k)*delp(i,j,k) + enddo + enddo + if ( hydrostatic ) then + do i=is,ie + zsum0(i,j) = ptop*(pk(i,j,1)-pk(i,j,km+1)) + zsum1(i,j) + enddo + endif + enddo + + E_Flux = consv +!$OMP single + if ( hydrostatic ) then + zsum = mpp_global_sum(domain, zsum0*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = E_Flux*(grav*pdt*4.*pi*radius**2) / (cp*zsum) + else + zsum = mpp_global_sum(domain, zsum1*gridstruct%area_64(is:ie,js:je), & + flags=BITWISE_EFP_SUM) + dtmp = E_Flux*(grav*pdt*4.*pi*radius**2) / (cv_air*zsum) + endif +!$OMP end single + endif ! end consv check + endif ! end last_step check + +! Note: pt at this stage is T_v + if ( do_sat_adj .and. nwat>=6 ) then + print*,'CODE NOT TESTED HERE 15' + call timing_on('sat_adj2') +!$OMP do + do k=kmp,km + do j=js,je + do i=is,ie + dpln(i,j) = peln(i,k+1,j) - peln(i,k,j) + enddo + enddo + call fv_sat_adj(abs(mdt), r_vir, is, ie, js, je, ng, hydrostatic, fast_mp_consv, & + te(isd,jsd,k), q(isd,jsd,k,sphum), q(isd,jsd,k,liq_wat), & + q(isd,jsd,k,ice_wat), q(isd,jsd,k,rainwat), & + q(isd,jsd,k,snowwat), q(isd,jsd,k,graupel), & + hs ,dpln, delz(isd:,jsd:,k), pt(isd,jsd,k), delp(isd,jsd,k), & + cappa(isd:,jsd:,k), gridstruct%area_64, dtdt(is:,js:,k), out_dt, last_step, q(isd,jsd,k,cld_amt)) + if ( .not. hydrostatic ) then + do j=js,je + do i=is,ie +#ifdef MOIST_CAPPA + pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) +#else + pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) +#endif + enddo + enddo + endif + enddo ! OpenMP k-loop + + if ( fast_mp_consv ) then +!$OMP do + do j=js,je + do i=is,ie + do k=kmp,km + te0_2d(i,j) = te0_2d(i,j) + te(i,j,k) + enddo + enddo + enddo + endif + call timing_off('sat_adj2') + endif ! do_sat_adj + + + if ( last_step .and. (.not. adiabatic) ) then + ! Output temperature if last_step + if ( .not. hydrostatic ) then + ! print*,'CODE EXECUTED' + !$ser savepoint MoistCVPlusPt_2d_last_step-In + !$ser data qvapor=q(:,:,:,sphum) + !$ser data qliquid=q(:,:,:,liq_wat) + !$ser data qice=q(:,:,:,ice_wat) + !$ser data qrain=q(:,:,:,rainwat) + !$ser data qsnow=q(:,:,:,snowwat) + !$ser data qgraupel=q(:,:,:,graupel) + !$ser data r_vir=r_vir + !$ser data dtmp=dtmp + !$ser data pt=pt + !$ser data pkz=pkz +!$OMP do + do k=1,km + do j=js,je +#ifdef USE_COND + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, gz, cvm) + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / ((1.+r_vir*q(i,j,k,sphum))*(1.-gz(i))) + enddo +#else + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / (1.+r_vir*q(i,j,k,sphum)) + enddo +#endif + enddo ! j-loop + enddo ! k-loop + !$ser savepoint MoistCVPlusPt_2d_last_step-Out + !$ser data pt=pt +#ifdef USE_COND + ! print*, "USE_COND active" + !$ser savepoint Cond_output-In + !$ser data q_con=q_con + !$ser data qliquid=q(:,:,:,liq_wat) + !$ser data qice=q(:,:,:,ice_wat) + !$ser data qrain=q(:,:,:,rainwat) + !$ser data qsnow=q(:,:,:,snowwat) + !$ser data qgraupel=q(:,:,:,graupel) +! Fill condensate output +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + q_con(i,j,k) = 0.0 + if (liq_wat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,liq_wat) + if (ice_wat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,ice_wat) + if (rainwat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,rainwat) + if (snowwat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,snowwat) + if (graupel > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,graupel) + enddo + enddo ! j-loop + enddo ! k-loop +#endif + !$ser savepoint Cond_output-Out + !$ser data q_con=q_con + + else + print*,'CODE NOT TESTED HERE 16' +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / (1.+r_vir*q(i,j,k,sphum)) + enddo + enddo ! j-loop + enddo ! k-loop + endif + + elseif ( last_step .and. adiabatic ) then + print*,'CODE NOT TESTED HERE 17' +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) + enddo + enddo ! j-loop + enddo ! k-loop + + else + print*,'CODE NOT TESTED HERE 18' + ! Top of the loop expects PT to be Theta_V +!$OMP do + do k=1,km + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k)/pkz(i,j,k) + enddo + enddo + enddo + + endif +!$OMP end parallel + +!$ser savepoint GetMPIProp-In +!$ser data delz=delz +!$ser savepoint GetMPIProp-Out +!$ser data delz=delz + end subroutine Lagrangian_to_Eulerian + +!>@brief The subroutine 'compute_total_energy' performs the FV3-consistent computation of the global total energy. +!>@details It includes the potential, internal (latent and sensible heat), kinetic terms. + subroutine compute_total_energy(is, ie, js, je, isd, ied, jsd, jed, km, & + u, v, w, delz, pt, delp, q, qc, pe, peln, hs, & + rsin2_l, cosa_s_l, & + r_vir, cp, rg, hlv, te_2d, ua, va, teq, & + moist_phys, nwat, sphum, liq_wat, rainwat, ice_wat, snowwat, graupel, hydrostatic, id_te) +!------------------------------------------------------ +! Compute vertically integrated total energy per column +!------------------------------------------------------ +! !INPUT PARAMETERS: + integer, intent(in):: km, is, ie, js, je, isd, ied, jsd, jed, id_te + integer, intent(in):: sphum, liq_wat, ice_wat, rainwat, snowwat, graupel, nwat + real, intent(inout), dimension(isd:ied,jsd:jed,km):: ua, va + real, intent(in), dimension(isd:ied,jsd:jed,km):: pt, delp + real, intent(in), dimension(isd:ied,jsd:jed,km,*):: q + real, intent(in), dimension(isd:ied,jsd:jed,km):: qc + real, intent(inout):: u(isd:ied, jsd:jed+1,km) + real, intent(inout):: v(isd:ied+1,jsd:jed, km) + real, intent(in):: w(isd:,jsd:,1:) !< vertical velocity (m/s) + real, intent(in):: delz(isd:,jsd:,1:) + real, intent(in):: hs(isd:ied,jsd:jed) !< surface geopotential + real, intent(in):: pe(is-1:ie+1,km+1,js-1:je+1) !< pressure at layer edges + real, intent(in):: peln(is:ie,km+1,js:je) !< log(pe) + real, intent(in):: cp, rg, r_vir, hlv + real, intent(in) :: rsin2_l(isd:ied, jsd:jed) + real, intent(in) :: cosa_s_l(isd:ied, jsd:jed) + logical, intent(in):: moist_phys, hydrostatic +!! Output: + real(kind=8), intent(out):: te_2d(is:ie,js:je) !< vertically integrated TE + real, intent(out):: teq(is:ie,js:je) !< Moist TE +!! Local + real, dimension(is:ie,km):: tv + real phiz(is:ie,km+1) + real cvm(is:ie), qd(is:ie) + integer i, j, k + +!---------------------- +! Output lat-lon winds: +!---------------------- +! call cubed_to_latlon(u, v, ua, va, dx, dy, rdxa, rdya, km, flagstruct%c2l_ord) + +!$OMP parallel do default(none) shared(is,ie,js,je,isd,ied,jsd,jed,km,hydrostatic,hs,pt,qc,rg,peln,te_2d, & +!$OMP pe,delp,cp,rsin2_l,u,v,cosa_s_l,delz,moist_phys,w, & +!$OMP q,nwat,liq_wat,rainwat,ice_wat,snowwat,graupel,sphum) & +!$OMP private(phiz, tv, cvm, qd) + do j=js,je + + if ( hydrostatic ) then + + do i=is,ie + phiz(i,km+1) = hs(i,j) + enddo + do k=km,1,-1 + do i=is,ie + tv(i,k) = pt(i,j,k)*(1.+qc(i,j,k)) + phiz(i,k) = phiz(i,k+1) + rg*tv(i,k)*(peln(i,k+1,j)-peln(i,k,j)) + enddo + enddo + + do i=is,ie + te_2d(i,j) = pe(i,km+1,j)*phiz(i,km+1) - pe(i,1,j)*phiz(i,1) + enddo + + do k=1,km + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cp*tv(i,k) + & + 0.25*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2 - & + (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j))) + enddo + enddo + + else +!----------------- +! Non-hydrostatic: +!----------------- + do i=is,ie + phiz(i,km+1) = hs(i,j) + do k=km,1,-1 + phiz(i,k) = phiz(i,k+1) - grav*delz(i,j,k) + enddo + enddo + do i=is,ie + te_2d(i,j) = 0. + enddo + if ( moist_phys ) then + do k=1,km +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, qd, cvm) +#endif + do i=is,ie +#ifdef MOIST_CAPPA + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cvm(i)*pt(i,j,k)*(1.+qc(i,j,k))*(1.-qd(i)) + & +#else + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cv_air*pt(i,j,k)*(1.+qc(i,j,k)) + & +#endif + 0.5*(phiz(i,k)+phiz(i,k+1)+w(i,j,k)**2+0.5*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2-(u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j)))) + enddo + enddo + else + do k=1,km + do i=is,ie + te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cv_air*pt(i,j,k) + & + 0.5*(phiz(i,k)+phiz(i,k+1)+w(i,j,k)**2+0.5*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & + v(i,j,k)**2+v(i+1,j,k)**2-(u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j)))) + enddo + enddo + endif + endif + enddo + +!------------------------------------- +! Diganostics computation for moist TE +!------------------------------------- + if( id_te>0 ) then +!$OMP parallel do default(none) shared(is,ie,js,je,teq,te_2d,moist_phys,km,hlv,sphum,q,delp) + do j=js,je + do i=is,ie + teq(i,j) = te_2d(i,j) + enddo + if ( moist_phys ) then + do k=1,km + do i=is,ie + teq(i,j) = teq(i,j) + hlv*q(i,j,k,sphum)*delp(i,j,k) + enddo + enddo + endif + enddo + endif + + end subroutine compute_total_energy + + + subroutine pkez(km, ifirst, ilast, jfirst, jlast, j, & + pe, pk, akap, peln, pkz, ptop) + +! INPUT PARAMETERS: + integer, intent(in):: km, j + integer, intent(in):: ifirst, ilast !< Latitude strip + integer, intent(in):: jfirst, jlast !< Latitude strip + real, intent(in):: akap + real, intent(in):: pe(ifirst-1:ilast+1,km+1,jfirst-1:jlast+1) + real, intent(in):: pk(ifirst:ilast,jfirst:jlast,km+1) + real, intent(IN):: ptop +! OUTPUT + real, intent(out):: pkz(ifirst:ilast,jfirst:jlast,km) + real, intent(inout):: peln(ifirst:ilast, km+1, jfirst:jlast) !< log (pe) +! Local + real pk2(ifirst:ilast, km+1) + real pek + real lnp + real ak1 + integer i, k + + ak1 = (akap + 1.) / akap + + pek = pk(ifirst,j,1) + do i=ifirst, ilast + pk2(i,1) = pek + enddo + + do k=2,km+1 + do i=ifirst, ilast +! peln(i,k,j) = log(pe(i,k,j)) + pk2(i,k) = pk(i,j,k) + enddo + enddo + +!---- GFDL modification + if( ptop < ptop_min ) then + do i=ifirst, ilast + peln(i,1,j) = peln(i,2,j) - ak1 + enddo + else + lnp = log( ptop ) + do i=ifirst, ilast + peln(i,1,j) = lnp + enddo + endif +!---- GFDL modification + + do k=1,km + do i=ifirst, ilast + pkz(i,j,k) = (pk2(i,k+1) - pk2(i,k) ) / & + (akap*(peln(i,k+1,j) - peln(i,k,j)) ) + enddo + enddo + + end subroutine pkez + + + + subroutine remap_z(km, pe1, q1, kn, pe2, q2, i1, i2, iv, kord) + +! INPUT PARAMETERS: + integer, intent(in) :: i1 !< Starting longitude + integer, intent(in) :: i2 !< Finishing longitude + integer, intent(in) :: kord !< Method order + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + integer, intent(in) :: iv + + real, intent(in) :: pe1(i1:i2,km+1) !< height at layer edges from model top to bottom surface + real, intent(in) :: pe2(i1:i2,kn+1) !< height at layer edges from model top to bottom surface + real, intent(in) :: q1(i1:i2,km) !< Field input + +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(i1:i2,kn) !< Field output + +! LOCAL VARIABLES: + real qs(i1:i2) + real dp1( i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, delp, esl + integer i, k, l, m, k0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) ! negative + q4(1,i,k) = q1(i,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then + call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + +! Mapping + do 3000 i=i1,i2 + k0 = 1 + do 555 k=1,kn + do 100 l=k0,km +! locate the top edge: pe2(i,k) + if(pe2(i,k) <= pe1(i,l) .and. pe2(i,k) >= pe1(i,l+1)) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) >= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) < pe1(i,m+1) ) then +! Whole layer.. + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + delp = pe2(i,k+1)-pe1(i,m) + esl = delp / dp1(i,m) + qsum = qsum + delp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif +100 continue +123 q2(i,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) +555 continue +3000 continue + + end subroutine remap_z + + subroutine map_scalar( km, pe1, q1, qs, & + kn, pe2, q2, i1, i2, & + j, ibeg, iend, jbeg, jend, iv, kord, q_min) +! iv=1 + integer, intent(in) :: i1 !< Starting longitude + integer, intent(in) :: i2 !< Finishing longitude + integer, intent(in) :: iv !< Mode: 0 == constituents 1 == temp 2 == remap temp with cs scheme + integer, intent(in) :: kord !< Method order + integer, intent(in) :: j !< Current latitude + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + real, intent(in) :: qs(i1:i2) !< bottom BC + real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) !< Field output + real, intent(in):: q_min + +! DESCRIPTION: +! IV = 0: constituents +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate +! LOCAL VARIABLES: + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + integer i, k, l, m, k0, jj + integer LM1,LP0,LP1 + +!$ser verbatim real q4_1_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_2_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_3_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_4_temp(i1:i2, i1:i2,km) +!$ser verbatim real dp1_temp(i1:i2, i1:i2,km) +!$ser verbatim real pe1_temp(i1:i2, i1:i2,km+1) +!$ser verbatim real pe2_temp(i1:i2, i1:i2,kn+1) + +!$ser verbatim real LM1_INDEX(i1:i2, i1:i2, km) +!$ser verbatim real LP0_INDEX(i1:i2, i1:i2, km) + + !$ser verbatim real :: qs_2d(i1:i2, i1:i2) + +!$ser verbatim q4_1_temp = 0.0 +!$ser verbatim q4_2_temp = 0.0 +!$ser verbatim q4_3_temp = 0.0 +!$ser verbatim q4_4_temp = 0.0 +!$ser verbatim dp1_temp = 0.0 +!$ser verbatim pe1_temp = 0.0 +!$ser verbatim pe2_temp = 0.0 +!$ser verbatim LM1_INDEX = 0 +!$ser verbatim LP0_INDEX = 0 +!$ser verbatim qs_2d = 0.0 + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,j,k) + enddo + enddo + +!$ser verbatim if (j == 1) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do i = i1,i2 +!$ser verbatim qs_2d(i,j) = qs(i) +!$ser verbatim enddo +!$ser savepoint Scalar_Profile-In +!$ser data qs_=qs_2d +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser data dp1_=dp1_temp +!$ser data q_min=q_min +!$ser verbatim endif + +! Compute vertical subgrid distribution + if ( kord > 7 ) then + call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + +!$ser verbatim if (j == 1) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Scalar_Profile-Out +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser data dp1_=dp1_temp +!$ser verbatim endif + +! NOTE : q1 and q2 fields being passed into map_scalar are identical variables +! even though q1 is declared at INTENT(IN) and q2 is declared as INTENT(IN/OUT). +!$ser verbatim if(j == 1) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim do jj = i1, i2 +!$ser verbatim dp1_temp(i,jj,k) = dp1(i,k) +!$ser verbatim pe1_temp(i,jj,k) = pe1(i,k) +!$ser verbatim pe2_temp(i,jj,k) = pe2(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do jj = i1, i2 +!$ser verbatim do i = i1, i2 +!$ser verbatim pe1_temp(i,jj,km+1) = pe1(i,km+1) +!$ser verbatim pe2_temp(i,jj,km+1) = pe2(i,km+1) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Lagrangian_Contribution_Interp-In +!$ser data q1=q1 pe1_=pe1_temp pe2_=pe2_temp q4_1=q4_1_temp q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp q4_4=q4_4_temp dp1_=dp1_temp +!$ser verbatim endif + +! Interpolate field onto target Pressures +! --------------------------------------- + do i=i1,i2 + k0 = 1 + do 555 k=1,kn + LM1 = 1 + LP0 = 1 + do while( LP0.le.km ) + if (pe1(i,LP0).lt.pe2(i,k)) then + LP0 = LP0+1 + else + exit + endif + enddo + LM1 = max(LP0-1,1) + LP0 = min(LP0, km) +! Extrapolate Linearly above first model level +! ---------------------------------------------------- + if( LM1.eq.1 .and. LP0.eq.1 ) then + q2(i,j,k) = q1(i,j,1) + ( q1(i,j,2)-q1(i,j,1) )*( pe2(i,k)-pe1(i,1) ) & + /( pe1(i,2)-pe1(i,1) ) +! Extrapolate Linearly below last model level +! --------------------------------------------------- + else if( LM1.eq.km .and. LP0.eq.km ) then + q2(i,j,k) = q1(i,j,km) + ( q1(i,j,km)-q1(i,j,km-1) )*( pe2(i,k )-pe1(i,km ) ) & + /( pe1(i,km)-pe1(i,km-1) ) +! Interpolate Linearly between levels 1 => 2 and km-1 => km +! ----------------------------------------------------------------- + else if( LM1.eq.1 .or. LP0.eq.km ) then + q2(i,j,k) = q1(i,j,LP0) + ( q1(i,j,LM1)-q1(i,j,LP0) )*( pe2(i,k )-pe1(i,LP0) ) & + /( pe1(i,LM1)-pe1(i,LP0) ) + else + do l=k0,km +! locate the top edge: pe2(i,k) + if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if( pe2(i,k+1) <= pe1(i,l+1) ) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,j,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if( pe2(i,k+1) > pe1(i,m+1) ) then +! Whole layer + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif + enddo +123 q2(i,j,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) + + endif +555 continue + enddo + +!$ser verbatim if(j == 1) then +!$ser savepoint Lagrangian_Contribution_Interp-Out +!$ser data q1=q2 +!$ser verbatim endif + + end subroutine map_scalar + + + subroutine map1_ppm( km, pe1, q1, qs, & + kn, pe2, q2, i1, i2, & + j, ibeg, iend, jbeg, jend, iv, kord) + integer, intent(in) :: i1 !< Starting longitude + integer, intent(in) :: i2 !< Finishing longitude + integer, intent(in) :: iv !< Mode: 0 == constituents 1 == ??? 2 == remap temp with cs scheme + integer, intent(in) :: kord !< Method order + integer, intent(in) :: j !< Current latitude + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + real, intent(in) :: qs(i1:i2) !< bottom BC + real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) !< Field output + +! DESCRIPTION: +! IV = 0: constituents +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate + +! LOCAL VARIABLES: + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + integer i, k, l, m, k0 + +!$ser verbatim real q4_1_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_2_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_3_temp(i1:i2, i1:i2,km) +!$ser verbatim real q4_4_temp(i1:i2, i1:i2,km) +!$ser verbatim real dp1_temp(i1:i2, i1:i2,km) +!$ser verbatim real :: qs_2d(i1:i2, i1:i2) + +!$ser verbatim q4_1_temp = 0.0 +!$ser verbatim q4_2_temp = 0.0 +!$ser verbatim q4_3_temp = 0.0 +!$ser verbatim q4_4_temp = 0.0 +!$ser verbatim qs_2d = 0.0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,j,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then +!$ser verbatim if (j == 1 .and. i2 == 24) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do i = i1,i2 +!$ser verbatim qs_2d(i,j) = qs(i) +!$ser verbatim enddo +!$ser savepoint CS_Profile-In +!$ser data qs_=qs_2d +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser data dp1_=dp1_temp +!$ser data kord_=kord +!$ser data iv_=iv +!$ser verbatim endif + call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) +!$ser verbatim if (j == 1 .and. i2 == 24) then +!$ser verbatim do k = 1,kn +!$ser verbatim do i = i1,i2 +!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) +!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) +!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) +!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint CS_Profile-Out +!$ser data q4_1=q4_1_temp +!$ser data q4_2=q4_2_temp +!$ser data q4_3=q4_3_temp +!$ser data q4_4=q4_4_temp +!$ser verbatim endif + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + + do i=i1,i2 + k0 = 1 + do 555 k=1,kn + do l=k0,km +! locate the top edge: pe2(i,k) + if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if( pe2(i,k+1) <= pe1(i,l+1) ) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,j,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if( pe2(i,k+1) > pe1(i,m+1) ) then +! Whole layer + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif + enddo +123 q2(i,j,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) +555 continue + enddo + + end subroutine map1_ppm + + + subroutine mapn_tracer(nq, km, pe1, pe2, q1, dp2, kord, j, & + i1, i2, isd, ied, jsd, jed, q_min, fill) +! INPUT PARAMETERS: + integer, intent(in):: km !< vertical dimension + integer, intent(in):: j, nq, i1, i2 + integer, intent(in):: isd, ied, jsd, jed + integer, intent(in):: kord(nq) + real, intent(in):: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in):: pe2(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in):: dp2(i1:i2,km) + real, intent(in):: q_min + logical, intent(in):: fill + real, intent(inout):: q1(isd:ied,jsd:jed,km,nq) ! Field input +! LOCAL VARIABLES: + real:: q4(4,i1:i2,km,nq) + real:: q2(i1:i2,km,nq) !< Field output + real:: qsum(nq) + real:: dp1(i1:i2,km) + real:: qs(i1:i2) + real:: pl, pr, dp, esl, fac1, fac2 + integer:: i, k, l, m, k0, iq + + !$ser verbatim integer:: kord_iq, iv, im, js2d + !$ser verbatim js2d=jsd+3 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + enddo + enddo + + do iq=1,nq + do k=1,km + do i=i1,i2 + q4(1,i,k,iq) = q1(i,j,k,iq) + enddo + enddo + call scalar_profile( qs, q4(1,i1,1,iq), dp1, km, i1, i2, 0, kord(iq), q_min ) + enddo + +! Mapping + do 4000 i=i1,i2 + k0 = 1 + do 555 k=1,km + do 100 l=k0,km +! locate the top edge: pe2(i,k) + if(pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1)) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) <= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + fac1 = pr + pl + fac2 = r3*(pr*fac1 + pl*pl) + fac1 = 0.5*fac1 + do iq=1,nq + q2(i,k,iq) = q4(2,i,l,iq) + (q4(4,i,l,iq)+q4(3,i,l,iq)-q4(2,i,l,iq))*fac1 & + - q4(4,i,l,iq)*fac2 + enddo + k0 = l + goto 555 + else +! Fractional area... + dp = pe1(i,l+1) - pe2(i,k) + fac1 = 1. + pl + fac2 = r3*(1.+pl*fac1) + fac1 = 0.5*fac1 + do iq=1,nq + qsum(iq) = dp*(q4(2,i,l,iq) + (q4(4,i,l,iq)+ & + q4(3,i,l,iq) - q4(2,i,l,iq))*fac1 - q4(4,i,l,iq)*fac2) + enddo + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) > pe1(i,m+1) ) then + ! Whole layer.. + do iq=1,nq + qsum(iq) = qsum(iq) + dp1(i,m)*q4(1,i,m,iq) + enddo + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + fac1 = 0.5*esl + fac2 = 1.-r23*esl + do iq=1,nq + qsum(iq) = qsum(iq) + dp*( q4(2,i,m,iq) + fac1*( & + q4(3,i,m,iq)-q4(2,i,m,iq)+q4(4,i,m,iq)*fac2 ) ) + enddo + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif +100 continue +123 continue + do iq=1,nq + q2(i,k,iq) = qsum(iq) / dp2(i,k) + enddo +555 continue +4000 continue + + !$ser verbatim if(j == js2d ) then + !$ser verbatim im = i2-i1+1 + !$ser verbatim iv = 9 + !$ser savepoint Fillz-In + ! Note : Currently, serializing nq=nq and q2tracers=q2(:,:,1:nq) will not run the translate test + ! To successfully run the translate test, serialize nq=iv and q2tracers=q2(:,:,1:iv) + !$ser data im=im km=km nq=iv dp2=dp2 q2tracers=q2(:,:,1:iv) + !$ser verbatim endif + + if (fill) call fillz(i2-i1+1, km, nq, q2, dp2) + + !$ser verbatim if(j == js2d ) then + !$ser savepoint Fillz-Out + !$ser data q2tracers=q2(:,:,1:iv) + !$ser verbatim endif + + do iq=1,nq +! if (fill) call fillz(i2-i1+1, km, 1, q2(i1,1,iq), dp2) + do k=1,km + do i=i1,i2 + q1(i,j,k,iq) = q2(i,k,iq) + enddo + enddo + enddo + + end subroutine mapn_tracer + + + subroutine map1_q2(km, pe1, q1, & + kn, pe2, q2, dp2, & + i1, i2, iv, kord, j, & + ibeg, iend, jbeg, jend, q_min ) + + +! INPUT PARAMETERS: + integer, intent(in) :: j + integer, intent(in) :: i1, i2 + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: iv !< Mode: 0 == constituents 1 == ??? + integer, intent(in) :: kord + integer, intent(in) :: km !< Original vertical dimension + integer, intent(in) :: kn !< Target vertical dimension + + real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input + real, intent(in) :: dp2(i1:i2,kn) + real, intent(in) :: q_min +! INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(i1:i2,kn) !< Field output +! LOCAL VARIABLES: + real qs(i1:i2) + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + + integer i, k, l, m, k0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,j,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then + call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + +! Mapping + do 5000 i=i1,i2 + k0 = 1 + do 555 k=1,kn + do 100 l=k0,km +! locate the top edge: pe2(i,k) + if(pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1)) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) <= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) > pe1(i,m+1) ) then + ! Whole layer.. + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif +100 continue +123 q2(i,k) = qsum / dp2(i,k) +555 continue +5000 continue + + end subroutine map1_q2 + + + + subroutine remap_2d(km, pe1, q1, & + kn, pe2, q2, & + i1, i2, iv, kord) + integer, intent(in):: i1, i2 + integer, intent(in):: iv !< Mode: 0 == constituents 1 ==others + integer, intent(in):: kord + integer, intent(in):: km !< Original vertical dimension + integer, intent(in):: kn !< Target vertical dimension + real, intent(in):: pe1(i1:i2,km+1) !< Pressure at layer edges from model top to bottom surface in the original vertical coordinate + real, intent(in):: pe2(i1:i2,kn+1) !< Pressure at layer edges from model top to bottom surface in the new vertical coordinate + real, intent(in) :: q1(i1:i2,km) !< Field input + real, intent(out):: q2(i1:i2,kn) !< Field output +! LOCAL VARIABLES: + real qs(i1:i2) + real dp1(i1:i2,km) + real q4(4,i1:i2,km) + real pl, pr, qsum, dp, esl + integer i, k, l, m, k0 + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + q4(1,i,k) = q1(i,k) + enddo + enddo + +! Compute vertical subgrid distribution + if ( kord >7 ) then + call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) + else + call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) + endif + + do i=i1,i2 + k0 = 1 + do 555 k=1,kn +#ifdef OLD_TOP_EDGE + if( pe2(i,k+1) <= pe1(i,1) ) then +! Entire grid above old ptop + q2(i,k) = q4(2,i,1) + elseif( pe2(i,k) < pe1(i,1) .and. pe2(i,k+1)>pe1(i,1) ) then +! Partially above old ptop: + q2(i,k) = q1(i,1) +#else + if( pe2(i,k) <= pe1(i,1) ) then +! above old ptop: + q2(i,k) = q1(i,1) +#endif + else + do l=k0,km +! locate the top edge: pe2(i,k) + if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then + pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) + if(pe2(i,k+1) <= pe1(i,l+1)) then +! entire new grid is within the original grid + pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) + q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & + *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) + k0 = l + goto 555 + else +! Fractional area... + qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & + q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & + (r3*(1.+pl*(1.+pl)))) + do m=l+1,km +! locate the bottom edge: pe2(i,k+1) + if(pe2(i,k+1) > pe1(i,m+1) ) then + ! Whole layer.. + qsum = qsum + dp1(i,m)*q4(1,i,m) + else + dp = pe2(i,k+1)-pe1(i,m) + esl = dp / dp1(i,m) + qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & + (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) + k0 = m + goto 123 + endif + enddo + goto 123 + endif + endif + enddo +123 q2(i,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) + endif +555 continue + enddo + + end subroutine remap_2d + + +!>@brief Optimized vertical profile reconstruction: +!> Latest: Apr 2008 S.-J. Lin, NOAA/GFDL + subroutine scalar_profile(qs, a4, delp, km, i1, i2, iv, kord, qmin) +! Optimized vertical profile reconstruction: +! Latest: Apr 2008 S.-J. Lin, NOAA/GFDL + integer, intent(in):: i1, i2 + integer, intent(in):: km !< vertical dimension + integer, intent(in):: iv !< iv =-1: winds iv = 0: positive definite scalars iv = 1: others + integer, intent(in):: kord + real, intent(in) :: qs(i1:i2) + real, intent(in) :: delp(i1:i2,km) !< Layer pressure thickness + real, intent(inout):: a4(4,i1:i2,km) !< Interpolated values + real, intent(in):: qmin +!----------------------------------------------------------------------- + logical, dimension(i1:i2,km):: extm, ext5, ext6 + real gam(i1:i2,km) + real q(i1:i2,km+1) + real d4(i1:i2) + real bet, a_bot, grat + real pmp_1, lac_1, pmp_2, lac_2, x0, x1 + integer i, k, im + + if ( iv .eq. -2 ) then + do i=i1,i2 + gam(i,2) = 0.5 + q(i,1) = 1.5*a4(1,i,1) + enddo + do k=2,km-1 + do i=i1, i2 + grat = delp(i,k-1) / delp(i,k) + bet = 2. + grat + grat - gam(i,k) + q(i,k) = (3.*(a4(1,i,k-1)+a4(1,i,k)) - q(i,k-1))/bet + gam(i,k+1) = grat / bet + enddo + enddo + do i=i1,i2 + grat = delp(i,km-1) / delp(i,km) + q(i,km) = (3.*(a4(1,i,km-1)+a4(1,i,km)) - grat*qs(i) - q(i,km-1)) / & + (2. + grat + grat - gam(i,km)) + q(i,km+1) = qs(i) + enddo + do k=km-1,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k+1)*q(i,k+1) + enddo + enddo + else + do i=i1,i2 + grat = delp(i,2) / delp(i,1) ! grid ratio + bet = grat*(grat+0.5) + q(i,1) = ( (grat+grat)*(grat+1.)*a4(1,i,1) + a4(1,i,2) ) / bet + gam(i,1) = ( 1. + grat*(grat+1.5) ) / bet + enddo + + do k=2,km + do i=i1,i2 + d4(i) = delp(i,k-1) / delp(i,k) + bet = 2. + d4(i) + d4(i) - gam(i,k-1) + q(i,k) = ( 3.*(a4(1,i,k-1)+d4(i)*a4(1,i,k)) - q(i,k-1) )/bet + gam(i,k) = d4(i) / bet + enddo + enddo + + do i=i1,i2 + a_bot = 1. + d4(i)*(d4(i)+1.5) + q(i,km+1) = (2.*d4(i)*(d4(i)+1.)*a4(1,i,km)+a4(1,i,km-1)-a_bot*q(i,km)) & + / ( d4(i)*(d4(i)+0.5) - a_bot*gam(i,km) ) + enddo + + do k=km,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k)*q(i,k+1) + enddo + enddo + endif + +!----- Perfectly linear scheme -------------------------------- + if ( abs(kord) > 16 ) then + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + enddo + return + endif +!----- Perfectly linear scheme -------------------------------- + +!------------------ +! Apply constraints +!------------------ + im = i2 - i1 + 1 + +! Apply *large-scale* constraints + do i=i1,i2 + q(i,2) = min( q(i,2), max(a4(1,i,1), a4(1,i,2)) ) + q(i,2) = max( q(i,2), min(a4(1,i,1), a4(1,i,2)) ) + enddo + + do k=2,km + do i=i1,i2 + gam(i,k) = a4(1,i,k) - a4(1,i,k-1) + enddo + enddo + +! Interior: + do k=3,km-1 + do i=i1,i2 + if ( gam(i,k-1)*gam(i,k+1)>0. ) then +! Apply large-scale constraint to ALL fields if not local max/min + q(i,k) = min( q(i,k), max(a4(1,i,k-1),a4(1,i,k)) ) + q(i,k) = max( q(i,k), min(a4(1,i,k-1),a4(1,i,k)) ) + else + if ( gam(i,k-1) > 0. ) then +! There exists a local max + q(i,k) = max(q(i,k), min(a4(1,i,k-1),a4(1,i,k))) + else +! There exists a local min + q(i,k) = min(q(i,k), max(a4(1,i,k-1),a4(1,i,k))) + if ( iv==0 ) q(i,k) = max(0., q(i,k)) + endif + endif + enddo + enddo + +! Bottom: + do i=i1,i2 + q(i,km) = min( q(i,km), max(a4(1,i,km-1), a4(1,i,km)) ) + q(i,km) = max( q(i,km), min(a4(1,i,km-1), a4(1,i,km)) ) + enddo + + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + enddo + enddo + + do k=1,km + if ( k==1 .or. k==km ) then + do i=i1,i2 + extm(i,k) = (a4(2,i,k)-a4(1,i,k)) * (a4(3,i,k)-a4(1,i,k)) > 0. + enddo + else + do i=i1,i2 + extm(i,k) = gam(i,k)*gam(i,k+1) < 0. + enddo + endif + if ( abs(kord) > 9 ) then + do i=i1,i2 + x0 = 2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k)) + x1 = abs(a4(2,i,k)-a4(3,i,k)) + a4(4,i,k) = 3.*x0 + ext5(i,k) = abs(x0) > x1 + ext6(i,k) = abs(a4(4,i,k)) > x1 + enddo + endif + enddo + +!--------------------------- +! Apply subgrid constraints: +!--------------------------- +! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) +! Top 2 and bottom 2 layers always use monotonic mapping + + if ( iv==0 ) then + do i=i1,i2 + a4(2,i,1) = max(0., a4(2,i,1)) + enddo + elseif ( iv==-1 ) then + do i=i1,i2 + if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. + enddo + elseif ( iv==2 ) then + do i=i1,i2 + a4(2,i,1) = a4(1,i,1) + a4(3,i,1) = a4(1,i,1) + a4(4,i,1) = 0. + enddo + endif + + if ( iv/=2 ) then + do i=i1,i2 + a4(4,i,1) = 3.*(2.*a4(1,i,1) - (a4(2,i,1)+a4(3,i,1))) + enddo + call cs_limiters(im, extm(i1,1), a4(1,i1,1), 1) + endif + +! k=2 + do i=i1,i2 + a4(4,i,2) = 3.*(2.*a4(1,i,2) - (a4(2,i,2)+a4(3,i,2))) + enddo + call cs_limiters(im, extm(i1,2), a4(1,i1,2), 2) + +!------------------------------------- +! Huynh's 2nd constraint for interior: +!------------------------------------- + do k=3,km-2 + if ( abs(kord)<9 ) then + do i=i1,i2 +! Left edges + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) +! Right edges + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==9 ) then + do i=i1,i2 + if ( extm(i,k) .and. extm(i,k-1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else if ( extm(i,k) .and. extm(i,k+1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else if ( extm(i,k) .and. a4(1,i,k) abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + endif + endif + enddo + elseif ( abs(kord)==10 ) then + do i=i1,i2 + if( ext5(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + elseif( ext6(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==12 ) then + do i=i1,i2 + if( extm(i,k) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else ! not a local extremum + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) +! Check within the smooth region if subgrid profile is non-monotonic + if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) + endif + endif + enddo + elseif ( abs(kord)==13 ) then + do i=i1,i2 + if( ext6(i,k) ) then + if ( ext6(i,k-1) .and. ext6(i,k+1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==14 ) then + + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==15 ) then ! Revised abs(kord)=9 scheme + do i=i1,i2 + if ( ext5(i,k) .and. ext5(i,k-1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + else if ( ext5(i,k) .and. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + else if ( ext5(i,k) .and. a4(1,i,k)@brief The subroutine 'cs_profile' performs the optimized vertical profile reconstruction: +!>@date April 2008 +!>@author S. J. Lin, NOAA/GFDL + subroutine cs_profile(qs, a4, delp, km, i1, i2, iv, kord) +! Optimized vertical profile reconstruction: +! Latest: Apr 2008 S.-J. Lin, NOAA/GFDL + integer, intent(in):: i1, i2 + integer, intent(in):: km !< vertical dimension + integer, intent(in):: iv !< iv =-1: winds + !< iv = 0: positive definite scalars + !< iv = 1: others + integer, intent(in):: kord + real, intent(in) :: qs(i1:i2) + real, intent(in) :: delp(i1:i2,km) !< layer pressure thickness + real, intent(inout):: a4(4,i1:i2,km) !< Interpolated values +!----------------------------------------------------------------------- + logical, dimension(i1:i2,km):: extm, ext5, ext6 + real gam(i1:i2,km) + real q(i1:i2,km+1) + real d4(i1:i2) + real bet, a_bot, grat + real pmp_1, lac_1, pmp_2, lac_2, x0, x1 + integer i, k, im + + if ( iv .eq. -2 ) then + do i=i1,i2 + gam(i,2) = 0.5 + q(i,1) = 1.5*a4(1,i,1) + enddo + do k=2,km-1 + do i=i1, i2 + grat = delp(i,k-1) / delp(i,k) + bet = 2. + grat + grat - gam(i,k) + q(i,k) = (3.*(a4(1,i,k-1)+a4(1,i,k)) - q(i,k-1))/bet + gam(i,k+1) = grat / bet + enddo + enddo + do i=i1,i2 + grat = delp(i,km-1) / delp(i,km) + q(i,km) = (3.*(a4(1,i,km-1)+a4(1,i,km)) - grat*qs(i) - q(i,km-1)) / & + (2. + grat + grat - gam(i,km)) + q(i,km+1) = qs(i) + enddo + do k=km-1,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k+1)*q(i,k+1) + enddo + enddo + else + do i=i1,i2 + grat = delp(i,2) / delp(i,1) ! grid ratio + bet = grat*(grat+0.5) + q(i,1) = ( (grat+grat)*(grat+1.)*a4(1,i,1) + a4(1,i,2) ) / bet + gam(i,1) = ( 1. + grat*(grat+1.5) ) / bet + enddo + + do k=2,km + do i=i1,i2 + d4(i) = delp(i,k-1) / delp(i,k) + bet = 2. + d4(i) + d4(i) - gam(i,k-1) + q(i,k) = ( 3.*(a4(1,i,k-1)+d4(i)*a4(1,i,k)) - q(i,k-1) )/bet + gam(i,k) = d4(i) / bet + enddo + enddo + + do i=i1,i2 + a_bot = 1. + d4(i)*(d4(i)+1.5) + q(i,km+1) = (2.*d4(i)*(d4(i)+1.)*a4(1,i,km)+a4(1,i,km-1)-a_bot*q(i,km)) & + / ( d4(i)*(d4(i)+0.5) - a_bot*gam(i,km) ) + enddo + + do k=km,1,-1 + do i=i1,i2 + q(i,k) = q(i,k) - gam(i,k)*q(i,k+1) + enddo + enddo + endif +!----- Perfectly linear scheme -------------------------------- + if ( abs(kord) > 16 ) then + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + enddo + return + endif +!----- Perfectly linear scheme -------------------------------- + +!------------------ +! Apply constraints +!------------------ + im = i2 - i1 + 1 + +! Apply *large-scale* constraints + do i=i1,i2 + q(i,2) = min( q(i,2), max(a4(1,i,1), a4(1,i,2)) ) + q(i,2) = max( q(i,2), min(a4(1,i,1), a4(1,i,2)) ) + enddo + + do k=2,km + do i=i1,i2 + gam(i,k) = a4(1,i,k) - a4(1,i,k-1) + enddo + enddo + +! Interior: + do k=3,km-1 + do i=i1,i2 + if ( gam(i,k-1)*gam(i,k+1)>0. ) then +! Apply large-scale constraint to ALL fields if not local max/min + q(i,k) = min( q(i,k), max(a4(1,i,k-1),a4(1,i,k)) ) + q(i,k) = max( q(i,k), min(a4(1,i,k-1),a4(1,i,k)) ) + else + if ( gam(i,k-1) > 0. ) then +! There exists a local max + q(i,k) = max(q(i,k), min(a4(1,i,k-1),a4(1,i,k))) + else +! There exists a local min + q(i,k) = min(q(i,k), max(a4(1,i,k-1),a4(1,i,k))) + if ( iv==0 ) q(i,k) = max(0., q(i,k)) + endif + endif + enddo + enddo + +! Bottom: + do i=i1,i2 + q(i,km) = min( q(i,km), max(a4(1,i,km-1), a4(1,i,km)) ) + q(i,km) = max( q(i,km), min(a4(1,i,km-1), a4(1,i,km)) ) + enddo + + do k=1,km + do i=i1,i2 + a4(2,i,k) = q(i,k ) + a4(3,i,k) = q(i,k+1) + enddo + enddo + + do k=1,km + if ( k==1 .or. k==km ) then + do i=i1,i2 + extm(i,k) = (a4(2,i,k)-a4(1,i,k)) * (a4(3,i,k)-a4(1,i,k)) > 0. + enddo + else + do i=i1,i2 + extm(i,k) = gam(i,k)*gam(i,k+1) < 0. + enddo + endif + if ( abs(kord) > 9 ) then + do i=i1,i2 + x0 = 2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k)) + x1 = abs(a4(2,i,k)-a4(3,i,k)) + a4(4,i,k) = 3.*x0 + ext5(i,k) = abs(x0) > x1 + ext6(i,k) = abs(a4(4,i,k)) > x1 + enddo + endif + enddo + +!--------------------------- +! Apply subgrid constraints: +!--------------------------- +! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) +! Top 2 and bottom 2 layers always use monotonic mapping + + if ( iv==0 ) then + do i=i1,i2 + a4(2,i,1) = max(0., a4(2,i,1)) + enddo + elseif ( iv==-1 ) then + do i=i1,i2 + if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. + enddo + elseif ( iv==2 ) then + do i=i1,i2 + a4(2,i,1) = a4(1,i,1) + a4(3,i,1) = a4(1,i,1) + a4(4,i,1) = 0. + enddo + endif + + if ( iv/=2 ) then + do i=i1,i2 + a4(4,i,1) = 3.*(2.*a4(1,i,1) - (a4(2,i,1)+a4(3,i,1))) + enddo + call cs_limiters(im, extm(i1,1), a4(1,i1,1), 1) + endif + +! k=2 + do i=i1,i2 + a4(4,i,2) = 3.*(2.*a4(1,i,2) - (a4(2,i,2)+a4(3,i,2))) + enddo + call cs_limiters(im, extm(i1,2), a4(1,i1,2), 2) + +!------------------------------------- +! Huynh's 2nd constraint for interior: +!------------------------------------- + do k=3,km-2 + if ( abs(kord)<9 ) then + do i=i1,i2 +! Left edges + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) +! Right edges + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==9 ) then + do i=i1,i2 + if ( extm(i,k) .and. extm(i,k-1) ) then ! c90_mp122 +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else if ( extm(i,k) .and. extm(i,k+1) ) then ! c90_mp122 +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) +! Check within the smooth region if subgrid profile is non-monotonic + if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) + endif + endif + enddo + elseif ( abs(kord)==10 ) then + do i=i1,i2 + if( ext5(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + elseif( ext6(i,k) ) then + if( ext5(i,k-1) .or. ext5(i,k+1) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==12 ) then + do i=i1,i2 + if( extm(i,k) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else ! not a local extremum + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) +! Check within the smooth region if subgrid profile is non-monotonic + if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) + endif + endif + enddo + elseif ( abs(kord)==13 ) then + do i=i1,i2 + if( ext6(i,k) ) then + if ( ext6(i,k-1) .and. ext6(i,k+1) ) then +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==14 ) then + + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + + elseif ( abs(kord)==15 ) then ! revised kord=9 scehem + do i=i1,i2 + if ( ext5(i,k) ) then ! c90_mp122 + if ( ext5(i,k-1) .or. ext5(i,k+1) ) then ! c90_mp122 +! grid-scale 2-delta-z wave detected + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + endif + elseif( ext6(i,k) ) then +! Check within the smooth region if subgrid profile is non-monotonic + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + elseif ( abs(kord)==16 ) then + do i=i1,i2 + if( ext5(i,k) ) then + if ( ext5(i,k-1) .or. ext5(i,k+1) ) then + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then + ! Left edges + pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) + lac_1 = pmp_1 + 1.5*gam(i,k+2) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & + max(a4(1,i,k), pmp_1, lac_1) ) + ! Right edges + pmp_2 = a4(1,i,k) + 2.*gam(i,k) + lac_2 = pmp_2 - 1.5*gam(i,k-1) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & + max(a4(1,i,k), pmp_2, lac_2) ) + endif + endif + enddo + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + else ! kord = 11 + do i=i1,i2 + if ( ext5(i,k) .and. (ext5(i,k-1) .or. ext5(i,k+1)) ) then +! Noisy region: + a4(2,i,k) = a4(1,i,k) + a4(3,i,k) = a4(1,i,k) + a4(4,i,k) = 0. + else + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + endif + enddo + endif + +! Additional constraint to ensure positivity + if ( iv==0 ) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 0) + + enddo ! k-loop + +!---------------------------------- +! Bottom layer subgrid constraints: +!---------------------------------- + if ( iv==0 ) then + do i=i1,i2 + a4(3,i,km) = max(0., a4(3,i,km)) + enddo + elseif ( iv .eq. -1 ) then + do i=i1,i2 + if ( a4(3,i,km)*a4(1,i,km) <= 0. ) a4(3,i,km) = 0. + enddo + endif + + do k=km-1,km + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + if(k==(km-1)) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 2) + if(k== km ) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 1) + enddo + + end subroutine cs_profile + + + subroutine cs_limiters(im, extm, a4, iv) + integer, intent(in) :: im + integer, intent(in) :: iv + logical, intent(in) :: extm(im) + real , intent(inout) :: a4(4,im) !< PPM array +! LOCAL VARIABLES: + real da1, da2, a6da + integer i + + if ( iv==0 ) then +! Positive definite constraint + do i=1,im + if( a4(1,i)<=0.) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + if( abs(a4(3,i)-a4(2,i)) < -a4(4,i) ) then + if( (a4(1,i)+0.25*(a4(3,i)-a4(2,i))**2/a4(4,i)+a4(4,i)*r12) < 0. ) then +! local minimum is negative + if( a4(1,i) a4(2,i) ) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + else + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + endif + endif + enddo + elseif ( iv==1 ) then + do i=1,im + if( (a4(1,i)-a4(2,i))*(a4(1,i)-a4(3,i))>=0. ) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + da1 = a4(3,i) - a4(2,i) + da2 = da1**2 + a6da = a4(4,i)*da1 + if(a6da < -da2) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + elseif(a6da > da2) then + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + enddo + else +! Standard PPM constraint + do i=1,im + if( extm(i) ) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + da1 = a4(3,i) - a4(2,i) + da2 = da1**2 + a6da = a4(4,i)*da1 + if(a6da < -da2) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + elseif(a6da > da2) then + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + enddo + endif + end subroutine cs_limiters + + + + subroutine ppm_profile(a4, delp, km, i1, i2, iv, kord) + +! INPUT PARAMETERS: + integer, intent(in):: iv !< iv =-1: winds iv = 0: positive definite scalars iv = 1: others iv = 2: temp (if remap_t) and w (iv=-2) + integer, intent(in):: i1 !< Starting longitude + integer, intent(in):: i2 !< Finishing longitude + integer, intent(in):: km !< Vertical dimension + integer, intent(in):: kord !< Order (or more accurately method no.): + ! + real , intent(in):: delp(i1:i2,km) !< Layer pressure thickness + +! !INPUT/OUTPUT PARAMETERS: + real , intent(inout):: a4(4,i1:i2,km) !< Interpolated values + +! DESCRIPTION: +! +! Perform the piecewise parabolic reconstruction +! +! !REVISION HISTORY: +! S.-J. Lin revised at GFDL 2007 +!----------------------------------------------------------------------- +! local arrays: + real dc(i1:i2,km) + real h2(i1:i2,km) + real delq(i1:i2,km) + real df2(i1:i2,km) + real d4(i1:i2,km) + +! local scalars: + integer i, k, km1, lmt, it + real fac + real a1, a2, c1, c2, c3, d1, d2 + real qm, dq, lac, qmp, pmp + + km1 = km - 1 + it = i2 - i1 + 1 + + do k=2,km + do i=i1,i2 + delq(i,k-1) = a4(1,i,k) - a4(1,i,k-1) + d4(i,k ) = delp(i,k-1) + delp(i,k) + enddo + enddo + + do k=2,km1 + do i=i1,i2 + c1 = (delp(i,k-1)+0.5*delp(i,k))/d4(i,k+1) + c2 = (delp(i,k+1)+0.5*delp(i,k))/d4(i,k) + df2(i,k) = delp(i,k)*(c1*delq(i,k) + c2*delq(i,k-1)) / & + (d4(i,k)+delp(i,k+1)) + dc(i,k) = sign( min(abs(df2(i,k)), & + max(a4(1,i,k-1),a4(1,i,k),a4(1,i,k+1))-a4(1,i,k), & + a4(1,i,k)-min(a4(1,i,k-1),a4(1,i,k),a4(1,i,k+1))), df2(i,k) ) + enddo + enddo + +!----------------------------------------------------------- +! 4th order interpolation of the provisional cell edge value +!----------------------------------------------------------- + + do k=3,km1 + do i=i1,i2 + c1 = delq(i,k-1)*delp(i,k-1) / d4(i,k) + a1 = d4(i,k-1) / (d4(i,k) + delp(i,k-1)) + a2 = d4(i,k+1) / (d4(i,k) + delp(i,k)) + a4(2,i,k) = a4(1,i,k-1) + c1 + 2./(d4(i,k-1)+d4(i,k+1)) * & + ( delp(i,k)*(c1*(a1 - a2)+a2*dc(i,k-1)) - & + delp(i,k-1)*a1*dc(i,k ) ) + enddo + enddo + +! if(km>8 .and. kord>4) call steepz(i1, i2, km, a4, df2, dc, delq, delp, d4) + +! Area preserving cubic with 2nd deriv. = 0 at the boundaries +! Top + do i=i1,i2 + d1 = delp(i,1) + d2 = delp(i,2) + qm = (d2*a4(1,i,1)+d1*a4(1,i,2)) / (d1+d2) + dq = 2.*(a4(1,i,2)-a4(1,i,1)) / (d1+d2) + c1 = 4.*(a4(2,i,3)-qm-d2*dq) / ( d2*(2.*d2*d2+d1*(d2+3.*d1)) ) + c3 = dq - 0.5*c1*(d2*(5.*d1+d2)-3.*d1*d1) + a4(2,i,2) = qm - 0.25*c1*d1*d2*(d2+3.*d1) +! Top edge: +!------------------------------------------------------- + a4(2,i,1) = d1*(2.*c1*d1**2-c3) + a4(2,i,2) +!------------------------------------------------------- +! a4(2,i,1) = (12./7.)*a4(1,i,1)-(13./14.)*a4(1,i,2)+(3./14.)*a4(1,i,3) +!------------------------------------------------------- +! No over- and undershoot condition + a4(2,i,2) = max( a4(2,i,2), min(a4(1,i,1), a4(1,i,2)) ) + a4(2,i,2) = min( a4(2,i,2), max(a4(1,i,1), a4(1,i,2)) ) + dc(i,1) = 0.5*(a4(2,i,2) - a4(1,i,1)) + enddo + +! Enforce monotonicity within the top layer + + if( iv==0 ) then + do i=i1,i2 + a4(2,i,1) = max(0., a4(2,i,1)) + a4(2,i,2) = max(0., a4(2,i,2)) + enddo + elseif( iv==-1 ) then + do i=i1,i2 + if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. + enddo + elseif( abs(iv)==2 ) then + do i=i1,i2 + a4(2,i,1) = a4(1,i,1) + a4(3,i,1) = a4(1,i,1) + enddo + endif + +! Bottom +! Area preserving cubic with 2nd deriv. = 0 at the surface + do i=i1,i2 + d1 = delp(i,km) + d2 = delp(i,km1) + qm = (d2*a4(1,i,km)+d1*a4(1,i,km1)) / (d1+d2) + dq = 2.*(a4(1,i,km1)-a4(1,i,km)) / (d1+d2) + c1 = (a4(2,i,km1)-qm-d2*dq) / (d2*(2.*d2*d2+d1*(d2+3.*d1))) + c3 = dq - 2.0*c1*(d2*(5.*d1+d2)-3.*d1*d1) + a4(2,i,km) = qm - c1*d1*d2*(d2+3.*d1) +! Bottom edge: +!----------------------------------------------------- + a4(3,i,km) = d1*(8.*c1*d1**2-c3) + a4(2,i,km) +! dc(i,km) = 0.5*(a4(3,i,km) - a4(1,i,km)) +!----------------------------------------------------- +! a4(3,i,km) = (12./7.)*a4(1,i,km)-(13./14.)*a4(1,i,km-1)+(3./14.)*a4(1,i,km-2) +! No over- and under-shoot condition + a4(2,i,km) = max( a4(2,i,km), min(a4(1,i,km), a4(1,i,km1)) ) + a4(2,i,km) = min( a4(2,i,km), max(a4(1,i,km), a4(1,i,km1)) ) + dc(i,km) = 0.5*(a4(1,i,km) - a4(2,i,km)) + enddo + + +! Enforce constraint on the "slope" at the surface + +#ifdef BOT_MONO + do i=i1,i2 + a4(4,i,km) = 0 + if( a4(3,i,km) * a4(1,i,km) <= 0. ) a4(3,i,km) = 0. + d1 = a4(1,i,km) - a4(2,i,km) + d2 = a4(3,i,km) - a4(1,i,km) + if ( d1*d2 < 0. ) then + a4(2,i,km) = a4(1,i,km) + a4(3,i,km) = a4(1,i,km) + else + dq = sign(min(abs(d1),abs(d2),0.5*abs(delq(i,km-1))), d1) + a4(2,i,km) = a4(1,i,km) - dq + a4(3,i,km) = a4(1,i,km) + dq + endif + enddo +#else + if( iv==0 ) then + do i=i1,i2 + a4(2,i,km) = max(0.,a4(2,i,km)) + a4(3,i,km) = max(0.,a4(3,i,km)) + enddo + elseif( iv<0 ) then + do i=i1,i2 + if( a4(1,i,km)*a4(3,i,km) <= 0. ) a4(3,i,km) = 0. + enddo + endif +#endif + + do k=1,km1 + do i=i1,i2 + a4(3,i,k) = a4(2,i,k+1) + enddo + enddo + +!----------------------------------------------------------- +! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) +!----------------------------------------------------------- +! Top 2 and bottom 2 layers always use monotonic mapping + do k=1,2 + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 0) + enddo + + if(kord >= 7) then +!----------------------- +! Huynh's 2nd constraint +!----------------------- + do k=2,km1 + do i=i1,i2 +! Method#1 +! h2(i,k) = delq(i,k) - delq(i,k-1) +! Method#2 - better + h2(i,k) = 2.*(dc(i,k+1)/delp(i,k+1) - dc(i,k-1)/delp(i,k-1)) & + / ( delp(i,k)+0.5*(delp(i,k-1)+delp(i,k+1)) ) & + * delp(i,k)**2 +! Method#3 +!!! h2(i,k) = dc(i,k+1) - dc(i,k-1) + enddo + enddo + + fac = 1.5 ! original quasi-monotone + + do k=3,km-2 + do i=i1,i2 +! Right edges +! qmp = a4(1,i,k) + 2.0*delq(i,k-1) +! lac = a4(1,i,k) + fac*h2(i,k-1) + 0.5*delq(i,k-1) +! + pmp = 2.*dc(i,k) + qmp = a4(1,i,k) + pmp + lac = a4(1,i,k) + fac*h2(i,k-1) + dc(i,k) + a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), qmp, lac)), & + max(a4(1,i,k), qmp, lac) ) +! Left edges +! qmp = a4(1,i,k) - 2.0*delq(i,k) +! lac = a4(1,i,k) + fac*h2(i,k+1) - 0.5*delq(i,k) +! + qmp = a4(1,i,k) - pmp + lac = a4(1,i,k) + fac*h2(i,k+1) - dc(i,k) + a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), qmp, lac)), & + max(a4(1,i,k), qmp, lac)) +!------------- +! Recompute A6 +!------------- + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo +! Additional constraint to ensure positivity when kord=7 + if (iv == 0 .and. kord >= 6 ) & + call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 2) + enddo + + else + + lmt = kord - 3 + lmt = max(0, lmt) + if (iv == 0) lmt = min(2, lmt) + + do k=3,km-2 + if( kord /= 4) then + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + endif + if(kord/=6) call ppm_limiters(dc(i1,k), a4(1,i1,k), it, lmt) + enddo + endif + + do k=km1,km + do i=i1,i2 + a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) + enddo + call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 0) + enddo + + end subroutine ppm_profile + + + subroutine ppm_limiters(dm, a4, itot, lmt) + +! INPUT PARAMETERS: + real , intent(in):: dm(*) !< Linear slope + integer, intent(in) :: itot !< Total Longitudes + integer, intent(in) :: lmt !< 0: Standard PPM constraint 1: Improved full monotonicity constraint + !< (Lin) 2: Positive definite constraint + !< 3: do nothing (return immediately) +! INPUT/OUTPUT PARAMETERS: + real , intent(inout) :: a4(4,*) !< PPM array AA <-- a4(1,i) AL <-- a4(2,i) AR <-- a4(3,i) A6 <-- a4(4,i) +! LOCAL VARIABLES: + real qmp + real da1, da2, a6da + real fmin + integer i + +! Developer: S.-J. Lin + + if ( lmt == 3 ) return + + if(lmt == 0) then +! Standard PPM constraint + do i=1,itot + if(dm(i) == 0.) then + a4(2,i) = a4(1,i) + a4(3,i) = a4(1,i) + a4(4,i) = 0. + else + da1 = a4(3,i) - a4(2,i) + da2 = da1**2 + a6da = a4(4,i)*da1 + if(a6da < -da2) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + elseif(a6da > da2) then + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + enddo + + elseif (lmt == 1) then + +! Improved full monotonicity constraint (Lin 2004) +! Note: no need to provide first guess of A6 <-- a4(4,i) + do i=1, itot + qmp = 2.*dm(i) + a4(2,i) = a4(1,i)-sign(min(abs(qmp),abs(a4(2,i)-a4(1,i))), qmp) + a4(3,i) = a4(1,i)+sign(min(abs(qmp),abs(a4(3,i)-a4(1,i))), qmp) + a4(4,i) = 3.*( 2.*a4(1,i) - (a4(2,i)+a4(3,i)) ) + enddo + + elseif (lmt == 2) then + +! Positive definite constraint + do i=1,itot + if( abs(a4(3,i)-a4(2,i)) < -a4(4,i) ) then + fmin = a4(1,i)+0.25*(a4(3,i)-a4(2,i))**2/a4(4,i)+a4(4,i)*r12 + if( fmin < 0. ) then + if(a4(1,i) a4(2,i)) then + a4(4,i) = 3.*(a4(2,i)-a4(1,i)) + a4(3,i) = a4(2,i) - a4(4,i) + else + a4(4,i) = 3.*(a4(3,i)-a4(1,i)) + a4(2,i) = a4(3,i) - a4(4,i) + endif + endif + endif + enddo + + endif + + end subroutine ppm_limiters + + + + subroutine steepz(i1, i2, km, a4, df2, dm, dq, dp, d4) + integer, intent(in) :: km, i1, i2 + real , intent(in) :: dp(i1:i2,km) !< Grid size + real , intent(in) :: dq(i1:i2,km) !< Backward diff of q + real , intent(in) :: d4(i1:i2,km) !< Backward sum: dp(k)+ dp(k-1) + real , intent(in) :: df2(i1:i2,km) !< First guess mismatch + real , intent(in) :: dm(i1:i2,km) !< Monotonic mismatch +! INPUT/OUTPUT PARAMETERS: + real , intent(inout) :: a4(4,i1:i2,km) !@brief The subroutine 'rst_remap' remaps all variables required for a restart. +!>@details npz_restart /= npz (i.e., when the number of vertical levels is +!! changed at restart) + subroutine rst_remap(km, kn, is,ie,js,je, isd,ied,jsd,jed, nq, ntp, & + delp_r, u_r, v_r, w_r, delz_r, pt_r, q_r, qdiag_r, & + delp, u, v, w, delz, pt, q, qdiag, & + ak_r, bk_r, ptop, ak, bk, hydrostatic, make_nh, & + domain, square_domain) +!------------------------------------ +! Assuming hybrid sigma-P coordinate: +!------------------------------------ +! INPUT PARAMETERS: + integer, intent(in):: km !< Restart z-dimension + integer, intent(in):: kn !< Run time dimension + integer, intent(in):: nq, ntp !< Number of tracers (including H2O) + integer, intent(in):: is,ie,isd,ied !< Starting & ending X-Dir index + integer, intent(in):: js,je,jsd,jed !< Starting & ending Y-Dir index + logical, intent(in):: hydrostatic, make_nh, square_domain + real, intent(IN) :: ptop + real, intent(in) :: ak_r(km+1) + real, intent(in) :: bk_r(km+1) + real, intent(in) :: ak(kn+1) + real, intent(in) :: bk(kn+1) + real, intent(in):: delp_r(is:ie,js:je,km) !< Pressure thickness + real, intent(in):: u_r(is:ie, js:je+1,km) !< u-wind (m/s) + real, intent(in):: v_r(is:ie+1,js:je ,km) !< v-wind (m/s) + real, intent(inout):: pt_r(is:ie,js:je,km) + real, intent(in):: w_r(is:ie,js:je,km) + real, intent(in):: q_r(is:ie,js:je,km,1:ntp) + real, intent(in):: qdiag_r(is:ie,js:je,km,ntp+1:nq) + real, intent(inout)::delz_r(is:ie,js:je,km) + type(domain2d), intent(INOUT) :: domain +! Output: + real, intent(out):: delp(isd:ied,jsd:jed,kn) !< Pressure thickness + real, intent(out):: u(isd:ied ,jsd:jed+1,kn) !< u-wind (m/s) + real, intent(out):: v(isd:ied+1,jsd:jed ,kn) !< v-wind (m/s) + real, intent(out):: w(isd: ,jsd: ,1:) !< Vertical velocity (m/s) + real, intent(out):: pt(isd:ied ,jsd:jed ,kn) !< Temperature + real, intent(out):: q(isd:ied,jsd:jed,kn,1:ntp) + real, intent(out):: qdiag(isd:ied,jsd:jed,kn,ntp+1:nq) + real, intent(out):: delz(isd:,jsd:,1:) !< Delta-height (m) +!----------------------------------------------------------------------- + real r_vir, rgrav + real ps(isd:ied,jsd:jed) !< Surface pressure + real pe1(is:ie,km+1) + real pe2(is:ie,kn+1) + real pv1(is:ie+1,km+1) + real pv2(is:ie+1,kn+1) + + integer i,j,k , iq + integer, parameter:: kord=4 + +#ifdef HYDRO_DELZ_REMAP + if (is_master() .and. .not. hydrostatic) then + print*, '' + print*, ' REMAPPING IC: INITIALIZING DELZ WITH HYDROSTATIC STATE ' + print*, '' + endif +#endif + +#ifdef HYDRO_DELZ_EXTRAP + if (is_master() .and. .not. hydrostatic) then + print*, '' + print*, ' REMAPPING IC: INITIALIZING DELZ WITH HYDROSTATIC STATE ABOVE INPUT MODEL TOP ' + print*, '' + endif +#endif + +#ifdef ZERO_W_EXTRAP + if (is_master() .and. .not. hydrostatic) then + print*, '' + print*, ' REMAPPING IC: INITIALIZING W TO ZERO ABOVE INPUT MODEL TOP ' + print*, '' + endif +#endif + + r_vir = rvgas/rdgas - 1. + rgrav = 1./grav + +!$OMP parallel do default(none) shared(is,ie,js,je,ps,ak_r) + do j=js,je + do i=is,ie + ps(i,j) = ak_r(1) + enddo + enddo + +! this OpenMP do-loop setup cannot work in it's current form.... +!$OMP parallel do default(none) shared(is,ie,js,je,km,ps,delp_r) + do j=js,je + do k=1,km + do i=is,ie + ps(i,j) = ps(i,j) + delp_r(i,j,k) + enddo + enddo + enddo + +! only one cell is needed + if ( square_domain ) then + call mpp_update_domains(ps, domain, whalo=1, ehalo=1, shalo=1, nhalo=1, complete=.true.) + else + call mpp_update_domains(ps, domain, complete=.true.) + endif + +! Compute virtual Temp +!$OMP parallel do default(none) shared(is,ie,js,je,km,pt_r,r_vir,q_r) + do k=1,km + do j=js,je + do i=is,ie + pt_r(i,j,k) = pt_r(i,j,k) * (1.+r_vir*q_r(i,j,k,1)) + enddo + enddo + enddo + +!$OMP parallel do default(none) shared(is,ie,js,je,km,ak_r,bk_r,ps,kn,ak,bk,u_r,u,delp, & +!$OMP ntp,nq,hydrostatic,make_nh,w_r,w,delz_r,delp_r,delz, & +!$OMP pt_r,pt,v_r,v,q,q_r,qdiag,qdiag_r) & +!$OMP private(pe1, pe2, pv1, pv2) + do 6000 j=js,je+1 +!------ +! map u +!------ + do k=1,km+1 + do i=is,ie + pe1(i,k) = ak_r(k) + 0.5*bk_r(k)*(ps(i,j-1)+ps(i,j)) + enddo + enddo + + do k=1,kn+1 + do i=is,ie + pe2(i,k) = ak(k) + 0.5*bk(k)*(ps(i,j-1)+ps(i,j)) + enddo + enddo + + call remap_2d(km, pe1, u_r(is:ie,j:j,1:km), & + kn, pe2, u(is:ie,j:j,1:kn), & + is, ie, -1, kord) + + if ( j /= (je+1) ) then + +!--------------- +! Hybrid sigma-p +!--------------- + do k=1,km+1 + do i=is,ie + pe1(i,k) = ak_r(k) + bk_r(k)*ps(i,j) + enddo + enddo + + do k=1,kn+1 + do i=is,ie + pe2(i,k) = ak(k) + bk(k)*ps(i,j) + enddo + enddo + +!------------- +! Compute delp +!------------- + do k=1,kn + do i=is,ie + delp(i,j,k) = pe2(i,k+1) - pe2(i,k) + enddo + enddo + +!---------------- +! Map constituents +!---------------- + if( nq /= 0 ) then + do iq=1,ntp + call remap_2d(km, pe1, q_r(is:ie,j:j,1:km,iq:iq), & + kn, pe2, q(is:ie,j:j,1:kn,iq:iq), & + is, ie, 0, kord) + enddo + do iq=ntp+1,nq + call remap_2d(km, pe1, qdiag_r(is:ie,j:j,1:km,iq:iq), & + kn, pe2, qdiag(is:ie,j:j,1:kn,iq:iq), & + is, ie, 0, kord) + enddo + endif + + if ( .not. hydrostatic .and. .not. make_nh) then +! Remap vertical wind: + call remap_2d(km, pe1, w_r(is:ie,j:j,1:km), & + kn, pe2, w(is:ie,j:j,1:kn), & + is, ie, -1, kord) + +#ifdef ZERO_W_EXTRAP + do k=1,kn + do i=is,ie + if (pe2(i,k) < pe1(i,1)) then + w(i,j,k) = 0. + endif + enddo + enddo +#endif + +#ifndef HYDRO_DELZ_REMAP +! Remap delz for hybrid sigma-p coordinate + do k=1,km + do i=is,ie + delz_r(i,j,k) = -delz_r(i,j,k)/delp_r(i,j,k) ! ="specific volume"/grav + enddo + enddo + call remap_2d(km, pe1, delz_r(is:ie,j:j,1:km), & + kn, pe2, delz(is:ie,j:j,1:kn), & + is, ie, 1, kord) + do k=1,kn + do i=is,ie + delz(i,j,k) = -delz(i,j,k)*delp(i,j,k) + enddo + enddo +#endif + endif + +! Geopotential conserving remap of virtual temperature: + do k=1,km+1 + do i=is,ie + pe1(i,k) = log(pe1(i,k)) + enddo + enddo + do k=1,kn+1 + do i=is,ie + pe2(i,k) = log(pe2(i,k)) + enddo + enddo + + call remap_2d(km, pe1, pt_r(is:ie,j:j,1:km), & + kn, pe2, pt(is:ie,j:j,1:kn), & + is, ie, 1, kord) + +#ifdef HYDRO_DELZ_REMAP + !initialize delz from the hydrostatic state + do k=1,kn + do i=is,ie + delz(i,j,k) = (rdgas*rgrav)*pt(i,j,k)*(pe2(i,k)-pe2(i,k+1)) + enddo + enddo +#endif +#ifdef HYDRO_DELZ_EXTRAP + !initialize delz from the hydrostatic state + do k=1,kn + do i=is,ie + if (pe2(i,k) < pe1(i,1)) then + delz(i,j,k) = (rdgas*rgrav)*pt(i,j,k)*(pe2(i,k)-pe2(i,k+1)) + endif + enddo + enddo +#endif +!------ +! map v +!------ + do k=1,km+1 + do i=is,ie+1 + pv1(i,k) = ak_r(k) + 0.5*bk_r(k)*(ps(i-1,j)+ps(i,j)) + enddo + enddo + do k=1,kn+1 + do i=is,ie+1 + pv2(i,k) = ak(k) + 0.5*bk(k)*(ps(i-1,j)+ps(i,j)) + enddo + enddo + + call remap_2d(km, pv1, v_r(is:ie+1,j:j,1:km), & + kn, pv2, v(is:ie+1,j:j,1:kn), & + is, ie+1, -1, kord) + + endif !(j < je+1) +6000 continue + +!$OMP parallel do default(none) shared(is,ie,js,je,kn,pt,r_vir,q) + do k=1,kn + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k) / (1.+r_vir*q(i,j,k,1)) + enddo + enddo + enddo + + end subroutine rst_remap + +!>@brief The subroutine 'mappm' is a general-purpose routine for remapping +!! one set of vertical levels to another. + subroutine mappm(km, pe1, q1, kn, pe2, q2, i1, i2, iv, kord, ptop) + +! IV = 0: constituents +! IV = 1: potential temp +! IV =-1: winds + +! Mass flux preserving mapping: q1(im,km) -> q2(im,kn) + +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate + + integer, intent(in):: i1, i2, km, kn, kord, iv + real, intent(in ):: pe1(i1:i2,km+1), pe2(i1:i2,kn+1) !< pe1: pressure at layer edges from model top to bottom + !! surface in the ORIGINAL vertical coordinate + !< pe2: pressure at layer edges from model top to bottom + !! surface in the NEW vertical coordinate +! Mass flux preserving mapping: q1(im,km) -> q2(im,kn) + real, intent(in ):: q1(i1:i2,km) + real, intent(out):: q2(i1:i2,kn) + real, intent(IN) :: ptop +! local + real qs(i1:i2) + real dp1(i1:i2,km) + real a4(4,i1:i2,km) + integer i, k, l + integer k0, k1 + real pl, pr, tt, delp, qsum, dpsum, esl + + do k=1,km + do i=i1,i2 + dp1(i,k) = pe1(i,k+1) - pe1(i,k) + a4(1,i,k) = q1(i,k) + enddo + enddo + + if ( kord >7 ) then + call cs_profile( qs, a4, dp1, km, i1, i2, iv, kord ) + else + call ppm_profile( a4, dp1, km, i1, i2, iv, kord ) + endif + +!------------------------------------ +! Lowest layer: constant distribution +!------------------------------------ +#ifdef NGGPS_SUBMITTED + do i=i1,i2 + a4(2,i,km) = q1(i,km) + a4(3,i,km) = q1(i,km) + a4(4,i,km) = 0. + enddo +#endif + + do 5555 i=i1,i2 + k0 = 1 + do 555 k=1,kn + + if(pe2(i,k) .le. pe1(i,1)) then +! above old ptop + q2(i,k) = q1(i,1) + elseif(pe2(i,k) .ge. pe1(i,km+1)) then +! Entire grid below old ps +#ifdef NGGPS_SUBMITTED + q2(i,k) = a4(3,i,km) ! this is not good. +#else + q2(i,k) = q1(i,km) +#endif + else + + do 45 L=k0,km +! locate the top edge at pe2(i,k) + if( pe2(i,k) .ge. pe1(i,L) .and. & + pe2(i,k) .le. pe1(i,L+1) ) then + k0 = L + PL = (pe2(i,k)-pe1(i,L)) / dp1(i,L) + if(pe2(i,k+1) .le. pe1(i,L+1)) then + +! entire new grid is within the original grid + PR = (pe2(i,k+1)-pe1(i,L)) / dp1(i,L) + TT = r3*(PR*(PR+PL)+PL**2) + q2(i,k) = a4(2,i,L) + 0.5*(a4(4,i,L)+a4(3,i,L) & + - a4(2,i,L))*(PR+PL) - a4(4,i,L)*TT + goto 555 + else +! Fractional area... + delp = pe1(i,L+1) - pe2(i,k) + TT = r3*(1.+PL*(1.+PL)) + qsum = delp*(a4(2,i,L)+0.5*(a4(4,i,L)+ & + a4(3,i,L)-a4(2,i,L))*(1.+PL)-a4(4,i,L)*TT) + dpsum = delp + k1 = L + 1 + goto 111 + endif + endif +45 continue + +111 continue + do 55 L=k1,km + if( pe2(i,k+1) .gt. pe1(i,L+1) ) then + +! Whole layer.. + + qsum = qsum + dp1(i,L)*q1(i,L) + dpsum = dpsum + dp1(i,L) + else + delp = pe2(i,k+1)-pe1(i,L) + esl = delp / dp1(i,L) + qsum = qsum + delp * (a4(2,i,L)+0.5*esl* & + (a4(3,i,L)-a4(2,i,L)+a4(4,i,L)*(1.-r23*esl)) ) + dpsum = dpsum + delp + k0 = L + goto 123 + endif +55 continue + delp = pe2(i,k+1) - pe1(i,km+1) + if(delp > 0.) then +! Extended below old ps +#ifdef NGGPS_SUBMITTED + qsum = qsum + delp * a4(3,i,km) ! not good. +#else + qsum = qsum + delp * q1(i,km) +#endif + dpsum = dpsum + delp + endif +123 q2(i,k) = qsum / dpsum + endif +555 continue +5555 continue + + end subroutine mappm + +!>@brief The subroutine 'moist_cv' computes the FV3-consistent moist heat capacity under constant volume, +!! including the heating capacity of water vapor and condensates. +!>@details See \cite emanuel1994atmospheric for information on variable heat capacities. + subroutine moist_cv(is,ie, isd,ied, jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, qd, cvm, t1) + integer, intent(in):: is, ie, isd,ied, jsd,jed, km, nwat, j, k + integer, intent(in):: sphum, liq_wat, rainwat, ice_wat, snowwat, graupel + real, intent(in), dimension(isd:ied,jsd:jed,km,nwat):: q + real, intent(out), dimension(is:ie):: cvm, qd + real, intent(in), optional:: t1(is:ie) +! + real, parameter:: t_i0 = 15. + real, dimension(is:ie):: qv, ql, qs + integer:: i + + select case (nwat) + + case(1) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + cvm(i) = (1.-qv(i))*cv_air + qv(i)*cv_vap + enddo + case(2) + if ( present(t1) ) then ! Special case for GFS physics + do i=is,ie + qd(i) = max(0., q(i,j,k,liq_wat)) + if ( t1(i) > tice ) then + qs(i) = 0. + elseif ( t1(i) < tice-t_i0 ) then + qs(i) = qd(i) + else + qs(i) = qd(i)*(tice-t1(i))/t_i0 + endif + ql(i) = qd(i) - qs(i) + qv(i) = max(0.,q(i,j,k,sphum)) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + else + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + qs(i) = max(q(i,j,k,liq_wat),0.0) + qd(i) = qs(i) + cvm(i) = (1.-qv(i))*cv_air + qv(i)*cv_vap + enddo + endif + case (3) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + qs(i) = max(q(i,j,k,ice_wat),0.0) + qd(i) = ql(i) + qs(i) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + case(4) ! K_warm_rain with fake ice + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + qd(i)*c_liq + enddo + case(5) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) + qs(i) = max(q(i,j,k,ice_wat),0.0) + max(q(i,j,k,snowwat),0.0) + qd(i) = ql(i) + qs(i) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + case(6:7) + do i=is,ie + qv(i) = max(q(i,j,k,sphum) ,0.0) + ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) + qs(i) = max(q(i,j,k,ice_wat),0.0) + max(q(i,j,k,snowwat),0.0) + max(q(i,j,k,graupel),0.0) + qd(i) = ql(i) + qs(i) + cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice + enddo + case default + !call mpp_error (NOTE, 'fv_mapz::moist_cv - using default cv_air') + do i=is,ie + qd(i) = 0. + cvm(i) = cv_air + enddo + end select + + end subroutine moist_cv + +!>@brief The subroutine 'moist_cp' computes the FV3-consistent moist heat capacity under constant pressure, +!! including the heating capacity of water vapor and condensates. + subroutine moist_cp(is,ie, isd,ied, jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, qd, cpm, t1) + + integer, intent(in):: is, ie, isd,ied, jsd,jed, km, nwat, j, k + integer, intent(in):: sphum, liq_wat, rainwat, ice_wat, snowwat, graupel + real, intent(in), dimension(isd:ied,jsd:jed,km,nwat):: q + real, intent(out), dimension(is:ie):: cpm, qd + real, intent(in), optional:: t1(is:ie) +! + real, parameter:: t_i0 = 15. + real, dimension(is:ie):: qv, ql, qs + integer:: i + + select case (nwat) + + case(2) + if ( present(t1) ) then ! Special case for GFS physics + do i=is,ie + qd(i) = max(0., q(i,j,k,liq_wat)) + if ( t1(i) > tice ) then + qs(i) = 0. + elseif ( t1(i) < tice-t_i0 ) then + qs(i) = qd(i) + else + qs(i) = qd(i)*(tice-t1(i))/t_i0 + endif + ql(i) = qd(i) - qs(i) + qv(i) = max(0.,q(i,j,k,sphum)) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + else + do i=is,ie + qv(i) = max(0.,q(i,j,k,sphum)) + qs(i) = max(0.,q(i,j,k,liq_wat)) + qd(i) = qs(i) + cpm(i) = (1.-qv(i))*cp_air + qv(i)*cp_vapor + enddo + endif + + case(3) + do i=is,ie + qv(i) = q(i,j,k,sphum) + ql(i) = q(i,j,k,liq_wat) + qs(i) = q(i,j,k,ice_wat) + qd(i) = ql(i) + qs(i) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + case(4) ! K_warm_rain scheme with fake ice + do i=is,ie + qv(i) = q(i,j,k,sphum) + qd(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + qd(i)*c_liq + enddo + case(5) + do i=is,ie + qv(i) = q(i,j,k,sphum) + ql(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) + qs(i) = q(i,j,k,ice_wat) + q(i,j,k,snowwat) + qd(i) = ql(i) + qs(i) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + case(6:7) + do i=is,ie + qv(i) = q(i,j,k,sphum) + ql(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) + qs(i) = q(i,j,k,ice_wat) + q(i,j,k,snowwat) + q(i,j,k,graupel) + qd(i) = ql(i) + qs(i) + cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice + enddo + case default + ! call mpp_error (NOTE, 'fv_mapz::moist_cp - using default cp_air') + do i=is,ie + qd(i) = 0. + cpm(i) = cp_air + enddo + end select + + end subroutine moist_cp + +!----------------------------------------------------------------------- +!BOP +! !ROUTINE: map1_gmao --- GMAO Interpolation for vertical re-mapping +! +! !INTERFACE: + subroutine map1_gmao( km, pe1, q1, & + kn, pe2, q2, i1, i2, & + j, ibeg, iend, jbeg, jend, akap, gmao_remap, P_MAP, conserv) + implicit none + +! !INPUT PARAMETERS: + integer, intent(in) :: i1 ! Starting longitude + integer, intent(in) :: i2 ! Finishing longitude + real, intent(in) :: akap + integer, intent(in) :: P_MAP ! Thermodynamic variable to remap + ! 1:TE-logP 2:PT-PK 3:T-logP + integer, intent(in) :: gmao_remap ! 3:cubic 2:quadratic 1:linear + logical, intent(in) :: conserv + integer, intent(in) :: j ! Current latitude + integer, intent(in) :: ibeg, iend, jbeg, jend + integer, intent(in) :: km ! Original vertical dimension + integer, intent(in) :: kn ! Target vertical dimension + + real, intent(in) :: pe1(i1:i2,km+1) ! pressure at layer edges + ! (from model top to bottom surface) + ! in the original vertical coordinate + real, intent(in) :: pe2(i1:i2,kn+1) ! pressure at layer edges + ! (from model top to bottom surface) + ! in the new vertical coordinate + + real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) ! Field input +! !INPUT/OUTPUT PARAMETERS: + real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) ! Field output + +! !DESCRIPTION: +! +! Perform Cubic Interpolation a given latitude +! pe1: pressure at layer edges (from model top to bottom surface) +! in the original vertical coordinate +! pe2: pressure at layer edges (from model top to bottom surface) +! in the new vertical coordinate +! +! !REVISION HISTORY: +! 2005.11.14 Takacs Initial Code +! 2016.07.20 Putman Modified to make genaric for any thermodynamic variable +! +!EOP +!----------------------------------------------------------------------- +!BOC +! +! !LOCAL VARIABLES: + real qx(i1:i2,km) + real logpl1(i1:i2,km) + real logpl2(i1:i2,kn) + real dlogp1(i1:i2,km) + real vsum1(i1:i2) + real vsum2(i1:i2) + real am2,am1,ap0,ap1,P,PLP1,PLP0,PLM1,PLM2,DLP0,DLM1,DLM2 + + integer i, k, LM2,LM1,LP0,LP1 + +! Initialization +! -------------- + + select case (P_MAP) + case(1) + ! Remapping in Log(P) + do k=1,km + qx(:,k) = q1(i1:i2,j,k) + logpl1(:,k) = log( r2*(pe1(:,k)+pe1(:,k+1)) ) + enddo + do k=1,kn + logpl2(:,k) = log( r2*(pe2(:,k)+pe2(:,k+1)) ) + enddo + do k=1,km-1 + dlogp1(:,k) = logpl1(:,k+1)-logpl1(:,k) + enddo + case(2) + ! Remapping in P^KAPPA + do k=1,km + qx(:,k) = q1(i1:i2,j,k) + logpl1(:,k) = exp( akap*log( r2*(pe1(:,k)+pe1(:,k+1))) ) + enddo + do k=1,kn + logpl2(:,k) = exp( akap*log( r2*(pe2(:,k)+pe2(:,k+1))) ) + enddo + do k=1,km-1 + dlogp1(:,k) = logpl1(:,k+1)-logpl1(:,k) + enddo + end select + + if (conserv) then +! Compute vertical integral of Input field +! ---------------------------------------- + vsum1(:) = r0 + do i=i1,i2 + do k=1,km + vsum1(i) = vsum1(i) + qx(i,k)*( pe1(i,k+1)-pe1(i,k) ) + enddo + vsum1(i) = vsum1(i) / ( pe1(i,km+1)-pe1(i,1) ) + enddo + endif + +! Interpolate field onto target Pressures +! --------------------------------------- + do i=i1,i2 + do k=1,kn + LM1 = 1 + LP0 = 1 + do while( LP0.le.km ) + if (logpl1(i,LP0).lt.logpl2(i,k)) then + LP0 = LP0+1 + else + exit + endif + enddo + LM1 = max(LP0-1,1) + LP0 = min(LP0, km) + +! Extrapolate Linearly above first model level +! ---------------------------------------------------- + if( LM1.eq.1 .and. LP0.eq.1 ) then + q2(i,j,k) = qx(i,1) + ( qx(i,2)-qx(i,1) )*( logpl2(i,k)-logpl1(i,1) ) & + /( logpl1(i,2)-logpl1(i,1) ) + +! Extrapolate Linearly below last model level +! --------------------------------------------------- + else if( LM1.eq.km .and. LP0.eq.km ) then + q2(i,j,k) = qx(i,km) + ( qx(i,km)-qx(i,km-1) )*( logpl2(i,k )-logpl1(i,km ) ) & + /( logpl1(i,km)-logpl1(i,km-1) ) + +! Interpolate Linearly between levels 1 => 2 and km-1 => km +! ----------------------------------------------------------------- + else if( LM1.eq.1 .or. LP0.eq.km ) then + q2(i,j,k) = qx(i,LP0) + ( qx(i,LM1)-qx(i,LP0) )*( logpl2(i,k )-logpl1(i,LP0) ) & + /( logpl1(i,LM1)-logpl1(i,LP0) ) +! Interpolate Cubicly between other model levels +! ------------------------------------------------------ + else + LP1 = LP0+1 + LM2 = LM1-1 + P = logpl2(i,k) + PLP1 = logpl1(i,LP1) + PLP0 = logpl1(i,LP0) + PLM1 = logpl1(i,LM1) + PLM2 = logpl1(i,LM2) + DLP0 = dlogp1(i,LP0) + DLM1 = dlogp1(i,LM1) + DLM2 = dlogp1(i,LM2) + + ! Cubic Coefficients + ! ------------------ + if( gmao_remap .eq. 3 ) then + ap1 = (P-PLP0)*(P-PLM1)*(P-PLM2)/( DLP0*(DLP0+DLM1)*(DLP0+DLM1+DLM2) ) + ap0 = (PLP1-P)*(P-PLM1)*(P-PLM2)/( DLP0* DLM1 *( DLM1+DLM2) ) + am1 = (PLP1-P)*(PLP0-P)*(P-PLM2)/( DLM1* DLM2 *(DLP0+DLM1 ) ) + am2 = (PLP1-P)*(PLP0-P)*(PLM1-P)/( DLM2*(DLM1+DLM2)*(DLP0+DLM1+DLM2) ) + q2(i,j,k) = ap1*qx(i,LP1) + ap0*qx(i,LP0) + am1*qx(i,LM1) + am2*qx(i,LM2) + endif + + ! Quadratic Coefficients + ! ---------------------- + if( gmao_remap .eq. 2 ) then + ap1 = (P-PLP0)*(P-PLM1)/( (PLP1-PLP0)*(PLP1-PLM1) ) + ap0 = (PLP1-P)*(P-PLM1)/( (PLP1-PLP0)*(PLP0-PLM1) ) + am1 = (PLP1-P)*(PLP0-P)/( (PLP1-PLM1)*(PLP0-PLM1) ) + q2(i,j,k) = ap1*qx(i,LP1) + ap0*qx(i,LP0) + am1*qx(i,LM1) + endif + + ! Linear Coefficients + ! ------------------- + if( gmao_remap .eq. 1 ) then + q2(i,j,k) = qx(i,LP0) + ( qx(i,LM1)-qx(i,LP0) )*( logpl2(i,k )-logpl1(i,LP0) ) & + /( logpl1(i,LM1)-logpl1(i,LP0) ) + endif + + endif + + enddo + enddo + + if (conserv) then +! Compute vertical integral of Output field +! ----------------------------------------- + vsum2(:) = r0 + do i=i1,i2 + do k=1,kn + vsum2(i) = vsum2(i) + q2(i,j,k)*( pe2(i,k+1)-pe2(i,k) ) + enddo + vsum2(i) = vsum2(i) / ( pe2(i,kn+1)-pe2(i,1) ) + enddo +! Adjust Final field to conserve +! ------------------------------ + do i=i1,i2 + do k=1,kn + q2(i,j,k) = q2(i,j,k) + vsum1(i)-vsum2(i) + enddo + enddo + endif + + return +!EOC + end subroutine map1_gmao +!----------------------------------------------------------------------- + + +end module fv_mapz_mod diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 40b34799..2b99cefc 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -69,41 +69,6 @@ def __init__( "jend": grid.je+1, "kend": grid.npz + 1, }, - "pe1_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pe2_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pe0_": { - "istart": grid.is_, - "iend": grid.ie+1, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1 - }, - "pe3_": { - "istart": grid.is_, - "iend": grid.ie+1, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1 - }, - # "qvapor": {"serialname": "qvapor_js"}, - # "qliquid": {"serialname": "qliquid_js"}, - # "qice": {"serialname": "qice_js"}, - # "qrain": {"serialname": "qrain_js"}, - # "qsnow": {"serialname": "qsnow_js"}, - # "qgraupel": {"serialname": "qgraupel_js"}, - # "qcld": {"serialname": "qcld_js"}, "qvapor": { "kend": grid.npz-1, }, @@ -137,20 +102,6 @@ def __init__( "pt": {}, "cappa": {}, "ps": {}, - "pn1_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pn2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, "peln_3d": { "istart": grid.is_, "iend": grid.ie, @@ -160,13 +111,6 @@ def __init__( }, "ak": {}, "bk": {}, - "dp2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, "pk": { "istart": grid.is_, "iend": grid.ie, @@ -180,13 +124,6 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, - "pk2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, "w": { "kend": grid.npz-1, }, @@ -240,13 +177,6 @@ def __init__( "jend": grid.jed, "kend": grid.npz-1, }, - "phis_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz+1, - }, "cosa_s": { "istart": grid.isd, "iend": grid.ied, @@ -283,6 +213,7 @@ def __init__( "jstart": grid.jsd, "jend": grid.jed, }, + "te": {} } # self.write_vars = ["gz", "cvm"] self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] @@ -308,180 +239,181 @@ def __init__( # "nq", ] self.out_vars = { - "pe1_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pe2_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pe0_": { - "istart": grid.is_, - "iend": grid.ie+1, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1 - }, - "pe3_": { - "istart": grid.is_, - "iend": grid.ie+1, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1 - }, + # "pe1_": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "pe2_": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "pe0_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz+1 + # }, + # "pe3_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz+1 + # }, "pt": {}, - "cappa": {}, - "q_con": {}, - "delp": {}, - "delz": {}, - "ps": {}, - "dp2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - "pn1_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pn2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pk2_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - - "qvapor": { - "kend": grid.npz-1, - }, - "qliquid": { - "kend": grid.npz-1, - }, - "qice": { - "kend": grid.npz-1, - }, - "qrain": { - "kend": grid.npz-1, - }, - "qsnow": { - "kend": grid.npz-1, - }, - "qgraupel": { - "kend": grid.npz-1, - }, - "qcld": { - "kend": grid.npz-1, - }, - "qo3mr": { - "kend": grid.npz-1, - }, - "qsgs_tke": { - "kend": grid.npz-1, - }, - "w": { - "kend": grid.npz-1, - }, - "u": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.jsd, - "jend": grid.jed+1, - "kend": grid.npz-1, - }, - "v": { - "istart": grid.isd, - "iend": grid.ied+1, - "jstart": grid.jsd, - "jend": grid.jed, - "kend": grid.npz-1, - }, - - "mfy": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, - }, - - "cy": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, - }, - "mfx_": { - "istart": grid.is_, - "iend": grid.ie+1, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz-1, - }, - - "cx_": { - "istart": grid.is_, - "iend": grid.ie+1, - "jstart": grid.jsd, - "jend": grid.jed, - "kend": grid.npz-1, - }, - - "peln_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - - "pe_": { - "istart": grid.is_-1, - "iend": grid.ie+1, - "jstart": grid.js-1, - "jend": grid.je+1, - "kend": grid.npz + 1, - }, - - "pk": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz+1, - }, - "pkz": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - }, - "te_2d_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - }, + # "cappa": {}, + # "q_con": {}, + # "delp": {}, + # "delz": {}, + # "ps": {}, + # "dp2_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + # "pn1_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "pn2_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + # "pk2_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + + # "qvapor": { + # "kend": grid.npz-1, + # }, + # "qliquid": { + # "kend": grid.npz-1, + # }, + # "qice": { + # "kend": grid.npz-1, + # }, + # "qrain": { + # "kend": grid.npz-1, + # }, + # "qsnow": { + # "kend": grid.npz-1, + # }, + # "qgraupel": { + # "kend": grid.npz-1, + # }, + # "qcld": { + # "kend": grid.npz-1, + # }, + # "qo3mr": { + # "kend": grid.npz-1, + # }, + # "qsgs_tke": { + # "kend": grid.npz-1, + # }, + # "w": { + # "kend": grid.npz-1, + # }, + # "u": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed+1, + # "kend": grid.npz-1, + # }, + # "v": { + # "istart": grid.isd, + # "iend": grid.ied+1, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz-1, + # }, + + # "mfy": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, + + # "cy": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.js, + # "jend": grid.je+1, + # "kend": grid.npz-1, + # }, + # "mfx_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz-1, + # }, + + # "cx_": { + # "istart": grid.is_, + # "iend": grid.ie+1, + # "jstart": grid.jsd, + # "jend": grid.jed, + # "kend": grid.npz-1, + # }, + + # "peln_3d": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz + 1, + # }, + + # "pe_": { + # "istart": grid.is_-1, + # "iend": grid.ie+1, + # "jstart": grid.js-1, + # "jend": grid.je+1, + # "kend": grid.npz + 1, + # }, + + # "pk": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # "kend": grid.npz+1, + # }, + # "pkz": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # }, + # "te_2d_": { + # "istart": grid.is_, + # "iend": grid.ie, + # "jstart": grid.js, + # "jend": grid.je, + # }, + # "te": {} } self.stencil_factory = stencil_factory @@ -544,10 +476,91 @@ def __init__( ), dtype=bool, ) + self._ps = self._pe1 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + ), dtype=Float, + ) + + self._pe0 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._pe1 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._pe2 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._pe3 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._pn1 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._pn2 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._dp2 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + + ) + + self._pk2 = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._phis = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + grid.npz+1, + ), dtype=Float, + ) + + self._init_pe = stencil_factory.from_origin_domain( init_pe, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1), + domain=(grid.nic, grid.njc, grid.npz + 1), ) self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( @@ -657,8 +670,8 @@ def compute_from_storage(self, inputs): self._init_pe( inputs["pe_"], - inputs["pe1_"], - inputs["pe2_"], + self._pe1, + self._pe2, inputs["ptop"], ) @@ -675,253 +688,262 @@ def compute_from_storage(self, inputs): inputs["delp"], inputs["delz"], inputs["pe_"], - inputs["pe2_"], + self._pe2, inputs["ak"], inputs["bk"], - inputs["dp2_3d"], - inputs["ps"], - inputs["pn1_3d"], - inputs["pn2_3d"], + self._dp2, + self._ps, + self._pn1, + self._pn2, inputs["peln_3d"], True, Float(inputs["r_vir"]), ) - self._pn2_pk_delp( - inputs["pe2_"], - inputs["pn2_3d"], - inputs["pk2_3d"], - Float(inputs["akap"]), - ) - - self._map_scalar( - inputs["pt"], - inputs["pn1_3d"], - inputs["pn2_3d"], - qmin=inputs["t_min"], - interp=True, - ) - - tracers = { "qvapor": inputs["qvapor"], - "qliquid": inputs["qliquid"], - "qice": inputs["qice"], - "qrain": inputs["qrain"], - "qsnow": inputs["qsnow"], - "qgraupel": inputs["qgraupel"], - "qcld": inputs["qcld"], - "qo3mr": inputs["qo3mr"], - "qsgs_tke": inputs["qsgs_tke"], - } - - self._mapn_tracer = MapNTracer( - self.stencil_factory, - self.quantity_factory, - abs(self.kord), - self.nq, - fill=self.fill, - tracers=tracers, - ) - - self._mapn_tracer(inputs["pe1_"], inputs["pe2_"], inputs["dp2_3d"], tracers) - - self._map_single_w = MapSingle( - self.stencil_factory, - self.quantity_factory, - inputs["kord_wz"], - -2, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - - self._map_single_delz = MapSingle( - self.stencil_factory, - self.quantity_factory, - inputs["kord_wz"], - 1, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - self._map_single_w(inputs["w"], inputs["pe1_"], inputs["pe2_"], qs=inputs["ws_"], interp=False) + # self._pn2_pk_delp( + # self._pe2, + # self._pn2, + # self._pk2, + # Float(inputs["akap"]), + # ) + + # self._map_scalar( + # inputs["pt"], + # self._pn1, + # self._pn2, + # qmin=inputs["t_min"], + # interp=True, + # ) + + # tracers = { "qvapor": inputs["qvapor"], + # "qliquid": inputs["qliquid"], + # "qice": inputs["qice"], + # "qrain": inputs["qrain"], + # "qsnow": inputs["qsnow"], + # "qgraupel": inputs["qgraupel"], + # "qcld": inputs["qcld"], + # "qo3mr": inputs["qo3mr"], + # "qsgs_tke": inputs["qsgs_tke"], + # } + + # self._mapn_tracer = MapNTracer( + # self.stencil_factory, + # self.quantity_factory, + # abs(self.kord), + # self.nq, + # fill=self.fill, + # tracers=tracers, + # ) + + # self._mapn_tracer(self._pe1, + # self._pe2, + # self._dp2, + # tracers) + + # self._map_single_w = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # inputs["kord_wz"], + # -2, + # dims=[X_DIM, Y_DIM, Z_DIM], + # ) + + # self._map_single_delz = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # inputs["kord_wz"], + # 1, + # dims=[X_DIM, Y_DIM, Z_DIM], + # ) + # self._map_single_w(inputs["w"], + # self._pe1, + # self._pe2, + # qs=inputs["ws_"], + # interp=False) - self._rescale_delz_1( - inputs["delz"], - inputs["delp"], - ) + # self._rescale_delz_1( + # inputs["delz"], + # inputs["delp"], + # ) - self._map_single_delz(inputs["delz"], inputs["pe1_"], inputs["pe2_"]) - - self._rescale_delz_2( - inputs["delz"], - inputs["dp2_3d"], - ) + # self._map_single_delz(inputs["delz"], + # self._pe1, + # self._pe2) + + # self._rescale_delz_2( + # inputs["delz"], + # self._dp2, + # ) - self._w_fix_consrv_moment( - inputs["w"], - self._w2, - inputs["dp2_3d"], - self._gz, - inputs["w_max"], - inputs["w_min"], - self._compute_performed - ) - - self._map1_ppm_u = MapSingle( - self.stencil_factory, - self.quantity_factory, - inputs["kord_mt"], - -1, - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], - ) - - self._pressures_mapu( - inputs["pe_"], - inputs["ak"], - inputs["bk"], - inputs["pe0_"], - inputs["pe3_"], - inputs["ptop"], - ) + # self._w_fix_consrv_moment( + # inputs["w"], + # self._w2, + # self._dp2, + # self._gz, + # inputs["w_max"], + # inputs["w_min"], + # self._compute_performed + # ) + + # self._map1_ppm_u = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # inputs["kord_mt"], + # -1, + # dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + # ) + + # self._pressures_mapu( + # inputs["pe_"], + # inputs["ak"], + # inputs["bk"], + # self._pe0, + # self._pe3, + # inputs["ptop"], + # ) - self._pe0_ptop_xmax( - inputs["pe0_"], - inputs["ptop"], - ) - - self._map1_ppm_u( - inputs["u"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + # self._pe0_ptop_xmax( + # self._pe0, + # inputs["ptop"], + # ) + + # self._map1_ppm_u( + # inputs["u"], + # self._pe0, + # self._pe3, + # interp=False, + # ) - self._map1_ppm_u( - inputs["mfy"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + # self._map1_ppm_u( + # inputs["mfy"], + # self._pe0, + # self._pe3, + # interp=False, + # ) - self._map1_ppm_u( - inputs["cy"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + # self._map1_ppm_u( + # inputs["cy"], + # self._pe0, + # self._pe3, + # interp=False, + # ) - self._map1_ppm_v = MapSingle( - self.stencil_factory, - self.quantity_factory, - inputs["kord_mt"], - -1, - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], - ) - - self._pressures_mapv( - inputs["pe_"], - inputs["ak"], - inputs["bk"], - inputs["pe0_"], - inputs["pe3_"], - ) - - self._map1_ppm_v( - inputs["v"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + # self._map1_ppm_v = MapSingle( + # self.stencil_factory, + # self.quantity_factory, + # inputs["kord_mt"], + # -1, + # dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + # ) + + # self._pressures_mapv( + # inputs["pe_"], + # inputs["ak"], + # inputs["bk"], + # self._pe0, + # self._pe3, + # ) + + # self._map1_ppm_v( + # inputs["v"], + # self._pe0, + # self._pe3, + # interp=False, + # ) - self._map1_ppm_v( - inputs["mfx_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + # self._map1_ppm_v( + # inputs["mfx_"], + # self._pe0, + # self._pe3, + # interp=False, + # ) - self._map1_ppm_v( - inputs["cx_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) - - self._pe_pk_delp_peln(inputs["pe_"], - inputs["pk"], - inputs["delp"], - inputs["peln_3d"], - inputs["pe2_"], - inputs["pk2_3d"], - inputs["pn2_3d"], - inputs["ak"], - inputs["bk"], - inputs["akap"], - inputs["ptop"], - ) - - self._moist_cv_pkz( - inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - inputs["pkz"], - inputs["pt"], - inputs["cappa"], - inputs["delp"], - inputs["delz"], - Float(inputs["r_vir"]), - ) - - # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then - - self._moist_cv_te(inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - inputs["u"], - inputs["v"], - inputs["w"], - inputs["te_2d_"], - inputs["pt"], - inputs["phis_"], - inputs["delp"], - inputs["rsin2"], - inputs["cosa_s"], - inputs["hs"], - inputs["delz"], - inputs["grav"], - ) - - self._te_zsum(inputs["te_2d_"], - inputs["te0_2d_"], - inputs["delp"], - inputs["pkz"], - self._zsum1, - ) + # self._map1_ppm_v( + # inputs["cx_"], + # self._pe0, + # self._pe3, + # interp=False, + # ) + + # self._pe_pk_delp_peln(inputs["pe_"], + # inputs["pk"], + # inputs["delp"], + # inputs["peln_3d"], + # inputs["pe2_"], + # inputs["pk2_3d"], + # inputs["pn2_3d"], + # inputs["ak"], + # inputs["bk"], + # inputs["akap"], + # inputs["ptop"], + # ) + + # self._moist_cv_pkz( + # inputs["qvapor"], + # inputs["qliquid"], + # inputs["qrain"], + # inputs["qsnow"], + # inputs["qice"], + # inputs["qgraupel"], + # inputs["pkz"], + # inputs["pt"], + # inputs["cappa"], + # inputs["delp"], + # inputs["delz"], + # Float(inputs["r_vir"]), + # ) + + # # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then + + # self._moist_cv_te(inputs["qvapor"], + # inputs["qliquid"], + # inputs["qrain"], + # inputs["qsnow"], + # inputs["qice"], + # inputs["qgraupel"], + # inputs["u"], + # inputs["v"], + # inputs["w"], + # inputs["te_2d_"], + # inputs["pt"], + # inputs["phis_"], + # inputs["delp"], + # inputs["rsin2"], + # inputs["cosa_s"], + # inputs["hs"], + # inputs["delz"], + # inputs["grav"], + # ) + + # self._te_zsum(inputs["te_2d_"], + # inputs["te0_2d_"], + # inputs["delp"], + # inputs["pkz"], + # self._zsum1, + # ) - # Note, since this is a serial translate test, mpp_global_sum won't work without the communicator. - # Also, mpp_global_sum is currently set up for the C24 TBC setup - - # tesum = mpp_global_sum(inputs["te_2d_"]*inputs["area_64"], communicator, self.stencil_factory) - - # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - - # zsum = mpp_global_sum(self._zsum1*inputs["area_64"], communicator, self.stencil_factory) - # dtmp = tesum / (cv_air*zsum) - - # If loop based on if ( last_step .and. (.not. adiabatic) ) then - # self._most_cv_pt_last_step(inputs["qvapor"], - # inputs["qliquid"], - # inputs["qrain"], - # inputs["qsnow"], - # inputs["qice"], - # inputs["qgraupel"], - # self._gz, - # inputs["pt"], - # inputs["pkz"], - # dtmp, - # inputs["r_vir], - # ) + # # Note, since this is a serial translate test, mpp_global_sum won't work without the communicator. + # # Also, mpp_global_sum is currently set up for the C24 TBC setup + + # # tesum = mpp_global_sum(inputs["te_2d_"]*inputs["area_64"], communicator, self.stencil_factory) + + # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed + + # # zsum = mpp_global_sum(self._zsum1*inputs["area_64"], communicator, self.stencil_factory) + # # dtmp = tesum / (cv_air*zsum) + + # # If loop based on if ( last_step .and. (.not. adiabatic) ) then + # # self._most_cv_pt_last_step(inputs["qvapor"], + # # inputs["qliquid"], + # # inputs["qrain"], + # # inputs["qsnow"], + # # inputs["qice"], + # # inputs["qgraupel"], + # # self._gz, + # # inputs["pt"], + # # inputs["pkz"], + # # dtmp, + # # inputs["r_vir], + # # ) return inputs From 3dd17d5c5965e8f899a9a5721652c6aa065a187b Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 16 Jan 2025 10:23:56 -0800 Subject: [PATCH 115/252] Modified remapping translate test to compute over entire domain. Tested over first remapping j-loop --- .../translate/translate_remapping_GEOS.py | 204 +++++++++--------- 1 file changed, 102 insertions(+), 102 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 2b99cefc..4ca92ef1 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -271,7 +271,7 @@ def __init__( # "cappa": {}, # "q_con": {}, # "delp": {}, - # "delz": {}, + "delz": {}, # "ps": {}, # "dp2_3d": { # "istart": grid.is_, @@ -302,36 +302,36 @@ def __init__( # "kend": grid.npz + 1, # }, - # "qvapor": { - # "kend": grid.npz-1, - # }, - # "qliquid": { - # "kend": grid.npz-1, - # }, - # "qice": { - # "kend": grid.npz-1, - # }, - # "qrain": { - # "kend": grid.npz-1, - # }, - # "qsnow": { - # "kend": grid.npz-1, - # }, - # "qgraupel": { - # "kend": grid.npz-1, - # }, - # "qcld": { - # "kend": grid.npz-1, - # }, - # "qo3mr": { - # "kend": grid.npz-1, - # }, - # "qsgs_tke": { - # "kend": grid.npz-1, - # }, - # "w": { - # "kend": grid.npz-1, - # }, + "qvapor": { + "kend": grid.npz-1, + }, + "qliquid": { + "kend": grid.npz-1, + }, + "qice": { + "kend": grid.npz-1, + }, + "qrain": { + "kend": grid.npz-1, + }, + "qsnow": { + "kend": grid.npz-1, + }, + "qgraupel": { + "kend": grid.npz-1, + }, + "qcld": { + "kend": grid.npz-1, + }, + "qo3mr": { + "kend": grid.npz-1, + }, + "qsgs_tke": { + "kend": grid.npz-1, + }, + "w": { + "kend": grid.npz-1, + }, # "u": { # "istart": grid.isd, # "iend": grid.ied, @@ -569,13 +569,13 @@ def __init__( externals={"hydrostatic": hydrostatic}, origin=grid_indexing.origin_compute(), # domain=grid_indexing.domain_compute(add=(0, 0, 1)), - domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2]+1), # Note : Many intervals go from (0,-1) in this stencil + domain=(grid_indexing.domain[0], grid_indexing.domain[1], grid_indexing.domain[2]+1), # Note : Many intervals go from (0,-1) in this stencil ) self._pn2_pk_delp = stencil_factory.from_origin_domain( pn2_pk_delp, origin=grid_indexing.origin_compute(add=(0,0,1)), - domain=(grid.nic, 1, grid.npz-1), + domain=(grid.nic, grid.njc, grid.npz-1), ) self._map_scalar = MapSingle( @@ -589,13 +589,13 @@ def __init__( self._rescale_delz_1 = stencil_factory.from_origin_domain( rescale_delz_1, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz), + domain=(grid.nic, grid.njc, grid.npz), ) self._rescale_delz_2 = stencil_factory.from_origin_domain( rescale_delz_2, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz), + domain=(grid.nic, grid.njc, grid.npz), ) self._w_fix_consrv_moment = stencil_factory.from_origin_domain( @@ -700,80 +700,80 @@ def compute_from_storage(self, inputs): Float(inputs["r_vir"]), ) - # self._pn2_pk_delp( - # self._pe2, - # self._pn2, - # self._pk2, - # Float(inputs["akap"]), - # ) - - # self._map_scalar( - # inputs["pt"], - # self._pn1, - # self._pn2, - # qmin=inputs["t_min"], - # interp=True, - # ) + self._pn2_pk_delp( + self._pe2, + self._pn2, + self._pk2, + Float(inputs["akap"]), + ) + + self._map_scalar( + inputs["pt"], + self._pn1, + self._pn2, + qmin=inputs["t_min"], + interp=True, + ) + + tracers = { "qvapor": inputs["qvapor"], + "qliquid": inputs["qliquid"], + "qice": inputs["qice"], + "qrain": inputs["qrain"], + "qsnow": inputs["qsnow"], + "qgraupel": inputs["qgraupel"], + "qcld": inputs["qcld"], + "qo3mr": inputs["qo3mr"], + "qsgs_tke": inputs["qsgs_tke"], + } - # tracers = { "qvapor": inputs["qvapor"], - # "qliquid": inputs["qliquid"], - # "qice": inputs["qice"], - # "qrain": inputs["qrain"], - # "qsnow": inputs["qsnow"], - # "qgraupel": inputs["qgraupel"], - # "qcld": inputs["qcld"], - # "qo3mr": inputs["qo3mr"], - # "qsgs_tke": inputs["qsgs_tke"], - # } - - # self._mapn_tracer = MapNTracer( - # self.stencil_factory, - # self.quantity_factory, - # abs(self.kord), - # self.nq, - # fill=self.fill, - # tracers=tracers, - # ) + self._mapn_tracer = MapNTracer( + self.stencil_factory, + self.quantity_factory, + abs(self.kord), + self.nq, + fill=self.fill, + tracers=tracers, + ) - # self._mapn_tracer(self._pe1, - # self._pe2, - # self._dp2, - # tracers) + self._mapn_tracer(self._pe1, + self._pe2, + self._dp2, + tracers) - # self._map_single_w = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # inputs["kord_wz"], - # -2, - # dims=[X_DIM, Y_DIM, Z_DIM], - # ) + self._map_single_w = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_wz"], + -2, + dims=[X_DIM, Y_DIM, Z_DIM], + ) - # self._map_single_delz = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # inputs["kord_wz"], - # 1, - # dims=[X_DIM, Y_DIM, Z_DIM], - # ) - # self._map_single_w(inputs["w"], - # self._pe1, - # self._pe2, - # qs=inputs["ws_"], - # interp=False) + self._map_single_delz = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_wz"], + 1, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + self._map_single_w(inputs["w"], + self._pe1, + self._pe2, + qs=inputs["ws_"], + interp=False) - # self._rescale_delz_1( - # inputs["delz"], - # inputs["delp"], - # ) + self._rescale_delz_1( + inputs["delz"], + inputs["delp"], + ) - # self._map_single_delz(inputs["delz"], - # self._pe1, - # self._pe2) + self._map_single_delz(inputs["delz"], + self._pe1, + self._pe2) - # self._rescale_delz_2( - # inputs["delz"], - # self._dp2, - # ) + self._rescale_delz_2( + inputs["delz"], + self._dp2, + ) # self._w_fix_consrv_moment( # inputs["w"], From c2d75106bcf5d7dae0f130677a2eb376684f0a46 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 16 Jan 2025 14:24:02 -0800 Subject: [PATCH 116/252] Verified additional portions of remapping using full j. Note that there may be an issue with computations using map1_ppm_u --- .../translate/translate_remapping_GEOS.py | 274 +++++++++--------- 1 file changed, 137 insertions(+), 137 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 4ca92ef1..5bd4e70b 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -332,51 +332,51 @@ def __init__( "w": { "kend": grid.npz-1, }, - # "u": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.jsd, - # "jend": grid.jed+1, - # "kend": grid.npz-1, - # }, - # "v": { - # "istart": grid.isd, - # "iend": grid.ied+1, - # "jstart": grid.jsd, - # "jend": grid.jed, - # "kend": grid.npz-1, - # }, + "u": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed+1, + "kend": grid.npz-1, + }, + "v": { + "istart": grid.isd, + "iend": grid.ied+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, - # "mfy": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, + "mfy": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, - # "cy": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz-1, - # }, - # "mfx_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, + "cy": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je+1, + "kend": grid.npz-1, + }, + "mfx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz-1, + }, - # "cx_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.jsd, - # "jend": grid.jed, - # "kend": grid.npz-1, - # }, + "cx_": { + "istart": grid.is_, + "iend": grid.ie+1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz-1, + }, # "peln_3d": { # "istart": grid.is_, @@ -386,13 +386,13 @@ def __init__( # "kend": grid.npz + 1, # }, - # "pe_": { - # "istart": grid.is_-1, - # "iend": grid.ie+1, - # "jstart": grid.js-1, - # "jend": grid.je+1, - # "kend": grid.npz + 1, - # }, + "pe_": { + "istart": grid.is_-1, + "iend": grid.ie+1, + "jstart": grid.js-1, + "jend": grid.je+1, + "kend": grid.npz + 1, + }, # "pk": { # "istart": grid.is_, @@ -601,19 +601,19 @@ def __init__( self._w_fix_consrv_moment = stencil_factory.from_origin_domain( func=W_fix_consrv_moment, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz), + domain=(grid.nic, grid.njc, grid.npz), ) self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1) + domain=(grid_indexing.domain[0],grid_indexing.domain[1],grid_indexing.domain[2] + 1) ) self._pe0_ptop_xmax = stencil_factory.from_origin_domain( pe0_ptop_xmax, - origin=(grid_indexing.domain[0]+3,3,0), - domain=(1,1,grid_indexing.domain[2] + 1) + origin=(grid_indexing.n_halo+grid_indexing.domain[0],grid_indexing.n_halo,0), + domain=(1,grid_indexing.domain[1], 1) ) self._pressures_mapv = stencil_factory.from_origin_domain( @@ -621,7 +621,7 @@ def __init__( origin=grid_indexing.origin_compute(), domain=( grid_indexing.domain[0] + 1, - 1, + grid_indexing.domain[1], grid_indexing.domain[2] + 1, ), ) @@ -775,95 +775,95 @@ def compute_from_storage(self, inputs): self._dp2, ) - # self._w_fix_consrv_moment( - # inputs["w"], - # self._w2, - # self._dp2, - # self._gz, - # inputs["w_max"], - # inputs["w_min"], - # self._compute_performed - # ) - - # self._map1_ppm_u = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # inputs["kord_mt"], - # -1, - # dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], - # ) - - # self._pressures_mapu( - # inputs["pe_"], - # inputs["ak"], - # inputs["bk"], - # self._pe0, - # self._pe3, - # inputs["ptop"], - # ) + self._w_fix_consrv_moment( + inputs["w"], + self._w2, + self._dp2, + self._gz, + inputs["w_max"], + inputs["w_min"], + self._compute_performed + ) + + self._map1_ppm_u = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_mt"], + -1, + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + ) + + self._pressures_mapu( + inputs["pe_"], + inputs["ak"], + inputs["bk"], + self._pe0, + self._pe3, + inputs["ptop"], + ) - # self._pe0_ptop_xmax( - # self._pe0, - # inputs["ptop"], - # ) - - # self._map1_ppm_u( - # inputs["u"], - # self._pe0, - # self._pe3, - # interp=False, - # ) + self._pe0_ptop_xmax( + self._pe0, + inputs["ptop"], + ) + + self._map1_ppm_u( + inputs["u"], + self._pe0, + self._pe3, + interp=False, + ) - # self._map1_ppm_u( - # inputs["mfy"], - # self._pe0, - # self._pe3, - # interp=False, - # ) + self._map1_ppm_u( + inputs["mfy"], + self._pe0, + self._pe3, + interp=False, + ) - # self._map1_ppm_u( - # inputs["cy"], - # self._pe0, - # self._pe3, - # interp=False, - # ) + self._map1_ppm_u( + inputs["cy"], + self._pe0, + self._pe3, + interp=False, + ) - # self._map1_ppm_v = MapSingle( - # self.stencil_factory, - # self.quantity_factory, - # inputs["kord_mt"], - # -1, - # dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], - # ) - - # self._pressures_mapv( - # inputs["pe_"], - # inputs["ak"], - # inputs["bk"], - # self._pe0, - # self._pe3, - # ) - - # self._map1_ppm_v( - # inputs["v"], - # self._pe0, - # self._pe3, - # interp=False, - # ) + self._map1_ppm_v = MapSingle( + self.stencil_factory, + self.quantity_factory, + inputs["kord_mt"], + -1, + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + ) + + self._pressures_mapv( + inputs["pe_"], + inputs["ak"], + inputs["bk"], + self._pe0, + self._pe3, + ) + + self._map1_ppm_v( + inputs["v"], + self._pe0, + self._pe3, + interp=False, + ) - # self._map1_ppm_v( - # inputs["mfx_"], - # self._pe0, - # self._pe3, - # interp=False, - # ) + self._map1_ppm_v( + inputs["mfx_"], + self._pe0, + self._pe3, + interp=False, + ) - # self._map1_ppm_v( - # inputs["cx_"], - # self._pe0, - # self._pe3, - # interp=False, - # ) + self._map1_ppm_v( + inputs["cx_"], + self._pe0, + self._pe3, + interp=False, + ) # self._pe_pk_delp_peln(inputs["pe_"], # inputs["pk"], From 9f6603c19fc76f2b37ccb1113d185339ec58ce18 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 16 Jan 2025 14:53:06 -0800 Subject: [PATCH 117/252] Verified mapz stencils within second j loop when executing with entire j range --- fv_mapz.F90.SER.v2 | 134 +++++++++--------- .../translate/translate_remapping_GEOS.py | 102 ++++++------- 2 files changed, 118 insertions(+), 118 deletions(-) diff --git a/fv_mapz.F90.SER.v2 b/fv_mapz.F90.SER.v2 index 85f57ff5..338cef68 100644 --- a/fv_mapz.F90.SER.v2 +++ b/fv_mapz.F90.SER.v2 @@ -1076,73 +1076,6 @@ contains 1000 continue -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js, je -!$ser verbatim do i = is, ie -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js-1, je+1 -!$ser verbatim do i = is-1, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe_=pe_3d -!$ser data qvapor=q(:,:,:,sphum) -!$ser data qliquid=q(:,:,:,liq_wat) -!$ser data qice=q(:,:,:,ice_wat) -!$ser data qrain=q(:,:,:,rainwat) -!$ser data qsnow=q(:,:,:,snowwat) -!$ser data qgraupel=q(:,:,:,graupel) -!$ser data qcld=q(:,:,:,cld_amt) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data peln_3d=peln_3d -!$ser data r_vir=r_vir -!$ser data cvm=cvm -!$ser data pk=pk -!$ser data pkz=pkz -!$ser data akap=akap -!$ser data t_min=qmin -!$ser data kord_wz=kord_wz -!$ser data ws_=ws -!$ser data w=w -!$ser data u=u -!$ser data v=v -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser data w_max=w_max_ -!$ser data w_min=w_min_ -!$ser data kord_mt=kord_mt -!$ser data cosa_s=cosa_s -!$ser data rsin2=rsin2 -!$ser data hs=hs -!$ser data grav=grav_ -!$ser data te_2d_=te_2d_f32 -!$ser data te0_2d_=te0_2d_f32 -!$ser data area_64=area_64 -!$ser data te=te - ! Update pressure variables and get new pkz, T_v, and omega !$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & @@ -1430,6 +1363,73 @@ contains 2000 continue +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js, je +!$ser verbatim do i = is, ie +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js-1, je+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe_=pe_3d +!$ser data qvapor=q(:,:,:,sphum) +!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qice=q(:,:,:,ice_wat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qgraupel=q(:,:,:,graupel) +!$ser data qcld=q(:,:,:,cld_amt) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data peln_3d=peln_3d +!$ser data r_vir=r_vir +!$ser data cvm=cvm +!$ser data pk=pk +!$ser data pkz=pkz +!$ser data akap=akap +!$ser data t_min=qmin +!$ser data kord_wz=kord_wz +!$ser data ws_=ws +!$ser data w=w +!$ser data u=u +!$ser data v=v +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser data w_max=w_max_ +!$ser data w_min=w_min_ +!$ser data kord_mt=kord_mt +!$ser data cosa_s=cosa_s +!$ser data rsin2=rsin2 +!$ser data hs=hs +!$ser data grav=grav_ +!$ser data te_2d_=te_2d_f32 +!$ser data te0_2d_=te0_2d_f32 +!$ser data area_64=area_64 +!$ser data te=te + ! Do total energy conservation and fast saturation adjustment as requested ! and fill new PT (Theta_V) for next k_split step or export dry T diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 5bd4e70b..b8bdc0f4 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -268,9 +268,9 @@ def __init__( # "kend": grid.npz+1 # }, "pt": {}, - # "cappa": {}, + "cappa": {}, # "q_con": {}, - # "delp": {}, + "delp": {}, "delz": {}, # "ps": {}, # "dp2_3d": { @@ -378,13 +378,13 @@ def __init__( "kend": grid.npz-1, }, - # "peln_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, + "peln_3d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz + 1, + }, "pe_": { "istart": grid.is_-1, @@ -394,19 +394,19 @@ def __init__( "kend": grid.npz + 1, }, - # "pk": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz+1, - # }, - # "pkz": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # }, + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz+1, + }, + "pkz": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, # "te_2d_": { # "istart": grid.is_, # "iend": grid.ie, @@ -629,14 +629,14 @@ def __init__( self._pe_pk_delp_peln = stencil_factory.from_origin_domain( pe_pk_delp_peln, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2] + 1, + domain=(grid_indexing.domain[0], grid_indexing.domain[1], grid_indexing.domain[2] + 1, ), ) self._moist_cv_pkz = stencil_factory.from_origin_domain( moist_cv.moist_pkz, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz), + domain=(grid.nic, grid.njc, grid.npz), ) self._moist_cv_te = stencil_factory.from_origin_domain( @@ -865,33 +865,33 @@ def compute_from_storage(self, inputs): interp=False, ) - # self._pe_pk_delp_peln(inputs["pe_"], - # inputs["pk"], - # inputs["delp"], - # inputs["peln_3d"], - # inputs["pe2_"], - # inputs["pk2_3d"], - # inputs["pn2_3d"], - # inputs["ak"], - # inputs["bk"], - # inputs["akap"], - # inputs["ptop"], - # ) - - # self._moist_cv_pkz( - # inputs["qvapor"], - # inputs["qliquid"], - # inputs["qrain"], - # inputs["qsnow"], - # inputs["qice"], - # inputs["qgraupel"], - # inputs["pkz"], - # inputs["pt"], - # inputs["cappa"], - # inputs["delp"], - # inputs["delz"], - # Float(inputs["r_vir"]), - # ) + self._pe_pk_delp_peln(inputs["pe_"], + inputs["pk"], + inputs["delp"], + inputs["peln_3d"], + self._pe2, + self._pk2, + self._pn2, + inputs["ak"], + inputs["bk"], + inputs["akap"], + inputs["ptop"], + ) + + self._moist_cv_pkz( + inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + inputs["pkz"], + inputs["pt"], + inputs["cappa"], + inputs["delp"], + inputs["delz"], + Float(inputs["r_vir"]), + ) # # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then From b7be2055c4d12244f9a16f9550ca506ce72a9f6a Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 17 Jan 2025 08:23:52 -0800 Subject: [PATCH 118/252] Removed extraneous code from remapping_GEOS translate test --- .../translate/translate_remapping_GEOS.py | 45 ------------------- 1 file changed, 45 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index b8bdc0f4..811bcaf4 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -21,36 +21,6 @@ from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.mapn_tracer import MapNTracer -def moist_pt( - qvapor: FloatField, - qliquid: FloatField, - qrain: FloatField, - qsnow: FloatField, - qice: FloatField, - qgraupel: FloatField, - q_con: FloatField, - pt: FloatField, - cappa: FloatField, - delp: FloatField, - delz: FloatField, - r_vir: Float, -): - with computation(PARALLEL), interval(...): - cvm, gz, q_con, cappa, pt = moist_cv.moist_pt_func( - qvapor, - qliquid, - qrain, - qsnow, - qice, - qgraupel, - q_con, - pt, - cappa, - delp, - delz, - r_vir, - ) - class TranslateRemapping_GEOS(TranslateDycoreFortranData2Py): def __init__( self, @@ -215,12 +185,6 @@ def __init__( }, "te": {} } - # self.write_vars = ["gz", "cvm"] - self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel","qcld"] - for k, v in self.in_vars["data_vars"].items(): - # if k not in self.write_vars: - if k in self.write_vars: - v["axis"] = 1 self.in_vars["parameters"] = [ "ptop", "r_vir", @@ -659,15 +623,6 @@ def __init__( def compute_from_storage(self, inputs): - # Replicates tracer values in I along the J direction - for name, value in inputs.items(): - if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: - inputs[name] = self.make_storage_data( - pad_field_in_j( - value, self.grid.njd, backend=self.stencil_factory.backend - ) - ) - self._init_pe( inputs["pe_"], self._pe1, From 6217b10b9728f3674cf55f7db81625c9090bc32b Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 17 Jan 2025 12:52:55 -0800 Subject: [PATCH 119/252] Converted GEOS remapping translate test into a parallel translate test. --- .../translate/translate_remapping_GEOS.py | 865 +++++++++++++++--- 1 file changed, 728 insertions(+), 137 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 811bcaf4..acd1d364 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,4 +1,3 @@ -import ndsl.dsl.gt4py_utils as utils from ndsl import Namelist, StencilFactory from pyFV3 import DynamicalCoreConfig from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure, pn2_pk_delp @@ -8,8 +7,8 @@ from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax, pressures_mapv, pe_pk_delp_peln from pyFV3.stencils.mpp_global_sum import mpp_global_sum -from ndsl.stencils.testing import pad_field_in_j, Grid -from pyFV3.testing import TranslateDycoreFortranData2Py +from ndsl.stencils.testing import Grid +from ndsl.stencils.testing import ParallelTranslateBaseSlicing from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -18,10 +17,315 @@ Z_DIM, Z_INTERFACE_DIM, ) -from ndsl.dsl.typing import Float, FloatField +from ndsl.dsl.typing import Float from pyFV3.stencils.mapn_tracer import MapNTracer - -class TranslateRemapping_GEOS(TranslateDycoreFortranData2Py): +from types import SimpleNamespace + +class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): + inputs = { + "pe_": { + "name": "pe_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qvapor": { + "name": "qvapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qliquid": { + "name": "qliquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qice": { + "name": "qice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qrain": { + "name": "qrain", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsnow": { + "name": "qsnow", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qgraupel": { + "name": "qgraupel", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qcld": { + "name": "qcld", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qo3mr": { + "name": "qo3mr", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsgs_tke": { + "name": "qsgs_tke", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delp": { + "name": "delp", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delz": { + "name": "delz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "q_con": { + "name": "q_con", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "pt": { + "name": "pt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cappa": { + "name": "cappa", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "ps": { + "name": "ps", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "peln_3d": { + "name": "peln_3d", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "ak": { + "name": "pt", + "dims": [X_DIM], + "units": "No Units", + }, + "bk": { + "name": "bk", + "dims": [X_DIM], + "units": "No Units", + }, + "pk": { + "name": "pk", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "pkz": { + "name": "pkz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "w": { + "name": "w", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "ws_": { + "name": "ws_", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "u": { + "name": "u", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "v": { + "name": "v", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfy": { + "name": "mfy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cy": { + "name": "cy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfx_": { + "name": "pt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cx_": { + "name": "cx_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cosa_s": { + "name": "cosa_s", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "rsin2": { + "name": "rsin2", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "te_2d_": { + "name": "te_2d_", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "hs": { + "name": "hs", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "te0_2d_": { + "name": "te0_2d", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "area_64": { + "name": "area_64", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "te": { + "name": "te", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + } + outputs = { + "pt": { + "name": "pt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cappa": { + "name": "cappa", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delp": { + "name": "delp", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delz": { + "name": "delz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qvapor": { + "name": "qvapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qliquid": { + "name": "qliquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qice": { + "name": "qice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qrain": { + "name": "qrain", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsnow": { + "name": "qsnow", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qgraupel": { + "name": "qgraupel", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qcld": { + "name": "qcld", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qo3mr": { + "name": "qo3mr", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsgs_tke": { + "name": "qsgs_tke", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "w": { + "name": "w", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "u": { + "name": "u", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "v": { + "name": "v", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfy": { + "name": "mfy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cy": { + "name": "cy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfx_": { + "name": "pt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cx_": { + "name": "cx_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "peln_3d": { + "name": "peln_3d", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "pe_": { + "name": "pe_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "pk": { + "name": "pk", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "pkz": { + "name": "pkz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + } def __init__( self, grid: Grid, @@ -30,7 +334,7 @@ def __init__( ): super().__init__(grid, namelist, stencil_factory) - self.in_vars["data_vars"] = { + self._base.in_vars["data_vars"] = { "pe_": { "istart": grid.is_-1, @@ -185,7 +489,7 @@ def __init__( }, "te": {} } - self.in_vars["parameters"] = [ + self._base.in_vars["parameters"] = [ "ptop", "r_vir", # "remap_t", # For some reason, translate test can't accept a logical variable @@ -202,7 +506,7 @@ def __init__( # "mdt", # "nq", ] - self.out_vars = { + self._base.out_vars = { # "pe1_": { # "istart": grid.is_, # "iend": grid.ie, @@ -621,64 +925,351 @@ def __init__( domain=(grid.nic, grid.njc, grid.npz), ) - def compute_from_storage(self, inputs): + def compute_sequential(self, inputs_list, communicator_list): + state_list = self.state_list_from_inputs_list(inputs_list) + for state in state_list: + state_namespace = SimpleNamespace(**state) + + self._init_pe( + state_namespace.pe_, + self._pe1, + self._pe2, + state_namespace.ptop, + ) + + self._moist_cv_pt_pressure( + state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.q_con, + state_namespace.pt, + state_namespace.cappa, + state_namespace.delp, + state_namespace.delz, + state_namespace.pe_, + self._pe2, + state_namespace.ak, + state_namespace.bk, + self._dp2, + self._ps, + self._pn1, + self._pn2, + state_namespace.peln_3d, + True, + Float(state_namespace.r_vir), + ) + + self._pn2_pk_delp( + self._pe2, + self._pn2, + self._pk2, + Float(state_namespace.akap), + ) + + self._map_scalar( + state_namespace.pt, + self._pn1, + self._pn2, + qmin=state_namespace.t_min, + interp=True, + ) + + tracers = { "qvapor": state_namespace.qvapor, + "qliquid": state_namespace.qliquid, + "qice": state_namespace.qice, + "qrain": state_namespace.qrain, + "qsnow": state_namespace.qsnow, + "qgraupel": state_namespace.qgraupel, + "qcld": state_namespace.qcld, + "qo3mr": state_namespace.qo3mr, + "qsgs_tke": state_namespace.qsgs_tke, + } + + self._mapn_tracer = MapNTracer( + self.stencil_factory, + self.quantity_factory, + abs(self.kord), + self.nq, + fill=self.fill, + tracers=tracers, + ) + + self._mapn_tracer(self._pe1, + self._pe2, + self._dp2, + tracers) + + self._map_single_w = MapSingle( + self.stencil_factory, + self.quantity_factory, + state_namespace.kord_wz, + -2, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._map_single_delz = MapSingle( + self.stencil_factory, + self.quantity_factory, + state_namespace.kord_wz, + 1, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + self._map_single_w(state_namespace.w, + self._pe1, + self._pe2, + qs=state_namespace.ws_, + interp=False) + + self._rescale_delz_1( + state_namespace.delz, + state_namespace.delp, + ) + + self._map_single_delz(state_namespace.delz, + self._pe1, + self._pe2) + + self._rescale_delz_2( + state_namespace.delz, + self._dp2, + ) + + self._w_fix_consrv_moment( + state_namespace.w, + self._w2, + self._dp2, + self._gz, + state_namespace.w_max, + state_namespace.w_min, + self._compute_performed + ) + + self._map1_ppm_u = MapSingle( + self.stencil_factory, + self.quantity_factory, + state_namespace.kord_mt, + -1, + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + ) + + self._pressures_mapu( + state_namespace.pe_, + state_namespace.ak, + state_namespace.bk, + self._pe0, + self._pe3, + state_namespace.ptop, + ) + + self._pe0_ptop_xmax( + self._pe0, + state_namespace.ptop, + ) + + self._map1_ppm_u( + state_namespace.u, + self._pe0, + self._pe3, + interp=False, + ) + + self._map1_ppm_u( + state_namespace.mfy, + self._pe0, + self._pe3, + interp=False, + ) + + self._map1_ppm_u( + state_namespace.cy, + self._pe0, + self._pe3, + interp=False, + ) + + self._map1_ppm_v = MapSingle( + self.stencil_factory, + self.quantity_factory, + state_namespace.kord_mt, + -1, + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + ) + + self._pressures_mapv( + state_namespace.pe_, + state_namespace.ak, + state_namespace.bk, + self._pe0, + self._pe3, + ) + + self._map1_ppm_v( + state_namespace.v, + self._pe0, + self._pe3, + interp=False, + ) + + self._map1_ppm_v( + state_namespace.mfx_, + self._pe0, + self._pe3, + interp=False, + ) + + self._map1_ppm_v( + state_namespace.cx_, + self._pe0, + self._pe3, + interp=False, + ) + + self._pe_pk_delp_peln(state_namespace.pe_, + state_namespace.pk, + state_namespace.delp, + state_namespace.peln_3d, + self._pe2, + self._pk2, + self._pn2, + state_namespace.ak, + state_namespace.bk, + state_namespace.akap, + state_namespace.ptop, + ) + + self._moist_cv_pkz( + state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.pkz, + state_namespace.pt, + state_namespace.cappa, + state_namespace.delp, + state_namespace.delz, + Float(state_namespace.r_vir), + ) + + # # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then + + # self._moist_cv_te(state_namespace.qvapor, + # state_namespace.qliquid, + # state_namespace.qrain, + # state_namespace.qsnow, + # state_namespace.qice, + # state_namespace.qgraupel, + # state_namespace.u, + # state_namespace.v, + # state_namespace.w, + # state_namespace.te_2d_, + # state_namespace.pt, + # state_namespace.phis_, + # state_namespace.delp, + # state_namespace.rsin2, + # state_namespace.cosa_s, + # state_namespace.hs, + # state_namespace.delz, + # state_namespace.grav, + # ) + + # self._te_zsum(state_namespace.te_2d_, + # state_namespace.te0_2d_, + # state_namespace.delp, + # state_namespace.pkz, + # self._zsum1, + # ) + + # # Note, since this is a serial translate test, mpp_global_sum won't work without the communicator. + # # Also, mpp_global_sum is currently set up for the C24 TBC setup + + # # tesum = mpp_global_sum(state_namespace.te_2d_*state_namespace.area_64, communicator, self.stencil_factory) + + # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed + + # # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64, communicator, self.stencil_factory) + # # dtmp = tesum / (cv_air*zsum) + + # # If loop based on if ( last_step .and. (.not. adiabatic) ) then + # # self._most_cv_pt_last_step(state_namespace.qvapor, + # # state_namespace.qliquid, + # # state_namespace.qrain, + # # state_namespace.qsnow, + # # state_namespace.qice, + # # state_namespace.qgraupel, + # # self._gz, + # # state_namespace.pt, + # # state_namespace.pkz, + # # dtmp, + # # state_namespace.r_vir], + # # ) + + return self.outputs_list_from_state_list(state_list) + + def compute_parallel(self, inputs, communicator): + state = self.state_from_inputs(inputs) + state_namespace = SimpleNamespace(**state) self._init_pe( - inputs["pe_"], + state_namespace.pe_, self._pe1, self._pe2, - inputs["ptop"], + state_namespace.ptop, ) self._moist_cv_pt_pressure( - inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - inputs["q_con"], - inputs["pt"], - inputs["cappa"], - inputs["delp"], - inputs["delz"], - inputs["pe_"], + state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.q_con, + state_namespace.pt, + state_namespace.cappa, + state_namespace.delp, + state_namespace.delz, + state_namespace.pe_, self._pe2, - inputs["ak"], - inputs["bk"], + state_namespace.ak, + state_namespace.bk, self._dp2, self._ps, self._pn1, self._pn2, - inputs["peln_3d"], + state_namespace.peln_3d, True, - Float(inputs["r_vir"]), + Float(state_namespace.r_vir), ) self._pn2_pk_delp( self._pe2, self._pn2, self._pk2, - Float(inputs["akap"]), + Float(state_namespace.akap), ) self._map_scalar( - inputs["pt"], + state_namespace.pt, self._pn1, self._pn2, - qmin=inputs["t_min"], + qmin=state_namespace.t_min, interp=True, ) - tracers = { "qvapor": inputs["qvapor"], - "qliquid": inputs["qliquid"], - "qice": inputs["qice"], - "qrain": inputs["qrain"], - "qsnow": inputs["qsnow"], - "qgraupel": inputs["qgraupel"], - "qcld": inputs["qcld"], - "qo3mr": inputs["qo3mr"], - "qsgs_tke": inputs["qsgs_tke"], + tracers = { "qvapor": state_namespace.qvapor, + "qliquid": state_namespace.qliquid, + "qice": state_namespace.qice, + "qrain": state_namespace.qrain, + "qsnow": state_namespace.qsnow, + "qgraupel": state_namespace.qgraupel, + "qcld": state_namespace.qcld, + "qo3mr": state_namespace.qo3mr, + "qsgs_tke": state_namespace.qsgs_tke, } self._mapn_tracer = MapNTracer( @@ -691,14 +1282,14 @@ def compute_from_storage(self, inputs): ) self._mapn_tracer(self._pe1, - self._pe2, - self._dp2, - tracers) + self._pe2, + self._dp2, + tracers) self._map_single_w = MapSingle( self.stencil_factory, self.quantity_factory, - inputs["kord_wz"], + state_namespace.kord_wz, -2, dims=[X_DIM, Y_DIM, Z_DIM], ) @@ -706,78 +1297,78 @@ def compute_from_storage(self, inputs): self._map_single_delz = MapSingle( self.stencil_factory, self.quantity_factory, - inputs["kord_wz"], + state_namespace.kord_wz, 1, dims=[X_DIM, Y_DIM, Z_DIM], ) - self._map_single_w(inputs["w"], - self._pe1, - self._pe2, - qs=inputs["ws_"], - interp=False) + self._map_single_w(state_namespace.w, + self._pe1, + self._pe2, + qs=state_namespace.ws_, + interp=False) self._rescale_delz_1( - inputs["delz"], - inputs["delp"], + state_namespace.delz, + state_namespace.delp, ) - self._map_single_delz(inputs["delz"], - self._pe1, - self._pe2) + self._map_single_delz(state_namespace.delz, + self._pe1, + self._pe2) self._rescale_delz_2( - inputs["delz"], + state_namespace.delz, self._dp2, ) self._w_fix_consrv_moment( - inputs["w"], - self._w2, - self._dp2, - self._gz, - inputs["w_max"], - inputs["w_min"], - self._compute_performed - ) + state_namespace.w, + self._w2, + self._dp2, + self._gz, + state_namespace.w_max, + state_namespace.w_min, + self._compute_performed + ) self._map1_ppm_u = MapSingle( self.stencil_factory, self.quantity_factory, - inputs["kord_mt"], + state_namespace.kord_mt, -1, dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], ) self._pressures_mapu( - inputs["pe_"], - inputs["ak"], - inputs["bk"], + state_namespace.pe_, + state_namespace.ak, + state_namespace.bk, self._pe0, self._pe3, - inputs["ptop"], + state_namespace.ptop, ) self._pe0_ptop_xmax( self._pe0, - inputs["ptop"], + state_namespace.ptop, ) self._map1_ppm_u( - inputs["u"], + state_namespace.u, self._pe0, self._pe3, interp=False, ) self._map1_ppm_u( - inputs["mfy"], + state_namespace.mfy, self._pe0, self._pe3, interp=False, ) self._map1_ppm_u( - inputs["cy"], + state_namespace.cy, self._pe0, self._pe3, interp=False, @@ -786,119 +1377,119 @@ def compute_from_storage(self, inputs): self._map1_ppm_v = MapSingle( self.stencil_factory, self.quantity_factory, - inputs["kord_mt"], + state_namespace.kord_mt, -1, dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], ) self._pressures_mapv( - inputs["pe_"], - inputs["ak"], - inputs["bk"], + state_namespace.pe_, + state_namespace.ak, + state_namespace.bk, self._pe0, self._pe3, ) self._map1_ppm_v( - inputs["v"], + state_namespace.v, self._pe0, self._pe3, interp=False, ) self._map1_ppm_v( - inputs["mfx_"], + state_namespace.mfx_, self._pe0, self._pe3, interp=False, ) self._map1_ppm_v( - inputs["cx_"], + state_namespace.cx_, self._pe0, self._pe3, interp=False, ) - self._pe_pk_delp_peln(inputs["pe_"], - inputs["pk"], - inputs["delp"], - inputs["peln_3d"], - self._pe2, - self._pk2, - self._pn2, - inputs["ak"], - inputs["bk"], - inputs["akap"], - inputs["ptop"], + self._pe_pk_delp_peln(state_namespace.pe_, + state_namespace.pk, + state_namespace.delp, + state_namespace.peln_3d, + self._pe2, + self._pk2, + self._pn2, + state_namespace.ak, + state_namespace.bk, + state_namespace.akap, + state_namespace.ptop, ) self._moist_cv_pkz( - inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - inputs["pkz"], - inputs["pt"], - inputs["cappa"], - inputs["delp"], - inputs["delz"], - Float(inputs["r_vir"]), + state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.pkz, + state_namespace.pt, + state_namespace.cappa, + state_namespace.delp, + state_namespace.delz, + Float(state_namespace.r_vir), ) # # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then - # self._moist_cv_te(inputs["qvapor"], - # inputs["qliquid"], - # inputs["qrain"], - # inputs["qsnow"], - # inputs["qice"], - # inputs["qgraupel"], - # inputs["u"], - # inputs["v"], - # inputs["w"], - # inputs["te_2d_"], - # inputs["pt"], - # inputs["phis_"], - # inputs["delp"], - # inputs["rsin2"], - # inputs["cosa_s"], - # inputs["hs"], - # inputs["delz"], - # inputs["grav"], + # self._moist_cv_te(state_namespace.qvapor, + # state_namespace.qliquid, + # state_namespace.qrain, + # state_namespace.qsnow, + # state_namespace.qice, + # state_namespace.qgraupel, + # state_namespace.u, + # state_namespace.v, + # state_namespace.w, + # state_namespace.te_2d_, + # state_namespace.pt, + # state_namespace.phis_, + # state_namespace.delp, + # state_namespace.rsin2, + # state_namespace.cosa_s, + # state_namespace.hs, + # state_namespace.delz, + # state_namespace.grav, # ) - # self._te_zsum(inputs["te_2d_"], - # inputs["te0_2d_"], - # inputs["delp"], - # inputs["pkz"], + # self._te_zsum(state_namespace.te_2d_, + # state_namespace.te0_2d_, + # state_namespace.delp, + # state_namespace.pkz, # self._zsum1, # ) # # Note, since this is a serial translate test, mpp_global_sum won't work without the communicator. # # Also, mpp_global_sum is currently set up for the C24 TBC setup - # # tesum = mpp_global_sum(inputs["te_2d_"]*inputs["area_64"], communicator, self.stencil_factory) + # # tesum = mpp_global_sum(state_namespace.te_2d_*state_namespace.area_64, communicator, self.stencil_factory) # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - # # zsum = mpp_global_sum(self._zsum1*inputs["area_64"], communicator, self.stencil_factory) + # # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64, communicator, self.stencil_factory) # # dtmp = tesum / (cv_air*zsum) # # If loop based on if ( last_step .and. (.not. adiabatic) ) then - # # self._most_cv_pt_last_step(inputs["qvapor"], - # # inputs["qliquid"], - # # inputs["qrain"], - # # inputs["qsnow"], - # # inputs["qice"], - # # inputs["qgraupel"], + # # self._most_cv_pt_last_step(state_namespace.qvapor, + # # state_namespace.qliquid, + # # state_namespace.qrain, + # # state_namespace.qsnow, + # # state_namespace.qice, + # # state_namespace.qgraupel, # # self._gz, - # # inputs["pt"], - # # inputs["pkz"], + # # state_namespace.pt, + # # state_namespace.pkz, # # dtmp, - # # inputs["r_vir], + # # state_namespace.r_vir], # # ) - return inputs + return self.outputs_from_state(state) \ No newline at end of file From 240cf340b2859345f8b7ac90a075aa48b3cdfb0c Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 17 Jan 2025 15:20:05 -0800 Subject: [PATCH 120/252] Debugging parallel remapping translate test --- pyFV3/stencils/mpp_global_sum.py | 4 +- .../translate/translate_remapping_GEOS.py | 422 ++++-------------- 2 files changed, 82 insertions(+), 344 deletions(-) diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py index 01b8d3d9..6398c6db 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/stencils/mpp_global_sum.py @@ -15,14 +15,14 @@ def mpp_global_sum(inputArray, communicator, stencil_factory=None): mag_max_term = 0.0 ints_sum = Quantity( - data=np.zeros((NUMINT),dtype=np.float64), + data=np.zeros((NUMINT),dtype=np.float32), dims=["K"], units="dunno", gt4py_backend=stencil_factory.backend, ) ints_sum_reduce = Quantity( - data=np.zeros((NUMINT),dtype=np.float64), + data=np.zeros((NUMINT),dtype=np.float32), dims=["K"], units="dunno", gt4py_backend=stencil_factory.backend, diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index acd1d364..02fe9df1 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -193,8 +193,8 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "area_64": { - "name": "area_64", + "area_64_": { + "name": "area_64_", "dims": [X_DIM, Y_DIM], "units": "No Units", }, @@ -481,7 +481,7 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, - "area_64": { + "area_64_": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.jsd, @@ -500,11 +500,12 @@ def __init__( "w_min", "kord_mt", "grav", - # "zvir", - # "last_step", - # "consv_te", - # "mdt", - # "nq", + "last_step", + "do_adiabatic_init", + "consv", + "consv_min", + "cv_air", + "adiabatic", ] self._base.out_vars = { # "pe1_": { @@ -910,13 +911,13 @@ def __init__( self._moist_cv_te = stencil_factory.from_origin_domain( moist_cv.moist_te, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz+1), + domain=(grid.nic, grid.njc, grid.npz+1), ) self._te_zsum = stencil_factory.from_origin_domain( moist_cv.te_zsum, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz), + domain=(grid.nic, grid.njc, grid.npz), ) self._most_cv_pt_last_step = stencil_factory.from_origin_domain( @@ -926,289 +927,7 @@ def __init__( ) def compute_sequential(self, inputs_list, communicator_list): - state_list = self.state_list_from_inputs_list(inputs_list) - for state in state_list: - state_namespace = SimpleNamespace(**state) - - self._init_pe( - state_namespace.pe_, - self._pe1, - self._pe2, - state_namespace.ptop, - ) - - self._moist_cv_pt_pressure( - state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, - state_namespace.q_con, - state_namespace.pt, - state_namespace.cappa, - state_namespace.delp, - state_namespace.delz, - state_namespace.pe_, - self._pe2, - state_namespace.ak, - state_namespace.bk, - self._dp2, - self._ps, - self._pn1, - self._pn2, - state_namespace.peln_3d, - True, - Float(state_namespace.r_vir), - ) - - self._pn2_pk_delp( - self._pe2, - self._pn2, - self._pk2, - Float(state_namespace.akap), - ) - - self._map_scalar( - state_namespace.pt, - self._pn1, - self._pn2, - qmin=state_namespace.t_min, - interp=True, - ) - - tracers = { "qvapor": state_namespace.qvapor, - "qliquid": state_namespace.qliquid, - "qice": state_namespace.qice, - "qrain": state_namespace.qrain, - "qsnow": state_namespace.qsnow, - "qgraupel": state_namespace.qgraupel, - "qcld": state_namespace.qcld, - "qo3mr": state_namespace.qo3mr, - "qsgs_tke": state_namespace.qsgs_tke, - } - - self._mapn_tracer = MapNTracer( - self.stencil_factory, - self.quantity_factory, - abs(self.kord), - self.nq, - fill=self.fill, - tracers=tracers, - ) - - self._mapn_tracer(self._pe1, - self._pe2, - self._dp2, - tracers) - - self._map_single_w = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_wz, - -2, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - - self._map_single_delz = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_wz, - 1, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - self._map_single_w(state_namespace.w, - self._pe1, - self._pe2, - qs=state_namespace.ws_, - interp=False) - - self._rescale_delz_1( - state_namespace.delz, - state_namespace.delp, - ) - - self._map_single_delz(state_namespace.delz, - self._pe1, - self._pe2) - - self._rescale_delz_2( - state_namespace.delz, - self._dp2, - ) - - self._w_fix_consrv_moment( - state_namespace.w, - self._w2, - self._dp2, - self._gz, - state_namespace.w_max, - state_namespace.w_min, - self._compute_performed - ) - - self._map1_ppm_u = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_mt, - -1, - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], - ) - - self._pressures_mapu( - state_namespace.pe_, - state_namespace.ak, - state_namespace.bk, - self._pe0, - self._pe3, - state_namespace.ptop, - ) - - self._pe0_ptop_xmax( - self._pe0, - state_namespace.ptop, - ) - - self._map1_ppm_u( - state_namespace.u, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_u( - state_namespace.mfy, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_u( - state_namespace.cy, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_v = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_mt, - -1, - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], - ) - - self._pressures_mapv( - state_namespace.pe_, - state_namespace.ak, - state_namespace.bk, - self._pe0, - self._pe3, - ) - - self._map1_ppm_v( - state_namespace.v, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_v( - state_namespace.mfx_, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_v( - state_namespace.cx_, - self._pe0, - self._pe3, - interp=False, - ) - - self._pe_pk_delp_peln(state_namespace.pe_, - state_namespace.pk, - state_namespace.delp, - state_namespace.peln_3d, - self._pe2, - self._pk2, - self._pn2, - state_namespace.ak, - state_namespace.bk, - state_namespace.akap, - state_namespace.ptop, - ) - - self._moist_cv_pkz( - state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, - state_namespace.pkz, - state_namespace.pt, - state_namespace.cappa, - state_namespace.delp, - state_namespace.delz, - Float(state_namespace.r_vir), - ) - - # # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then - - # self._moist_cv_te(state_namespace.qvapor, - # state_namespace.qliquid, - # state_namespace.qrain, - # state_namespace.qsnow, - # state_namespace.qice, - # state_namespace.qgraupel, - # state_namespace.u, - # state_namespace.v, - # state_namespace.w, - # state_namespace.te_2d_, - # state_namespace.pt, - # state_namespace.phis_, - # state_namespace.delp, - # state_namespace.rsin2, - # state_namespace.cosa_s, - # state_namespace.hs, - # state_namespace.delz, - # state_namespace.grav, - # ) - - # self._te_zsum(state_namespace.te_2d_, - # state_namespace.te0_2d_, - # state_namespace.delp, - # state_namespace.pkz, - # self._zsum1, - # ) - - # # Note, since this is a serial translate test, mpp_global_sum won't work without the communicator. - # # Also, mpp_global_sum is currently set up for the C24 TBC setup - - # # tesum = mpp_global_sum(state_namespace.te_2d_*state_namespace.area_64, communicator, self.stencil_factory) - - # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - - # # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64, communicator, self.stencil_factory) - # # dtmp = tesum / (cv_air*zsum) - - # # If loop based on if ( last_step .and. (.not. adiabatic) ) then - # # self._most_cv_pt_last_step(state_namespace.qvapor, - # # state_namespace.qliquid, - # # state_namespace.qrain, - # # state_namespace.qsnow, - # # state_namespace.qice, - # # state_namespace.qgraupel, - # # self._gz, - # # state_namespace.pt, - # # state_namespace.pkz, - # # dtmp, - # # state_namespace.r_vir], - # # ) - - return self.outputs_list_from_state_list(state_list) + print("No serial test available") def compute_parallel(self, inputs, communicator): state = self.state_from_inputs(inputs) @@ -1438,58 +1157,77 @@ def compute_parallel(self, inputs, communicator): state_namespace.delz, Float(state_namespace.r_vir), ) - - # # May need if loop here based on if( last_step .and. (.not.do_adiabatic_init) ) then - - # self._moist_cv_te(state_namespace.qvapor, - # state_namespace.qliquid, - # state_namespace.qrain, - # state_namespace.qsnow, - # state_namespace.qice, - # state_namespace.qgraupel, - # state_namespace.u, - # state_namespace.v, - # state_namespace.w, - # state_namespace.te_2d_, - # state_namespace.pt, - # state_namespace.phis_, - # state_namespace.delp, - # state_namespace.rsin2, - # state_namespace.cosa_s, - # state_namespace.hs, - # state_namespace.delz, - # state_namespace.grav, - # ) - - # self._te_zsum(state_namespace.te_2d_, - # state_namespace.te0_2d_, - # state_namespace.delp, - # state_namespace.pkz, - # self._zsum1, - # ) + + if state_namespace.last_step and not state_namespace.do_adiabatic_init: + + if state_namespace.consv > state_namespace.consv_min: + + self._moist_cv_te(state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.u, + state_namespace.v, + state_namespace.w, + state_namespace.te_2d_, + state_namespace.pt, + self._phis, + state_namespace.delp, + state_namespace.rsin2, + state_namespace.cosa_s, + state_namespace.hs, + state_namespace.delz, + state_namespace.grav, + ) + + print("sum(te_2d): ", sum(state_namespace.te_2d_.data)) # nans here + print("sum(qliquid): ", sum(sum(state_namespace.qliquid.data))) + print("sum(qrain): ", sum(sum(state_namespace.qrain.data))) + print("sum(qsnow): ", sum(sum(state_namespace.qsnow.data))) + print("sum(qice): ", sum(sum(state_namespace.qice.data))) + print("sum(qgraupel): ", sum(sum(state_namespace.qgraupel.data))) + print("sum(u): ", sum(sum(state_namespace.u.data))) # nans here + print("sum(v): ", sum(sum(state_namespace.v.data))) + print("sum(w): ", sum(sum(state_namespace.w.data))) + print("sum(pt): ", sum(sum(state_namespace.pt.data))) + print("sum(delp): ", sum(sum(state_namespace.delp.data))) + print("sum(rsin2): ", sum(state_namespace.rsin2.data)) + print("sum(cosa_s): ", sum(state_namespace.cosa_s.data)) + print("sum(hs): ", sum(state_namespace.hs.data)) + print("sum(delz): ", sum(sum(state_namespace.delz.data))) + + # self._te_zsum(state_namespace.te_2d_, + # state_namespace.te0_2d_, + # state_namespace.delp, + # state_namespace.pkz, + # self._zsum1, + # ) - # # Note, since this is a serial translate test, mpp_global_sum won't work without the communicator. - # # Also, mpp_global_sum is currently set up for the C24 TBC setup - - # # tesum = mpp_global_sum(state_namespace.te_2d_*state_namespace.area_64, communicator, self.stencil_factory) + # # Note, mpp_global_sum is currently set up for the C24 TBC setup + # print("te_2d type :", type(state_namespace.te_2d_.data[0,0])) + # print("area_64 type :", type(state_namespace.area_64_.data[0,0])) + # tesum = mpp_global_sum(state_namespace.te_2d_.data*state_namespace.area_64_.data, communicator, self.stencil_factory) + # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64_.data, communicator, self.stencil_factory) + # dtmp = tesum / (state_namespace.cv_air.data * zsum) # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - # # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64, communicator, self.stencil_factory) - # # dtmp = tesum / (cv_air*zsum) - - # # If loop based on if ( last_step .and. (.not. adiabatic) ) then - # # self._most_cv_pt_last_step(state_namespace.qvapor, - # # state_namespace.qliquid, - # # state_namespace.qrain, - # # state_namespace.qsnow, - # # state_namespace.qice, - # # state_namespace.qgraupel, - # # self._gz, - # # state_namespace.pt, - # # state_namespace.pkz, - # # dtmp, - # # state_namespace.r_vir], - # # ) + + # if state_namespace.last_step and not state_namespace.adiabatic: + + # self._most_cv_pt_last_step(state_namespace.qvapor, + # state_namespace.qliquid, + # state_namespace.qrain, + # state_namespace.qsnow, + # state_namespace.qice, + # state_namespace.qgraupel, + # self._gz, + # state_namespace.pt, + # state_namespace.pkz, + # dtmp, + # state_namespace.r_vir, + # ) return self.outputs_from_state(state) \ No newline at end of file From ac0774744ca37df9fb08a5630c41500a8f760485 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 17 Jan 2025 15:21:41 -0800 Subject: [PATCH 121/252] Adding serialized fv_mpaz.F90.SER.v2 --- fv_mapz.F90.SER.v2 | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/fv_mapz.F90.SER.v2 b/fv_mapz.F90.SER.v2 index 338cef68..cb56d14a 100644 --- a/fv_mapz.F90.SER.v2 +++ b/fv_mapz.F90.SER.v2 @@ -230,11 +230,11 @@ contains real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32, te0_2d_f32, zsum1_f32 - real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64 + real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64_ logical :: serial_flag - real :: grav_ + real :: grav_, consv_min_, cv_air_ !$ser verbatim real :: w_max_, w_min_ @@ -255,6 +255,8 @@ contains !$ser verbatim pk2_3d = pk !$ser verbatim gz_2d = 0.0 !$ser verbatim phis_3d = 0.0 + !$ser verbatim consv_min_=consv_min + !$ser verbatim cv_air_=cv_air remap_t = .false. remap_pt = .false. @@ -378,7 +380,7 @@ contains !$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & !$OMP peln_3d, gz_2d, pe0_3d, pe3_3d, phis_3d, te_2d, te_2d_f32, te0_2d, te0_2d_f32) & !$OMP private(gz,cvm,bkh,dp2, & -!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64) +!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64_) !$ser verbatim do k = 1, km+1 !$ser verbatim do j = js, je @@ -401,7 +403,7 @@ contains !$ser verbatim cosa_s=gridstruct%cosa_s !$ser verbatim grav_=grav !$ser verbatim te_2d_f32 = real(te_2d) -!$ser verbatim area_64 = real(gridstruct%area_64) +!$ser verbatim area_64_ = real(gridstruct%area_64) !$ser verbatim te0_2d_f32 = real(te0_2d) !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d @@ -456,9 +458,15 @@ contains !$ser data grav=grav_ !$ser data te_2d_=te_2d_f32 !$ser data te0_2d_=te0_2d_f32 -!$ser data area_64=area_64 +!$ser data area_64_=area_64_ !$ser data te=te - +!$ser data last_step=last_step +!$ser data do_adiabatic_init=do_adiabatic_init +!$ser data consv=consv +!$ser data consv_min=consv_min_ +!$ser data cv_air=cv_air_ +!$ser data adiabatic=adiabatic + do 1000 j=js,je+1 do k=1,km+1 @@ -1427,7 +1435,7 @@ contains !$ser data grav=grav_ !$ser data te_2d_=te_2d_f32 !$ser data te0_2d_=te0_2d_f32 -!$ser data area_64=area_64 +!$ser data area_64_=area_64_ !$ser data te=te ! Do total energy conservation and fast saturation adjustment as requested From 4757f6136195eac0802f029b5fae544f1b96dffc Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 17 Jan 2025 17:25:39 -0800 Subject: [PATCH 122/252] Unified fv_mapz serialization file --- fv_mapz.F90.SER | 233 +-- fv_mapz.F90.SER.v2 | 4836 -------------------------------------------- 2 files changed, 101 insertions(+), 4968 deletions(-) delete mode 100644 fv_mapz.F90.SER.v2 diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index ddd4b1c7..cb56d14a 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -230,11 +230,11 @@ contains real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32, te0_2d_f32, zsum1_f32 - real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64 + real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64_ logical :: serial_flag - real :: grav_ + real :: grav_, consv_min_, cv_air_ !$ser verbatim real :: w_max_, w_min_ @@ -255,6 +255,8 @@ contains !$ser verbatim pk2_3d = pk !$ser verbatim gz_2d = 0.0 !$ser verbatim phis_3d = 0.0 + !$ser verbatim consv_min_=consv_min + !$ser verbatim cv_air_=cv_air remap_t = .false. remap_pt = .false. @@ -378,59 +380,33 @@ contains !$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & !$OMP peln_3d, gz_2d, pe0_3d, pe3_3d, phis_3d, te_2d, te_2d_f32, te0_2d, te0_2d_f32) & !$OMP private(gz,cvm,bkh,dp2, & -!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64) - do 1000 j=js,je+1 +!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64_) -!$ser verbatim if(j == js) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe1_3d(i,j,k) = pe1(i,k) -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is-1, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim pe_3d(i,j-1,k) = pe(i,k,j-1) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim dp2_3d(i,j,k) = dp2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo !$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) +!$ser verbatim do j = js, je +!$ser verbatim do i = is, ie +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo !$ser verbatim enddo !$ser verbatim enddo -!$ser verbatim qmin=t_min -!$ser verbatim w_max_ = w_max -!$ser verbatim w_min_ = w_min -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim phis_3d(i,j,k) = phis(i,k) -!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js-1, je+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim enddo !$ser verbatim enddo -!$ser verbatim rsin2=gridstruct%rsin2 -!$ser verbatim cosa_s=gridstruct%cosa_s -!!$ser verbatim do i = is, ie -!!$ser verbatim te_2d_f32(i,:) = te_2d(i,:) -!!$ser verbatim area_64(i,:) = gridstruct%area_64(i,:) -!!$ser verbatim enddo -!$ser verbatim grav_=grav -!$ser verbatim te_2d_f32 = real(te_2d) -!$ser verbatim area_64 = real(gridstruct%area_64) -!$ser verbatim te0_2d_f32 = real(te0_2d) +!$ser verbatim enddo +!$ser verbatim qmin=t_min +!$ser verbatim w_max_ = w_max +!$ser verbatim w_min_ = w_min +!$ser verbatim rsin2=gridstruct%rsin2 +!$ser verbatim cosa_s=gridstruct%cosa_s +!$ser verbatim grav_=grav +!$ser verbatim te_2d_f32 = real(te_2d) +!$ser verbatim area_64_ = real(gridstruct%area_64) +!$ser verbatim te0_2d_f32 = real(te0_2d) !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d -!$ser data pe1_=pe1_3d -!$ser data pe2_=pe2_3d -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d !$ser data ptop=ptop !$ser data qvapor=q(:,:,:,sphum) !$ser data qliquid=q(:,:,:,liq_wat) @@ -448,12 +424,8 @@ contains !$ser data delz=delz !$ser data ak=ak !$ser data bk=bk -!$ser data dp2_3d=dp2_3d !$ser data ps=ps -!$ser data pn1_3d=pn1_3d -!$ser data pn2_3d=pn2_3d !$ser data peln_3d=peln_3d -!$ser data pk2_3d=pk2_3d !$ser data r_vir=r_vir !$ser data cvm=cvm !$ser data pk=pk @@ -486,12 +458,16 @@ contains !$ser data grav=grav_ !$ser data te_2d_=te_2d_f32 !$ser data te0_2d_=te0_2d_f32 -!$ser data phis_=phis_3d -!$ser data area_64=area_64 -!!$ser data gz=gz -!!$ser data k1k=k1k -!!$ser data remap_t=remap_t -!$ser verbatim endif +!$ser data area_64_=area_64_ +!$ser data te=te +!$ser data last_step=last_step +!$ser data do_adiabatic_init=do_adiabatic_init +!$ser data consv=consv +!$ser data consv_min=consv_min_ +!$ser data cv_air=cv_air_ +!$ser data adiabatic=adiabatic + + do 1000 j=js,je+1 do k=1,km+1 do i=is,ie @@ -1395,6 +1371,73 @@ contains 2000 continue +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js, je +!$ser verbatim do i = is, ie +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js-1, je+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser savepoint Remapping_GEOS-Out +!$ser data pe_=pe_3d +!$ser data qvapor=q(:,:,:,sphum) +!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qice=q(:,:,:,ice_wat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qgraupel=q(:,:,:,graupel) +!$ser data qcld=q(:,:,:,cld_amt) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data peln_3d=peln_3d +!$ser data r_vir=r_vir +!$ser data cvm=cvm +!$ser data pk=pk +!$ser data pkz=pkz +!$ser data akap=akap +!$ser data t_min=qmin +!$ser data kord_wz=kord_wz +!$ser data ws_=ws +!$ser data w=w +!$ser data u=u +!$ser data v=v +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser data w_max=w_max_ +!$ser data w_min=w_min_ +!$ser data kord_mt=kord_mt +!$ser data cosa_s=cosa_s +!$ser data rsin2=rsin2 +!$ser data hs=hs +!$ser data grav=grav_ +!$ser data te_2d_=te_2d_f32 +!$ser data te0_2d_=te0_2d_f32 +!$ser data area_64_=area_64_ +!$ser data te=te + ! Do total energy conservation and fast saturation adjustment as requested ! and fill new PT (Theta_V) for next k_split step or export dry T @@ -1513,80 +1556,6 @@ contains !$ser data te_2d_=te_2d_f32 !$ser verbatim endif - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim pk2_3d(i,j,k) = pk2(i,k) - !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) - !$ser verbatim peln_3d(i,j,k) = peln(i,k,j) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie+1 - !$ser verbatim pe0_3d(i,j,k) = pe0(i,k) - !$ser verbatim pe3_3d(i,j,k) = pe3(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do i = is, ie - !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) - !$ser verbatim enddo - !$ser savepoint Remapping_GEOS-Out - !$ser data pe1_=pe1_3d pe2_=pe2_3d - !$ser data pe_=pe_3d - !$ser data peln_3d=peln_3d - !$ser data pk=pk - !$ser data pkz=pkz - !$ser data q_con=q_con - !$ser data pt=pt - !$ser data cappa=cappa - !$ser data delp=delp - !$ser data delz=delz - !$ser data ps=ps - !$ser data dp2_3d=dp2_3d - !$ser data pn1_3d=pn1_3d - !$ser data pn2_3d=pn2_3d - !$ser data pk2_3d=pk2_3d - - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) - - !$ser data w=w - - !$ser data pe0_=pe0_3d - !$ser data pe3_=pe3_3d - !$ser data u=u - !$ser verbatim if (present(mfy)) then - !$ser data mfy=mfy - !$ser verbatim endif - !$ser verbatim if (present(cy)) then - !$ser data cy=cy - !$ser verbatim endif - !$ser data v=v - !$ser verbatim if (present(mfx)) then - !$ser data mfx_=mfx - !$ser verbatim endif - !$ser verbatim if (present(cx)) then - !$ser data cx_=cx - !$ser verbatim endif - !$ser data te_2d_=te_2d_f32 - !$ser verbatim endif - endif ! end non-hydro !$ser verbatim if(j == js2d) then diff --git a/fv_mapz.F90.SER.v2 b/fv_mapz.F90.SER.v2 deleted file mode 100644 index cb56d14a..00000000 --- a/fv_mapz.F90.SER.v2 +++ /dev/null @@ -1,4836 +0,0 @@ -!*********************************************************************** -!* GNU Lesser General Public License -!* -!* This file is part of the FV3 dynamical core. -!* -!* The FV3 dynamical core is free software: you can redistribute it -!* and/or modify it under the terms of the -!* GNU Lesser General Public License as published by the -!* Free Software Foundation, either version 3 of the License, or -!* (at your option) any later version. -!* -!* The FV3 dynamical core is distributed in the hope that it will be -!* useful, but WITHOUT ANYWARRANTY; without even the implied warranty -!* of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. -!* See the GNU General Public License for more details. -!* -!* You should have received a copy of the GNU Lesser General Public -!* License along with the FV3 dynamical core. -!* If not, see . -!*********************************************************************** - -!>@brief The module 'fv_mapz' contains the vertical mapping routines \cite lin2004vertically -!>@note April 12, 2012 -SJL: This revision may actually produce rounding level differences -!! due to the elimination of KS to compute pressure level for remapping. - -module fv_mapz_mod - -! Modules Included: -! -! -! -! -! -!
Module NameFunctions Included
-! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -! -!
constants_modradius, pi=>pi_8, rvgas, rdgas, grav, hlv, hlf, cp_air, cp_vapor
field_manager_modMODEL_ATMOS
fv_arrays_modfv_grid_type
fv_cmp_modqs_init, fv_sat_adj
fv_fill_modfillz
fv_grid_utils_modg_sum, ptop_min
fv_mp_modis_master
fv_timing_modtiming_on, timing_off
fv_tracer2d_modtracer_2d, tracer_2d_1L, tracer_2d_nested
mpp_mod/td> -! NOTE, mpp_error, get_unit, mpp_root_pe, mpp_pe
mpp_domains_mod/td> -! mpp_update_domains, domain2d
tracer_manager_modget_tracer_index
- !$ser verbatim USE m_serialize, ONLY: fs_is_serialization_on - use constants_mod, only: radius, pi=>pi_8, rvgas, rdgas, grav, hlv, hlf, hls, cp_air, cp_vapor - use tracer_manager_mod,only: get_tracer_index - use field_manager_mod, only: MODEL_ATMOS - use fv_grid_utils_mod, only: g_sum, ptop_min - use fv_fill_mod, only: fillz - use mpp_domains_mod, only: mpp_update_domains, domain2d, mpp_global_sum, BITWISE_EFP_SUM, BITWISE_EXACT_SUM - use mpp_mod, only: NOTE, mpp_error, get_unit, mpp_root_pe, mpp_pe - use fv_arrays_mod, only: fv_grid_type, fv_flags_type - use fv_timing_mod, only: timing_on, timing_off - use fv_mp_mod, only: is_master - use fv_cmp_mod, only: qs_init, fv_sat_adj - - implicit none - real, parameter:: consv_min= 0.001 !< below which no correction applies - real, parameter:: te_min= -1.e25 - real, parameter:: t_min= 184. !< below which applies stricter constraint - real, parameter:: r2=1./2., r0=0.0 - real, parameter:: r3 = 1./3., r23 = 2./3., r12 = 1./12. - real, parameter:: cv_vap = 3.*rvgas !< 1384.5 - real, parameter:: cv_air = cp_air - rdgas !< = rdgas * (7/2-1) = 2.5*rdgas=717.68 -! real, parameter:: c_ice = 2106. !< heat capacity of ice at 0.C - real, parameter:: c_ice = 1972. !< heat capacity of ice at -15.C - real, parameter:: c_liq = 4.1855e+3 !< GFS: heat capacity of water at 0C -! real, parameter:: c_liq = 4218. !< ECMWF-IFS - real, parameter:: cp_vap = cp_vapor !< 1846. - real, parameter:: tice = 273.16 - - logical, parameter :: w_limiter = .true. - real, parameter :: w_max = 90. - real, parameter :: w_min = -60. - - real(kind=8) :: E_Flux = 0. - private - - public compute_total_energy, Lagrangian_to_Eulerian, moist_cv, moist_cp, & - rst_remap, mappm, E_Flux, mapn_tracer, map1_q2 - -!---- version number ----- - character(len=128) :: version = '$Id$' - character(len=128) :: tagname = '$Name$' - -contains - -!>@brief The subroutine 'Lagrangian_to_Eulerian' remaps deformed Lagrangian layers back to the reference Eulerian coordinate. -!>@details It also includes the entry point for calling fast microphysical processes. This is typically calle on the k_split loop. - subroutine Lagrangian_to_Eulerian(last_step, consv, ps, pe, delp, pkz, pk, & - mdt, pdt, km, is,ie,js,je, isd,ied,jsd,jed, & - nq, nwat, sphum, q_con, u, v, w, delz, pt, q, hs, r_vir, cp, & - akap, cappa, kord_mt, kord_wz, kord_tr, kord_tm, peln, te0_2d, & - ng, ua, va, omga, te, ws, fill, reproduce_sum, out_dt, dtdt, & - ptop, ak, bk, pfull, flagstruct, gridstruct, domain, do_sat_adj, & - hydrostatic, hybrid_z, do_omega, adiabatic, do_adiabatic_init, & - mfx, mfy, cx, cy, remap_option, gmao_remap) - logical, intent(in):: last_step - real, intent(in):: mdt !< remap time step - real, intent(in):: pdt !< phys time step - integer, intent(in):: km - integer, intent(in):: nq !< number of tracers (including h2o) - integer, intent(in):: nwat - integer, intent(in):: sphum !< index for water vapor (specific humidity) - integer, intent(in):: ng - integer, intent(in):: is,ie,isd,ied !< starting & ending X-Dir index - integer, intent(in):: js,je,jsd,jed !< starting & ending Y-Dir index - integer, intent(in):: kord_mt !< Mapping order for the vector winds - integer, intent(in):: kord_wz !< Mapping order/option for w - integer, intent(in):: kord_tr(nq) !< Mapping order for tracers - integer, intent(in):: kord_tm !< Mapping order for thermodynamics - - real, intent(in):: consv !< factor for TE conservation - real, intent(in):: r_vir - real, intent(in):: cp - real, intent(in):: akap - real, intent(in):: hs(isd:ied,jsd:jed) !< surface geopotential - real(kind=8), intent(inout):: te0_2d(is:ie,js:je) - real, intent(in):: ws(is:ie,js:je) - - logical, intent(in):: do_sat_adj - logical, intent(in):: fill !< fill negative tracers - logical, intent(in):: reproduce_sum - logical, intent(in):: do_omega, adiabatic, do_adiabatic_init - real, intent(in) :: ptop - real, intent(in) :: ak(km+1) - real, intent(in) :: bk(km+1) - real, intent(in):: pfull(km) - type(fv_grid_type), intent(IN), target :: gridstruct - type(fv_flags_type), intent(INOUT) :: flagstruct - type(domain2d), intent(INOUT) :: domain - - ! INPUT/OUTPUT - real, intent(inout):: pk(is:ie,js:je,km+1) !< pe to the kappa - real, intent(inout):: q(isd:ied,jsd:jed,km,*) - real, intent(inout):: delp(isd:ied,jsd:jed,km) !< pressure thickness - real, intent(inout):: pe(is-1:ie+1,km+1,js-1:je+1) !< pressure at layer edges - real, intent(inout):: ps(isd:ied,jsd:jed) !< surface pressure - - ! u-wind will be ghosted one latitude to the north upon exit - real, intent(inout):: u(isd:ied ,jsd:jed+1,km) !< u-wind (m/s) - real, intent(inout):: v(isd:ied+1,jsd:jed ,km) !< v-wind (m/s) - real, intent(inout):: w(isd: ,jsd: ,1:) !< vertical velocity (m/s) - real, intent(inout):: pt(isd:ied ,jsd:jed ,km) !< cp*virtual potential temperature - !< as input; output: temperature - real, intent(inout), dimension(isd:,jsd:,1:)::delz, q_con, cappa - logical, intent(in):: hydrostatic - logical, intent(in):: hybrid_z - logical, intent(in):: out_dt - - real, intent(inout):: ua(isd:ied,jsd:jed,km) !< u-wind (m/s) on physics grid - real, intent(inout):: va(isd:ied,jsd:jed,km) !< v-wind (m/s) on physics grid - real, intent(inout):: omga(isd:ied,jsd:jed,km) !< vertical press. velocity (pascal/sec) - real, intent(inout):: peln(is:ie,km+1,js:je) !< log(pe) - real, intent(inout):: dtdt(is:ie,js:je,km) - real, intent(out):: pkz(is:ie,js:je,km) !< layer-mean pk for converting t to pt - real, intent(out):: te(isd:ied,jsd:jed,km) - ! Mass fluxes - real, optional, intent(inout):: mfx(is:ie+1,js:je ,km) ! X-dir Mass Flux - real, optional, intent(inout):: mfy(is:ie ,js:je+1,km) ! Y-dir Mass Flux - ! Courant numbers - real, optional, intent(inout):: cx(is:ie+1, jsd:jed,km) - real, optional, intent(inout):: cy(isd:ied ,js:je+1,km) - - integer, intent(in):: remap_option, gmao_remap - - ! !DESCRIPTION: - ! - ! !REVISION HISTORY: - ! SJL 03.11.04: Initial version for partial remapping - ! - !----------------------------------------------------------------------- - real(kind=8), dimension(is:ie,js:je):: te_2d, zsum0, zsum1 - real, dimension(is:ie,js:je):: dpln - real, dimension(is:ie,km) :: q2, dp2, w2 - real, dimension(is:ie,km+1):: pe1, pe2, pk1, pk2, pn1, pn2, phis - real, dimension(is:ie+1,km+1):: pe0, pe3 - real, dimension(is:ie):: gz, cvm - real(kind=8):: tesum, zsum, dtmp - real :: rcp, rg, tmp, tpe, rrg, bkh, k1k, dlnp - logical:: fast_mp_consv - integer:: i,j,k - integer:: nt, liq_wat, ice_wat, rainwat, snowwat, cld_amt, graupel, iq, n, kmp, kp, k_next - logical:: remap_t, remap_pt, remap_te - - real, dimension(is:ie, js:je, km+1) :: pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, peln_3d - real, dimension(is:ie, js:je, km+1) :: phis_3d - real, dimension(is-1:ie+1,js-1:je+1,km+1) :: pe_3d - real, dimension(is:ie+1, js:je+1, km+1) :: pe0_3d, pe3_3d - - real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d - - real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32, te0_2d_f32, zsum1_f32 - real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64_ - - logical :: serial_flag - - real :: grav_, consv_min_, cv_air_ - -!$ser verbatim real :: w_max_, w_min_ - - !$ser verbatim integer:: mode, abskord, iep1, iedp1, jedp1, js2d, o3mr, sgs_tke - !$ser verbatim real :: qmin - !$ser verbatim qmin = 0.0 - !$ser verbatim iep1=ie+1 - !$ser verbatim iedp1=ied+1 - !$ser verbatim jedp1=jed+1 - !$ser verbatim js2d=js - !$ser verbatim pe_3d = 0.0 - !$ser verbatim pe1_3d = 0.0 - !$ser verbatim pe2_3d = 0.0 - !$ser verbatim peln_3d = 0.0 - !$ser verbatim dp2_3d = 0.0 - !$ser verbatim pn1_3d = 0.0 - !$ser verbatim pn2_3d = 0.0 - !$ser verbatim pk2_3d = pk - !$ser verbatim gz_2d = 0.0 - !$ser verbatim phis_3d = 0.0 - !$ser verbatim consv_min_=consv_min - !$ser verbatim cv_air_=cv_air - - remap_t = .false. - remap_pt = .false. - remap_te = .false. - select case (remap_option) - case(0) - remap_t = .true. - case(1) - remap_pt = .true. - case(2) - remap_te = .true. - case default - print*, ' INVALID REMAP_OPTION ' - stop - end select - - select case (gmao_remap) - case(0) - ! use GFDL schemes - case(1) - ! GMAO linear remap - case(2) - ! GMAO quadratic remap - case(3) - ! GMAO cubic remap - case default - print*, ' INVALID GMAO_REMAP' - stop - end select - - if (is_master() .and. flagstruct%fv_debug) then - print*, '' - select case (remap_option) - case(0) - print*, ' REMAPPING T in logP' - case(1) - print*, ' REMAPPING PT in P' - case(2) - print*, ' REMAPPING TE in logP' - end select - - print*, '' - select case (gmao_remap) - case(0) - print*, ' Using GFDL schemes' - case(1) - print*, ' Using GMAO linear scheme' - case(2) - print*, ' Using GMAO quadratic scheme' - case(3) - print*, ' Using GMAO cubic scheme' - end select - - ! Total eergy conservation - print*, '' - print*, ' REMAPPING CONSV: ', consv - print*, ' REMAPPING CONSV_MIN: ', consv_min - print*, '' - endif - - k1k = rdgas/cv_air ! akap / (1.-akap) = rg/Cv=0.4 - rg = rdgas - rcp = 1./ cp - rrg = -rdgas/grav - -#ifdef MAPL_MODE - select case(nwat) - case(1) - liq_wat = -1 - ice_wat = -1 - rainwat = -1 - snowwat = -1 - graupel = -1 - cld_amt = -1 - case(3) - liq_wat = 2 - ice_wat = 3 - rainwat = -1 - snowwat = -1 - graupel = -1 - cld_amt = -1 - case(6:7) - liq_wat = 2 - ice_wat = 3 - rainwat = 4 - snowwat = 5 - graupel = 6 - cld_amt = 7 - end select -#else - liq_wat = get_tracer_index (MODEL_ATMOS, 'liq_wat') - ice_wat = get_tracer_index (MODEL_ATMOS, 'ice_wat') - rainwat = get_tracer_index (MODEL_ATMOS, 'rainwat') - snowwat = get_tracer_index (MODEL_ATMOS, 'snowwat') - graupel = get_tracer_index (MODEL_ATMOS, 'graupel') - cld_amt = get_tracer_index (MODEL_ATMOS, 'cld_amt') - !$ser verbatim o3mr = get_tracer_index (MODEL_ATMOS, 'o3mr') - !$ser verbatim sgs_tke = get_tracer_index (MODEL_ATMOS, 'sgs_tke') -#endif - - if ( do_sat_adj .and. nwat>=6 ) then - print*,'CODE NOT TESTED HERE 1' - fast_mp_consv = (.not.do_adiabatic_init) .and. consv>consv_min - do k=1,km - kmp = k - if ( pfull(k) > 10.E2 ) exit - enddo - call qs_init(kmp) - endif - -!$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & -!$OMP remap_pt,remap_te,mfy,mfx,cx,cy,hydrostatic, & -!$OMP pt,pk,rg,peln,q,nwat,liq_wat,rainwat,ice_wat,snowwat, & -#ifdef SERIALIZE -!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, cld_amt, mode,qmin, abskord,iep1, iedp1, jedp1, js2d, o3mr, sgs_tke, & -#endif -!$OMP graupel,q_con, sphum,cappa,r_vir,rcp,cp,k1k,delp, & -!$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & -!$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, adiabatic, & -!$OMP hs,w,ws,kord_wz,rrg,kord_mt,consv,remap_option,gmao_remap,& -!$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & -!$OMP peln_3d, gz_2d, pe0_3d, pe3_3d, phis_3d, te_2d, te_2d_f32, te0_2d, te0_2d_f32) & -!$OMP private(gz,cvm,bkh,dp2, & -!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64_) - -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js, je -!$ser verbatim do i = is, ie -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js-1, je+1 -!$ser verbatim do i = is-1, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim qmin=t_min -!$ser verbatim w_max_ = w_max -!$ser verbatim w_min_ = w_min -!$ser verbatim rsin2=gridstruct%rsin2 -!$ser verbatim cosa_s=gridstruct%cosa_s -!$ser verbatim grav_=grav -!$ser verbatim te_2d_f32 = real(te_2d) -!$ser verbatim area_64_ = real(gridstruct%area_64) -!$ser verbatim te0_2d_f32 = real(te0_2d) -!$ser savepoint Remapping_GEOS-In -!$ser data pe_=pe_3d -!$ser data ptop=ptop -!$ser data qvapor=q(:,:,:,sphum) -!$ser data qliquid=q(:,:,:,liq_wat) -!$ser data qice=q(:,:,:,ice_wat) -!$ser data qrain=q(:,:,:,rainwat) -!$ser data qsnow=q(:,:,:,snowwat) -!$ser data qgraupel=q(:,:,:,graupel) -!$ser data qcld=q(:,:,:,cld_amt) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ak=ak -!$ser data bk=bk -!$ser data ps=ps -!$ser data peln_3d=peln_3d -!$ser data r_vir=r_vir -!$ser data cvm=cvm -!$ser data pk=pk -!$ser data pkz=pkz -!$ser data akap=akap -!$ser data t_min=qmin -!$ser data kord_wz=kord_wz -!$ser data ws_=ws -!$ser data w=w -!$ser data u=u -!$ser data v=v -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser data w_max=w_max_ -!$ser data w_min=w_min_ -!$ser data kord_mt=kord_mt -!$ser data cosa_s=cosa_s -!$ser data rsin2=rsin2 -!$ser data hs=hs -!$ser data grav=grav_ -!$ser data te_2d_=te_2d_f32 -!$ser data te0_2d_=te0_2d_f32 -!$ser data area_64_=area_64_ -!$ser data te=te -!$ser data last_step=last_step -!$ser data do_adiabatic_init=do_adiabatic_init -!$ser data consv=consv -!$ser data consv_min=consv_min_ -!$ser data cv_air=cv_air_ -!$ser data adiabatic=adiabatic - - do 1000 j=js,je+1 - - do k=1,km+1 - do i=is,ie - pe1(i,k) = pe(i,k,j) - enddo - enddo - - do i=is,ie - pe2(i, 1) = ptop - pe2(i,km+1) = pe(i,km+1,j) - enddo - - if ( j /= (je+1) ) then - - if (remap_t) then - ! Remap T in logP - ! Note: pt at this stage is Theta_v - if ( hydrostatic ) then - print*,'CODE NOT TESTED HERE 2' - ! Transform virtual pt to virtual Temp - do k=1,km - do i=is,ie - pt(i,j,k) = pt(i,j,k)*(pk(i,j,k+1)-pk(i,j,k))/(akap*(peln(i,k+1,j)-peln(i,k,j))) - enddo - enddo - else - ! Transform "density pt" to "density temp" - !$ser verbatim if(j == js2d) then - !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' - !$ser savepoint MoistCVPlusPt_2d-In - !$ser data qvapor_js=q(:,j,:,sphum) qliquid_js=q(:,j,:,liq_wat) qice_js=q(:,j,:,ice_wat) qrain_js=q(:,j,:,rainwat) qsnow_js=q(:,j,:,snowwat) qgraupel_js=q(:,j,:,graupel) qcld_js=q(:,j,:,cld_amt) gz1d=gz cvm=cvm r_vir=r_vir cappa=cappa rrg=rrg delp=delp delz=delz pt=pt k1k=k1k j_2d=js2d q_con=q_con - !$ser verbatim endif - do k=1,km -#ifdef MOIST_CAPPA - call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, gz, cvm) - do i=is,ie - !$ser verbatim q_con(i,j,k) = gz(i) - cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) - pt(i,j,k) = pt(i,j,k)*exp(cappa(i,j,k)/(1.-cappa(i,j,k))*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - enddo -#else - do i=is,ie - pt(i,j,k) = pt(i,j,k)*exp(k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - enddo -#endif - enddo - !$ser verbatim if(j == js2d) then - !$ser savepoint MoistCVPlusPt_2d-Out - !$ser data gz1d=gz cvm=cvm pt=pt cappa=cappa q_con=q_con - !$ser verbatim endif - - endif ! hydro test - elseif (remap_pt) then - print*,'CODE NOT TESTED HERE 3' - ! Remap PT in P - ! pt is already virtual PT - elseif (remap_te) then - print*,'CODE NOT TESTED HERE 4' - ! Remap TE in logP - ! Transform virtual pt to total energy - if ( hydrostatic ) then - call pkez(km, is, ie, js, je, j, pe, pk, akap, peln, pkz, ptop) - do i=is,ie - phis(i,km+1) = hs(i,j) - enddo - do k=km,1,-1 - do i=is,ie - phis(i,k) = phis(i,k+1) + cp_air*pt(i,j,k)*(pk(i,j,k+1)-pk(i,j,k)) - enddo - enddo - do k=1,km+1 - do i=is,ie - phis(i,k) = phis(i,k) * pe(i,k,j) - enddo - enddo - ! Compute cp*T + KE - do k=1,km - do i=is,ie - te(i,j,k) = 0.25*gridstruct%rsin2(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & - v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & - + cp_air*pt(i,j,k)*pkz(i,j,k) & - + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) - enddo - enddo - else - ! TE using 3D winds (pt is virtual potential temperature): - do i=is,ie - phis(i,km+1) = hs(i,j) - enddo - do k=km,1,-1 - do i=is,ie - phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) - enddo - enddo - do k=1,km+1 - do i=is,ie - phis(i,k) = phis(i,k) * pe(i,k,j) - enddo - enddo - do k=1,km -#ifdef MOIST_CAPPA - call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, gz, cvm) - do i=is,ie - cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) - pkz(i,j,k) = exp(cappa(i,j,k)/(1.-cappa(i,j,k))*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - ! TE = KE + Cv*T_v + PE - te(i,j,k) = 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & - u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & - + cvm(i)*pt(i,j,k)*pkz(i,j,k) & - + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) - enddo -#else - do i=is,ie - pkz(i,j,k) = exp(k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - ! TE = KE + Cv*T_v + PE - te(i,j,k) = 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & - u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & - + cv_air*pt(i,j,k)*pkz(i,j,k) & - + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) - enddo -#endif - enddo - endif ! hydro test - endif - - ! update ps - do i=is,ie - ps(i,j) = pe1(i,km+1) - enddo - ! - ! Hybrid sigma-P coordinate: - ! - do k=2,km - do i=is,ie - pe2(i,k) = ak(k) + bk(k)*pe(i,km+1,j) - enddo - enddo - do k=1,km - do i=is,ie - dp2(i,k) = pe2(i,k+1) - pe2(i,k) - enddo - enddo - !------------------ - ! Compute p**Kappa - !------------------ - do k=1,km+1 - do i=is,ie - pk1(i,k) = pk(i,j,k) - pn1(i,k) = peln(i,k,j) - enddo - enddo - - do i=is,ie - pn2(i, 1) = peln(i, 1,j) - pn2(i,km+1) = peln(i,km+1,j) - pk2(i, 1) = pk1(i, 1) - pk2(i,km+1) = pk1(i,km+1) - enddo - - do k=2,km - do i=is,ie - pn2(i,k) = log(pe2(i,k)) - pk2(i,k) = exp(akap*pn2(i,k)) - enddo - enddo - - if (remap_te) then - print*,'CODE NOT TESTED HERE 5' - !---------------------------------- - ! map TE in log P - !---------------------------------- - if ( gmao_remap > 0 ) then - call map1_gmao (km, pe1, te, & - km, pe2, te, & - is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.true.) - else - !$ser verbatim if(j == js2d) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint MapScalar_2d-In - !$ser verbatim mode=1 - !$ser data pe1_=pn1_3d pt=pt q1=te mode=mode j_2d=js2d q_min=-1.e25 - !$ser verbatim endif - call map_scalar(km, pn1, te, gz, & - km, pn2, te, & - is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), te_min) - !$ser verbatim if(j == js2d) then - !$ser savepoint MapScalar_2d-Out - !$ser data q1=te - !$ser verbatim endif - endif - else - !---------------------------------- - ! map T or PT in log P - !---------------------------------- - if ( gmao_remap > 0 ) then - print*,'CODE NOT TESTED HERE 6' - call map1_gmao (km, pe1, pt, & - km, pe2, pt, & - is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.false.) - else - !$ser verbatim if(j == js2d) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Map_Scalar-In - !$ser verbatim qmin=184.0 - !$ser data pe1_=pn1_3d pe2_=pn2_3d q1=pt j_2d=js2d q_min=qmin - !$ser verbatim endif - call map_scalar(km, pn1, pt, gz, & - km, pn2, pt, & - is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), t_min) - !$ser verbatim if(j == js2d) then - !$ser savepoint Map_Scalar-Out - !$ser data q1=pt - !$ser verbatim endif - endif - endif - - !---------------- - ! Map constituents - !---------------- - if( nq > 5 ) then - !$ser verbatim if(j == js2d) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) - !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim - !$ser verbatim if (k < km+1) dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint MapN_Tracer_2d-In - !$ser data j_2d=js2d nq=nq pe1_=pe1_3d pe2_=pe2_3d dp2_=dp2_3d - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) - !$ser verbatim endif - call mapn_tracer(nq, km, pe1, pe2, q, dp2, kord_tr, j, & - is, ie, isd, ied, jsd, jed, 0., fill) - !$ser verbatim if(j == js2d) then - !$ser savepoint MapN_Tracer_2d-Out - !$ser data qvapor=q(:,:,:,1) - !$ser data qliquid=q(:,:,:,2) - !$ser data qice=q(:,:,:,3) - !$ser data qrain=q(:,:,:,4) - !$ser data qsnow=q(:,:,:,5) - !$ser data qgraupel=q(:,:,:,6) - !$ser data qcld=q(:,:,:,7) - !$ser data qo3mr=q(:,:,:,8) - !$ser data qsgs_tke=q(:,:,:,9) - !$ser verbatim endif - - elseif ( nq > 0 ) then - print*,'CODE NOT TESTED HERE 7' - ! Remap one tracer at a time - do iq=1,nq - call map1_q2(km, pe1, q(isd,jsd,1,iq), & - km, pe2, q2, dp2, & - is, ie, 0, kord_tr(iq), j, isd, ied, jsd, jed, 0.) - if (fill) call fillz(ie-is+1, km, 1, q2, dp2) - do k=1,km - do i=is,ie - q(i,j,k,iq) = q2(i,k) - enddo - enddo - enddo - endif - - if ( .not. hydrostatic ) then - ! Remap vertical wind: - !$ser verbatim if(j == js2d) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Map1_PPM_W-In - !$ser data pe1_=pe1_3d - !$ser data pe2_=pe2_3d - !$ser data ws_=ws - !$ser data w_=w - !$ser data kord_wz=kord_wz - !$ser verbatim endif - call map1_ppm (km, pe1, w, ws(is,j), & - km, pe2, w, & - is, ie, j, isd, ied, jsd, jed, -2, kord_wz) - !$ser verbatim if(j == js2d) then - !$ser savepoint Map1_PPM_W-Out - !$ser data w_=w - !$ser verbatim endif - - !$ser verbatim if(j == js2d) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim do i = is, ie - !$ser verbatim gz_2d(i,j) = gz(i) - !$ser verbatim enddo - !$ser verbatim do k = 1, km - !$ser verbatim do i = is, ie - !$ser verbatim dp2_3d(i,j,k) = dp2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint Map1_PPM_delz-In - !$ser data pe1_=pe1_3d - !$ser data pe2_=pe2_3d - !$ser data dp2_3d=dp2_3d - !$ser data gz_=gz_2d - !$ser data delz_=delz - !$ser data delp=delp - !$ser data kord_wz=kord_wz - !$ser verbatim endif - ! Remap delz for hybrid sigma-p coordinate - do k=1,km - do i=is,ie - delz(i,j,k) = -delz(i,j,k) / delp(i,j,k) ! ="specific volume"/grav - enddo - enddo - - call map1_ppm (km, pe1, delz, gz, & - km, pe2, delz, & - is, ie, j, isd, ied, jsd, jed, 1, abs(kord_wz)) - - do k=1,km - do i=is,ie - delz(i,j,k) = -delz(i,j,k)*dp2(i,k) - enddo - enddo - !$ser verbatim if(j == js2d) then - !$ser savepoint Map1_PPM_delz-Out - !$ser data delz_=delz - !$ser verbatim endif - - !Fix excessive w - momentum conserving --- sjl - ! gz(:) used here as a temporary array - if ( w_limiter ) then -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km -!$ser verbatim do i = is, ie -!$ser verbatim dp2_3d(i,:,k) = dp2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim w_max_ = w_max -!$ser verbatim w_min_ = w_min - -!$ser savepoint W_fix_consrv_moment-In -!$ser data w=w dp2_W=dp2_3d w_max=w_max_ w_min=w_min_ -!$ser verbatim endif - do k=1,km - do i=is,ie - w2(i,k) = w(i,j,k) - enddo - enddo - do k=1, km-1 - do i=is,ie - if ( w2(i,k) > w_max ) then - if(j == js2d) then - print*,"ENTERED LOOP 1" - endif - gz(i) = (w2(i,k)-w_max) * dp2(i,k) - w2(i,k ) = w_max - w2(i,k+1) = w2(i,k+1) + gz(i)/dp2(i,k+1) - !print*, ' W_LIMITER down: ', i,j,k, w2(i,k:k+1), w(i,j,k:k+1) - elseif ( w2(i,k) < w_min ) then - if(j == js2d) then - print*,"ENTERED LOOP 2" - endif - gz(i) = (w2(i,k)-w_min) * dp2(i,k) - w2(i,k ) = w_min - w2(i,k+1) = w2(i,k+1) + gz(i)/dp2(i,k+1) - !print*, ' W_LIMITER down: ', i,j,k, w2(i,k:k+1), w(i,j,k:k+1) - endif - enddo - enddo - do k=km, 2, -1 - do i=is,ie - if ( w2(i,k) > w_max ) then - if(j == js2d) then - print*,"ENTERED LOOP 3" - endif - gz(i) = (w2(i,k)-w_max) * dp2(i,k) - w2(i,k ) = w_max - w2(i,k-1) = w2(i,k-1) + gz(i)/dp2(i,k-1) - !print*, ' W_LIMITER up: ', i,j,k, w2(i,k-1:k), w(i,j,k-1:k) - elseif ( w2(i,k) < w_min ) then - if(j == js2d) then - print*,"ENTERED LOOP 4" - endif - gz(i) = (w2(i,k)-w_min) * dp2(i,k) - w2(i,k ) = w_min - w2(i,k-1) = w2(i,k-1) + gz(i)/dp2(i,k-1) - !print*, ' W_LIMITER up: ', i,j,k, w2(i,k-1:k), w(i,j,k-1:k) - endif - enddo - enddo - do i=is,ie - if (w2(i,1) > w_max*2. ) then - if(j == js2d) then - print*,'ENTERED LOOP 5' - endif - w2(i,1) = w_max*2 ! sink out of the top of the domain - !print*, ' W_LIMITER top limited: ', i,j,1, w2(i,1), w(i,j,1) - elseif (w2(i,1) < w_min*2. ) then - if(j == js2d) then - print*,'ENTERED LOOP 6' - endif - w2(i,1) = w_min*2. - !print*, ' W_LIMITER top limited: ', i,j,1, w2(i,1), w(i,j,1) - endif - enddo - do k=1,km - do i=is,ie - w(i,j,k) = w2(i,k) - enddo - enddo -!$ser verbatim if(j == js2d) then -!$ser savepoint W_fix_consrv_moment-Out -!$ser data w=w -!$ser verbatim endif - endif - endif - - endif !(j < je+1) - -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is-1, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim pe_3d(i,j-1,k) = pe(i,k,j-1) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Pressures_mapU-In -!$ser data pe_=pe_3d -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data ak=ak -!$ser data bk=bk -!$ser data ptop=ptop -!$ser data u_=u -!$ser data kord_mt=kord_mt -!$ser verbatim if (present(mfy)) then -!$ser data mfy_=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy_=cy -!$ser verbatim endif -!$ser verbatim endif - !------ - ! map u - !------ - do i=is,ie+1 - pe0(i,1) = ptop - enddo - - do k=2,km+1 - do i=is,ie - pe0(i,k) = 0.5*(pe(i,k,j-1)+pe(i,k,j)) - enddo - enddo - - do k=1,km+1 - bkh = 0.5*bk(k) - do i=is,ie - pe3(i,k) = ak(k) + bkh*(pe(i,km+1,j-1)+pe(i,km+1,j)) - enddo - enddo - - call map1_ppm( km, pe0(is:ie,:), u, gz, & - km, pe3(is:ie,:), u, & - is, ie, j, isd, ied, jsd, jed+1, -1, kord_mt) - if (present(mfy)) then - ! print*, 'mfy present' - call map1_ppm( km, pe0(is:ie,:), mfy, gz, & - km, pe3(is:ie,:), mfy, & - is, ie, j, is, ie, js, je+1, -1, kord_mt) - endif - if (present(cy)) then - ! print*, 'cy present' - call map1_ppm( km, pe0(is:ie,:), cy, gz, & - km, pe3(is:ie,:), cy, & - is, ie, j, isd, ied, js, je+1, -1, kord_mt) - endif - -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Pressures_mapU-Out -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data u_=u -!$ser verbatim if (present(mfy)) then -!$ser data mfy_=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy_=cy -!$ser verbatim endif -!$ser verbatim endif - -! Note : This serialization portion will test the update of pe0 -! and pe3 -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim pe_3d(i-1,j,k) = pe(i-1,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Pressures_mapV-In -!$ser data pe_=pe_3d -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data ak=ak -!$ser data bk=bk -!$ser data v_=v -!$ser data kord_mt=kord_mt -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser verbatim endif - !------ - ! map v - !------ - if (j < je+1) then - do k=2,km+1 - do i=is,ie+1 - pe0(i,k) = 0.5*(pe(i-1,k, j)+pe(i,k, j)) - enddo - enddo - - do k=1,km+1 - bkh = 0.5*bk(k) - do i=is,ie+1 - pe3(i,k) = ak(k) + bkh*(pe(i-1,km+1,j)+pe(i,km+1,j)) - enddo - enddo - - call map1_ppm (km, pe0, v, gz, & - km, pe3, v, is, ie+1, & - j, isd, ied+1, jsd, jed, -1, kord_mt) - if (present(mfx)) then - ! print*, 'mfx present' - call map1_ppm (km, pe0, mfx, gz, & - km, pe3, mfx, is, ie+1, & - j, is, ie+1, js, je, -1, kord_mt) - endif - if (present(cx)) then - ! print*, 'cx present' - call map1_ppm (km, pe0, cx, gz, & - km, pe3, cx, is, ie+1, & - j, is, ie+1, jsd, jed, -1, kord_mt) - endif - endif ! (j < je+1) - -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie+1 -!$ser verbatim pe0_3d(i,j,k) = pe0(i,k) -!$ser verbatim pe3_3d(i,j,k) = pe3(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Pressures_mapV-Out -!$ser data pe0_=pe0_3d -!$ser data pe3_=pe3_3d -!$ser data v_=v -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser verbatim endif - -1000 continue - -! Update pressure variables and get new pkz, T_v, and omega - -!$OMP parallel do default(none) shared(is,ie,js,je,km,pe,ptop,kord_tm,remap_t, & -!$OMP remap_pt,remap_te,mfy,mfx,cx,cy,hydrostatic, & -!$OMP pt,pk,rg,peln,q,nwat,liq_wat,rainwat,ice_wat,snowwat, & -!$OMP graupel,sphum,cappa,r_vir,rcp,cp,k1k,delp, cld_amt, & -!$OMP delz,akap,pkz,te,u,v,ps, gridstruct, & -!$OMP ak,bk,nq,isd,ied,jsd,jed,kord_tr,fill, & -!$OMP hs,w,ws,kord_wz,do_omega,omga,rrg,kord_mt, & -#ifdef SERIALIZE -!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, js2d, & -#endif -!$OMP pe2_3d, pe_3d, peln_3d, pn2_3d, pk2_3d, pe0, pe0_3d, pe3_3d, & -!$OMP pe1_3d, pn1, dp2_3d, q_con, pe1, pn1_3d) & -!$OMP private(gz,cvm,kp,k_next,bkh,dp2, & -!$OMP pe2,pe3,pk2,pn2,phis,tpe,dlnp,tmp) - do 2000 j=js,je - -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint PE_pk_delp_peln-In -!$ser data pe2_=pe2_3d -!$ser data pe_=pe_3d -!$ser data peln_=peln_3d -!$ser data pn2_=pn2_3d -!$ser data pk2_=pk2_3d -!$ser data delp=delp -!$ser data pk=pk -!$ser data ak=ak -!$ser data bk=bk -!$ser data akap=akap -!$ser data ptop=ptop -!$ser verbatim endif - !---------- - ! Update pe - !---------- - do i=is,ie - pe2(i, 1) = ptop - pe2(i,km+1) = pe(i,km+1,j) - enddo - do k=2,km - do i=is,ie - pe2(i,k) = ak(k) + bk(k)*pe(i,km+1,j) - enddo - enddo - do k=1,km+1 - do i=is,ie - pe(i,k,j) = pe2(i,k) - enddo - enddo - - !---------- - ! Update pk - !---------- - do i=is,ie - pn2(i, 1) = peln(i, 1,j) - pn2(i,km+1) = peln(i,km+1,j) - pk2(i, 1) = pk(i,j, 1) - pk2(i,km+1) = pk(i,j,km+1) - enddo - do k=2,km - do i=is,ie - pn2(i,k) = log(pe2(i,k)) - pk2(i,k) = exp(akap*pn2(i,k)) - enddo - enddo - do k=1,km+1 - do i=is,ie - pk(i,j,k) = pk2(i,k) - enddo - enddo - - !------------ - ! update delp - !------------ - do k=1,km - do i=is,ie - delp(i,j,k) = pe2(i,k+1) - pe2(i,k) - enddo - enddo - - !------------ - ! update logP - !------------ - do k=1,km+1 - do i=is,ie - peln(i,k,j) = pn2(i,k) - enddo - enddo - -!$ser verbatim if(j == js2d) then -!$ser verbatim do k = 1, km+1 -!$ser verbatim do i = is, ie -!$ser verbatim pe2_3d(i,j,k) = pe2(i,k) -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim pn2_3d(i,j,k) = pn2(i,k) -!$ser verbatim pk2_3d(i,j,k) = pk2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint PE_pk_delp_peln-Out -!$ser data pe2_=pe2_3d -!$ser data pe_=pe_3d -!$ser data peln_=peln_3d -!$ser data pn2_=pn2_3d -!$ser data pk2_=pk2_3d -!$ser data delp=delp -!$ser data pk=pk -!$ser verbatim endif - - !--------------------- - ! Compute pkz and T_v - !--------------------- - if ( hydrostatic ) then - print*,'CODE NOT TESTED HERE 8' - do k=1,km - do i=is,ie - pkz(i,j,k) = (pk(i,j,k+1)-pk(i,j,k))/(akap*(peln(i,k+1,j)-peln(i,k,j))) - enddo - enddo - if (.not.remap_t) then - if (remap_te) then - ! Get updated T_v (store in pt) - do i=is,ie - gz(i) = hs(i,j) - enddo - do k=km,1,-1 - do i=is,ie - tpe = te(i,j,k) - gz(i) - 0.25*gridstruct%rsin2(i,j)*( & - u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j) ) - dlnp = rg*(peln(i,k+1,j) - peln(i,k,j)) - tmp = tpe / (cp - pe(i,k,j)*dlnp/delp(i,j,k)) - pt(i,j,k) = tmp - gz(i) = gz(i) + dlnp*tmp - enddo - enddo ! end k-loop - else - ! Make pt T_v - do k=1,km - do i=is,ie - pt(i,j,k) = pt(i,j,k)*pkz(i,j,k) - enddo - enddo - endif - endif - else - if (remap_te) then - print*,'CODE NOT TESTED HERE 9' - ! Invert TE using 3D winds to get pt (virtual temperature) and pkz: - do i=is,ie - phis(i,km+1) = hs(i,j) - enddo - do k=km,1,-1 - do i=is,ie - phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) - enddo - enddo - do k=1,km+1 - do i=is,ie - phis(i,k) = phis(i,k) * pe(i,k,j) - enddo - enddo - do k=1,km -#ifdef MOIST_CAPPA - call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, gz, cvm) - do i=is,ie - tpe = te(i,j,k) - & - ( 0.5*w(i,j,k)**2 + 0.25*gridstruct%rsin2(i,j)*( & - u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)) & - + (phis(i,k+1)-phis(i,k))/(pe(i,k+1,j)-pe(i,k,j)) ) - pt(i,j,k) = tpe/cvm(i) - cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) - pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - enddo -#else - do i=is,ie - tpe = te(i,j,k) - & - ( 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & - u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j))) ) - pt(i,j,k) = tpe/cv_air - pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - enddo -#endif - enddo - endif - if (remap_t) then - ! print*,'CODE EXECUTED' -! Note: pt at this stage is T_v or T_m - !$ser verbatim if(j == js2d) then - !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' - !$ser savepoint MoistCVPlusPkz_2d-In - !$ser data qvapor_js=q(:,j,:,sphum) - !$ser data qliquid_js=q(:,j,:,liq_wat) - !$ser data qice_js=q(:,j,:,ice_wat) - !$ser data qrain_js=q(:,j,:,rainwat) - !$ser data qsnow_js=q(:,j,:,snowwat) - !$ser data qgraupel_js=q(:,j,:,graupel) - !$ser data qcld_js=q(:,j,:,cld_amt) - !$ser data r_vir=r_vir - !$ser data cappa=cappa - !$ser data rrg=rrg - !$ser data delp=delp - !$ser data delz=delz - !$ser data pkz=pkz - !$ser data pt=pt - !$ser data k1k=k1k - !$ser verbatim endif - do k=1,km -#ifdef MOIST_CAPPA - call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, gz, cvm) - do i=is,ie - cappa(i,j,k) = rdgas / ( rdgas + cvm(i)/(1.+r_vir*q(i,j,k,sphum)) ) - pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - enddo -#else - do i=is,ie - pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) - enddo -#endif - enddo - !$ser verbatim if(j == js2d) then - !$ser savepoint MoistCVPlusPkz_2d-Out - !$ser data pkz=pkz - !$ser data cappa=cappa - !$ser verbatim endif - endif - if (remap_pt) then - print*,'CODE NOT TESTED HERE 10' -! Note: pt at this stage is Theta_v - do k=1,km - do i=is,ie - pkz(i,j,k) = exp( k1k*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k)) ) - ! Make pt T_v - pt(i,j,k) = pt(i,j,k)*pkz(i,j,k) - enddo - enddo - endif - endif - -! Interpolate omega/pe3 (defined at peln) to remapped cell center (dp2) - if ( do_omega ) then -! Copy omega field to pe3 - print*,'CODE NOT TESTED HERE 11' - do i=is,ie - pe3(i,1) = 0. - enddo - do k=2,km+1 - do i=is,ie - pe3(i,k) = omga(i,j,k-1) - enddo - enddo - do k=1,km - do i=is,ie - dp2(i,k) = 0.5*(peln(i,k,j) + peln(i,k+1,j)) - enddo - enddo - do i=is,ie - k_next = 1 - do n=1,km - kp = k_next - do k=kp,km - if( dp2(i,n) <= peln(i,k+1,j) .and. dp2(i,n) >= peln(i,k,j) ) then - omga(i,j,n) = pe3(i,k) + (pe3(i,k+1) - pe3(i,k)) * & - (dp2(i,n)-peln(i,k,j)) / (peln(i,k+1,j)-peln(i,k,j) ) - k_next = k - exit - endif - enddo - enddo - enddo - endif ! end do_omega - -2000 continue - -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js, je -!$ser verbatim do i = is, ie -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js-1, je+1 -!$ser verbatim do i = is-1, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe_=pe_3d -!$ser data qvapor=q(:,:,:,sphum) -!$ser data qliquid=q(:,:,:,liq_wat) -!$ser data qice=q(:,:,:,ice_wat) -!$ser data qrain=q(:,:,:,rainwat) -!$ser data qsnow=q(:,:,:,snowwat) -!$ser data qgraupel=q(:,:,:,graupel) -!$ser data qcld=q(:,:,:,cld_amt) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data peln_3d=peln_3d -!$ser data r_vir=r_vir -!$ser data cvm=cvm -!$ser data pk=pk -!$ser data pkz=pkz -!$ser data akap=akap -!$ser data t_min=qmin -!$ser data kord_wz=kord_wz -!$ser data ws_=ws -!$ser data w=w -!$ser data u=u -!$ser data v=v -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser data w_max=w_max_ -!$ser data w_min=w_min_ -!$ser data kord_mt=kord_mt -!$ser data cosa_s=cosa_s -!$ser data rsin2=rsin2 -!$ser data hs=hs -!$ser data grav=grav_ -!$ser data te_2d_=te_2d_f32 -!$ser data te0_2d_=te0_2d_f32 -!$ser data area_64_=area_64_ -!$ser data te=te - -! Do total energy conservation and fast saturation adjustment as requested -! and fill new PT (Theta_V) for next k_split step or export dry T - -!$OMP parallel default(none) shared(is,ie,js,je,km,kmp,ptop,u,v,pe,isd,ied,jsd,jed,kord_mt, & -!$OMP remap_t,remap_pt,remap_te, & -!$OMP te_2d,te,delp,hydrostatic,hs,rg,pt,peln, adiabatic, & -!$OMP cp,delz,nwat,rainwat,liq_wat,ice_wat,snowwat, & -!$OMP graupel,q_con,r_vir,sphum,w,pk,pkz,last_step,consv, & -!$OMP do_adiabatic_init,zsum1,zsum0,te0_2d,domain, & -!$OMP ng,gridstruct,E_Flux,pdt,dtmp,reproduce_sum,q, & -!$OMP mdt,cld_amt,cappa,dtdt,out_dt,rrg,akap,do_sat_adj, & -!$OMP fast_mp_consv,kord_tm, phis_3d, te_2d_f32, te0_2d_f32, zsum1_f32, & -#ifdef SERIALIZE -!$OMP ppser_savepoint, ppser_serializer, ppser_serializer_ref, ppser_zrperturb, serial_flag, & -#endif -!$OMP js2d, pe1_3d, pe2_3d, pn1, pn1_3d, pn2, pn2_3d, pk2, pk2_3d, pe_3d, peln_3d, dp2, dp2_3d, & -!$OMP pe0_3d, pe3_3d, ps, mfy, cy, mfx, cx) & -!$OMP private(pe0,pe1,pe2,pe3,cvm,gz,phis,tesum,zsum,dpln,dlnp,tmp, rsin2, cosa_s, grav_) - - dtmp = 0. - if( last_step .and. (.not.do_adiabatic_init) ) then - ! NOTE : Code can enter here since do_adiabatic_init can be False - if ( consv > consv_min ) then - ! print*, "consv > consv_min = entered", consv, consv_min -!$OMP do - do j=js,je - if ( hydrostatic ) then - print*,'CODE NOT TESTED HERE 12' - do i=is,ie - gz(i) = hs(i,j) - do k=1,km - gz(i) = gz(i) + rg*pt(i,j,k)*(peln(i,k+1,j)-peln(i,k,j)) - enddo - enddo - do i=is,ie - te_2d(i,j) = pe(i,km+1,j)*hs(i,j) - pe(i,1,j)*gz(i) - enddo - do k=1,km - do i=is,ie - te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cp*pt(i,j,k) + & - 0.25*gridstruct%rsin2(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & - v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j))) - enddo - enddo - else -! TE using 3D winds (pt is virtual temperature): - ! if(j==js) print*,"CODE EXECUTED 1" - !$ser verbatim if(j == js) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim phis_3d(i,j,k) = phis(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser verbatim rsin2=gridstruct%rsin2 - !$ser verbatim cosa_s=gridstruct%cosa_s - !$ser verbatim do i = is, ie - !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) - !$ser verbatim enddo - !$ser verbatim grav_=grav - !!$ser verbatim print*,'MoistCVPlusTe_2D serialization' - !$ser savepoint MoistCVPlusTe_2d-In - !$ser data qvapor_js=q(:,j,:,sphum) - !$ser data qliquid_js=q(:,j,:,liq_wat) - !$ser data qice_js=q(:,j,:,ice_wat) - !$ser data qrain_js=q(:,j,:,rainwat) - !$ser data qsnow_js=q(:,j,:,snowwat) - !$ser data qgraupel_js=q(:,j,:,graupel) - !$ser data delp=delp - !$ser data pt=pt - !$ser data phis_=phis_3d - !$ser data te_2d_=te_2d_f32 - !$ser data u=u - !$ser data v=v - !$ser data w=w - !$ser data cosa_s=cosa_s - !$ser data rsin2=rsin2 - !$ser data hs=hs - !$ser data grav=grav_ - !$ser data delz=delz - !$ser verbatim endif - do i=is,ie - te_2d(i,j) = 0. - phis(i,km+1) = hs(i,j) - enddo - do k=km,1,-1 - do i=is,ie - phis(i,k) = phis(i,k+1) - grav*delz(i,j,k) - enddo - enddo - do k=1,km -#ifdef MOIST_CAPPA - ! if(j==js) print*,"MOIST_CAPPA EXECUTED" - call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, gz, cvm) - do i=is,ie - te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cvm(i)*pt(i,j,k) + & - 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & - u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)))) - enddo -#else - do i=is,ie - te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cv_air*pt(i,j,k) + & - 0.5*(phis(i,k)+phis(i,k+1) + w(i,j,k)**2 + 0.5*gridstruct%rsin2(i,j)*( & - u(i,j,k)**2+u(i,j+1,k)**2 + v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*gridstruct%cosa_s(i,j)))) - enddo -#endif - enddo ! k-loop - !$ser verbatim if(j == js) then - !$ser verbatim do i = is, ie - !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) - !$ser verbatim enddo - !$ser savepoint MoistCVPlusTe_2d-Out - !$ser data te_2d_=te_2d_f32 - !$ser verbatim endif - - endif ! end non-hydro - - !$ser verbatim if(j == js2d) then - !$ser verbatim te_2d_f32 = real(te_2d) - !$ser verbatim te0_2d_f32 = real(te0_2d) - !$ser verbatim zsum1_f32 = real(zsum1) - !$ser savepoint Te_Zsum-In - !$ser data te_2d_=te_2d_f32 - !$ser data te0_2d_=te0_2d_f32 - !$ser data zsum1=zsum1_f32 - !$ser data delp=delp - !$ser data pkz=pkz - !$ser verbatim endif - do i=is,ie - te_2d(i,j) = te0_2d(i,j) - te_2d(i,j) - zsum1(i,j) = pkz(i,j,1)*delp(i,j,1) - enddo - do k=2,km - do i=is,ie - zsum1(i,j) = zsum1(i,j) + pkz(i,j,k)*delp(i,j,k) - enddo - enddo - if ( hydrostatic ) then - do i=is,ie - zsum0(i,j) = ptop*(pk(i,j,1)-pk(i,j,km+1)) + zsum1(i,j) - enddo - endif - !$ser verbatim if(j == js2d) then - !$ser verbatim te_2d_f32 = real(te_2d) - !$ser verbatim zsum1_f32 = real(zsum1) - !!$ser verbatim print*, 'sum te_2d: ', sum(te_2d_f32(:,1)), sum(te_2d(:,1)) - !$ser savepoint Te_Zsum-Out - !$ser data te_2d_=te_2d_f32 - !$ser data zsum1=zsum1_f32 - !$ser verbatim endif - - enddo ! j-loop - -!$OMP single - !print*,"MPP GLOBAL SUM CODE EXECUTED 1" - !$ser savepoint Mpp_global_sum-In - !$ser verbatim serial_flag=.true. - !$ser verbatim j=24 - !$ser data x_compute_size=j - !$ser data y_compute_size=j - !$ser verbatim j=1 - !$ser data x_compute_begin=j - !$ser data y_compute_begin=j - !$ser data max_ntile_pe=j - !$ser data tile=j - !$ser verbatim j=0 - !$ser data ioff=j - !$ser data joff=j - !$ser data serial_flag=serial_flag - !$ser data tesum=tesum - tesum = mpp_global_sum(domain, te_2d*gridstruct%area_64(is:ie,js:je), & - flags=BITWISE_EFP_SUM) - - !$ser verbatim te_2d=te_2d*gridstruct%area_64(is:ie,js:je) - !!$ser verbatim print*,'Sum of input into mpp_global_sum = ', sum(te_2d) - !$ser data inputArray=te_2d - - !$ser savepoint Mpp_global_sum-Out - !$ser data tesum=tesum - - !print*,'tesum = ', tesum - E_Flux = DBLE(consv)*tesum / DBLE(grav*pdt*4.*pi*radius**2) ! unit: W/m**2 - ! Note pdt is "phys" time step - if ( hydrostatic ) then - print*,'CODE NOT TESTED HERE 13' - zsum = mpp_global_sum(domain, zsum0*gridstruct%area_64(is:ie,js:je), & - flags=BITWISE_EFP_SUM) - dtmp = tesum / DBLE(cp*zsum) - else - zsum = mpp_global_sum(domain, zsum1*gridstruct%area_64(is:ie,js:je), & - flags=BITWISE_EFP_SUM) - dtmp = tesum / DBLE(cv_air*zsum) - endif -!$OMP end single - - elseif ( consv < -consv_min ) then - print*,'CODE NOT TESTED HERE 14' -!$OMP do - do j=js,je - do i=is,ie - zsum1(i,j) = pkz(i,j,1)*delp(i,j,1) - enddo - do k=2,km - do i=is,ie - zsum1(i,j) = zsum1(i,j) + pkz(i,j,k)*delp(i,j,k) - enddo - enddo - if ( hydrostatic ) then - do i=is,ie - zsum0(i,j) = ptop*(pk(i,j,1)-pk(i,j,km+1)) + zsum1(i,j) - enddo - endif - enddo - - E_Flux = consv -!$OMP single - if ( hydrostatic ) then - zsum = mpp_global_sum(domain, zsum0*gridstruct%area_64(is:ie,js:je), & - flags=BITWISE_EFP_SUM) - dtmp = E_Flux*(grav*pdt*4.*pi*radius**2) / (cp*zsum) - else - zsum = mpp_global_sum(domain, zsum1*gridstruct%area_64(is:ie,js:je), & - flags=BITWISE_EFP_SUM) - dtmp = E_Flux*(grav*pdt*4.*pi*radius**2) / (cv_air*zsum) - endif -!$OMP end single - endif ! end consv check - endif ! end last_step check - -! Note: pt at this stage is T_v - if ( do_sat_adj .and. nwat>=6 ) then - print*,'CODE NOT TESTED HERE 15' - call timing_on('sat_adj2') -!$OMP do - do k=kmp,km - do j=js,je - do i=is,ie - dpln(i,j) = peln(i,k+1,j) - peln(i,k,j) - enddo - enddo - call fv_sat_adj(abs(mdt), r_vir, is, ie, js, je, ng, hydrostatic, fast_mp_consv, & - te(isd,jsd,k), q(isd,jsd,k,sphum), q(isd,jsd,k,liq_wat), & - q(isd,jsd,k,ice_wat), q(isd,jsd,k,rainwat), & - q(isd,jsd,k,snowwat), q(isd,jsd,k,graupel), & - hs ,dpln, delz(isd:,jsd:,k), pt(isd,jsd,k), delp(isd,jsd,k), & - cappa(isd:,jsd:,k), gridstruct%area_64, dtdt(is:,js:,k), out_dt, last_step, q(isd,jsd,k,cld_amt)) - if ( .not. hydrostatic ) then - do j=js,je - do i=is,ie -#ifdef MOIST_CAPPA - pkz(i,j,k) = exp(cappa(i,j,k)*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) -#else - pkz(i,j,k) = exp(akap*log(rrg*delp(i,j,k)/delz(i,j,k)*pt(i,j,k))) -#endif - enddo - enddo - endif - enddo ! OpenMP k-loop - - if ( fast_mp_consv ) then -!$OMP do - do j=js,je - do i=is,ie - do k=kmp,km - te0_2d(i,j) = te0_2d(i,j) + te(i,j,k) - enddo - enddo - enddo - endif - call timing_off('sat_adj2') - endif ! do_sat_adj - - - if ( last_step .and. (.not. adiabatic) ) then - ! Output temperature if last_step - if ( .not. hydrostatic ) then - ! print*,'CODE EXECUTED' - !$ser savepoint MoistCVPlusPt_2d_last_step-In - !$ser data qvapor=q(:,:,:,sphum) - !$ser data qliquid=q(:,:,:,liq_wat) - !$ser data qice=q(:,:,:,ice_wat) - !$ser data qrain=q(:,:,:,rainwat) - !$ser data qsnow=q(:,:,:,snowwat) - !$ser data qgraupel=q(:,:,:,graupel) - !$ser data r_vir=r_vir - !$ser data dtmp=dtmp - !$ser data pt=pt - !$ser data pkz=pkz -!$OMP do - do k=1,km - do j=js,je -#ifdef USE_COND - call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, gz, cvm) - do i=is,ie - pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / ((1.+r_vir*q(i,j,k,sphum))*(1.-gz(i))) - enddo -#else - do i=is,ie - pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / (1.+r_vir*q(i,j,k,sphum)) - enddo -#endif - enddo ! j-loop - enddo ! k-loop - !$ser savepoint MoistCVPlusPt_2d_last_step-Out - !$ser data pt=pt -#ifdef USE_COND - ! print*, "USE_COND active" - !$ser savepoint Cond_output-In - !$ser data q_con=q_con - !$ser data qliquid=q(:,:,:,liq_wat) - !$ser data qice=q(:,:,:,ice_wat) - !$ser data qrain=q(:,:,:,rainwat) - !$ser data qsnow=q(:,:,:,snowwat) - !$ser data qgraupel=q(:,:,:,graupel) -! Fill condensate output -!$OMP do - do k=1,km - do j=js,je - do i=is,ie - q_con(i,j,k) = 0.0 - if (liq_wat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,liq_wat) - if (ice_wat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,ice_wat) - if (rainwat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,rainwat) - if (snowwat > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,snowwat) - if (graupel > 0) q_con(i,j,k) = q_con(i,j,k) + q(i,j,k,graupel) - enddo - enddo ! j-loop - enddo ! k-loop -#endif - !$ser savepoint Cond_output-Out - !$ser data q_con=q_con - - else - print*,'CODE NOT TESTED HERE 16' -!$OMP do - do k=1,km - do j=js,je - do i=is,ie - pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) / (1.+r_vir*q(i,j,k,sphum)) - enddo - enddo ! j-loop - enddo ! k-loop - endif - - elseif ( last_step .and. adiabatic ) then - print*,'CODE NOT TESTED HERE 17' -!$OMP do - do k=1,km - do j=js,je - do i=is,ie - pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) - enddo - enddo ! j-loop - enddo ! k-loop - - else - print*,'CODE NOT TESTED HERE 18' - ! Top of the loop expects PT to be Theta_V -!$OMP do - do k=1,km - do j=js,je - do i=is,ie - pt(i,j,k) = pt(i,j,k)/pkz(i,j,k) - enddo - enddo - enddo - - endif -!$OMP end parallel - -!$ser savepoint GetMPIProp-In -!$ser data delz=delz -!$ser savepoint GetMPIProp-Out -!$ser data delz=delz - end subroutine Lagrangian_to_Eulerian - -!>@brief The subroutine 'compute_total_energy' performs the FV3-consistent computation of the global total energy. -!>@details It includes the potential, internal (latent and sensible heat), kinetic terms. - subroutine compute_total_energy(is, ie, js, je, isd, ied, jsd, jed, km, & - u, v, w, delz, pt, delp, q, qc, pe, peln, hs, & - rsin2_l, cosa_s_l, & - r_vir, cp, rg, hlv, te_2d, ua, va, teq, & - moist_phys, nwat, sphum, liq_wat, rainwat, ice_wat, snowwat, graupel, hydrostatic, id_te) -!------------------------------------------------------ -! Compute vertically integrated total energy per column -!------------------------------------------------------ -! !INPUT PARAMETERS: - integer, intent(in):: km, is, ie, js, je, isd, ied, jsd, jed, id_te - integer, intent(in):: sphum, liq_wat, ice_wat, rainwat, snowwat, graupel, nwat - real, intent(inout), dimension(isd:ied,jsd:jed,km):: ua, va - real, intent(in), dimension(isd:ied,jsd:jed,km):: pt, delp - real, intent(in), dimension(isd:ied,jsd:jed,km,*):: q - real, intent(in), dimension(isd:ied,jsd:jed,km):: qc - real, intent(inout):: u(isd:ied, jsd:jed+1,km) - real, intent(inout):: v(isd:ied+1,jsd:jed, km) - real, intent(in):: w(isd:,jsd:,1:) !< vertical velocity (m/s) - real, intent(in):: delz(isd:,jsd:,1:) - real, intent(in):: hs(isd:ied,jsd:jed) !< surface geopotential - real, intent(in):: pe(is-1:ie+1,km+1,js-1:je+1) !< pressure at layer edges - real, intent(in):: peln(is:ie,km+1,js:je) !< log(pe) - real, intent(in):: cp, rg, r_vir, hlv - real, intent(in) :: rsin2_l(isd:ied, jsd:jed) - real, intent(in) :: cosa_s_l(isd:ied, jsd:jed) - logical, intent(in):: moist_phys, hydrostatic -!! Output: - real(kind=8), intent(out):: te_2d(is:ie,js:je) !< vertically integrated TE - real, intent(out):: teq(is:ie,js:je) !< Moist TE -!! Local - real, dimension(is:ie,km):: tv - real phiz(is:ie,km+1) - real cvm(is:ie), qd(is:ie) - integer i, j, k - -!---------------------- -! Output lat-lon winds: -!---------------------- -! call cubed_to_latlon(u, v, ua, va, dx, dy, rdxa, rdya, km, flagstruct%c2l_ord) - -!$OMP parallel do default(none) shared(is,ie,js,je,isd,ied,jsd,jed,km,hydrostatic,hs,pt,qc,rg,peln,te_2d, & -!$OMP pe,delp,cp,rsin2_l,u,v,cosa_s_l,delz,moist_phys,w, & -!$OMP q,nwat,liq_wat,rainwat,ice_wat,snowwat,graupel,sphum) & -!$OMP private(phiz, tv, cvm, qd) - do j=js,je - - if ( hydrostatic ) then - - do i=is,ie - phiz(i,km+1) = hs(i,j) - enddo - do k=km,1,-1 - do i=is,ie - tv(i,k) = pt(i,j,k)*(1.+qc(i,j,k)) - phiz(i,k) = phiz(i,k+1) + rg*tv(i,k)*(peln(i,k+1,j)-peln(i,k,j)) - enddo - enddo - - do i=is,ie - te_2d(i,j) = pe(i,km+1,j)*phiz(i,km+1) - pe(i,1,j)*phiz(i,1) - enddo - - do k=1,km - do i=is,ie - te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*(cp*tv(i,k) + & - 0.25*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & - v(i,j,k)**2+v(i+1,j,k)**2 - & - (u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j))) - enddo - enddo - - else -!----------------- -! Non-hydrostatic: -!----------------- - do i=is,ie - phiz(i,km+1) = hs(i,j) - do k=km,1,-1 - phiz(i,k) = phiz(i,k+1) - grav*delz(i,j,k) - enddo - enddo - do i=is,ie - te_2d(i,j) = 0. - enddo - if ( moist_phys ) then - do k=1,km -#ifdef MOIST_CAPPA - call moist_cv(is,ie,isd,ied,jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, qd, cvm) -#endif - do i=is,ie -#ifdef MOIST_CAPPA - te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cvm(i)*pt(i,j,k)*(1.+qc(i,j,k))*(1.-qd(i)) + & -#else - te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cv_air*pt(i,j,k)*(1.+qc(i,j,k)) + & -#endif - 0.5*(phiz(i,k)+phiz(i,k+1)+w(i,j,k)**2+0.5*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & - v(i,j,k)**2+v(i+1,j,k)**2-(u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j)))) - enddo - enddo - else - do k=1,km - do i=is,ie - te_2d(i,j) = te_2d(i,j) + delp(i,j,k)*( cv_air*pt(i,j,k) + & - 0.5*(phiz(i,k)+phiz(i,k+1)+w(i,j,k)**2+0.5*rsin2_l(i,j)*(u(i,j,k)**2+u(i,j+1,k)**2 + & - v(i,j,k)**2+v(i+1,j,k)**2-(u(i,j,k)+u(i,j+1,k))*(v(i,j,k)+v(i+1,j,k))*cosa_s_l(i,j)))) - enddo - enddo - endif - endif - enddo - -!------------------------------------- -! Diganostics computation for moist TE -!------------------------------------- - if( id_te>0 ) then -!$OMP parallel do default(none) shared(is,ie,js,je,teq,te_2d,moist_phys,km,hlv,sphum,q,delp) - do j=js,je - do i=is,ie - teq(i,j) = te_2d(i,j) - enddo - if ( moist_phys ) then - do k=1,km - do i=is,ie - teq(i,j) = teq(i,j) + hlv*q(i,j,k,sphum)*delp(i,j,k) - enddo - enddo - endif - enddo - endif - - end subroutine compute_total_energy - - - subroutine pkez(km, ifirst, ilast, jfirst, jlast, j, & - pe, pk, akap, peln, pkz, ptop) - -! INPUT PARAMETERS: - integer, intent(in):: km, j - integer, intent(in):: ifirst, ilast !< Latitude strip - integer, intent(in):: jfirst, jlast !< Latitude strip - real, intent(in):: akap - real, intent(in):: pe(ifirst-1:ilast+1,km+1,jfirst-1:jlast+1) - real, intent(in):: pk(ifirst:ilast,jfirst:jlast,km+1) - real, intent(IN):: ptop -! OUTPUT - real, intent(out):: pkz(ifirst:ilast,jfirst:jlast,km) - real, intent(inout):: peln(ifirst:ilast, km+1, jfirst:jlast) !< log (pe) -! Local - real pk2(ifirst:ilast, km+1) - real pek - real lnp - real ak1 - integer i, k - - ak1 = (akap + 1.) / akap - - pek = pk(ifirst,j,1) - do i=ifirst, ilast - pk2(i,1) = pek - enddo - - do k=2,km+1 - do i=ifirst, ilast -! peln(i,k,j) = log(pe(i,k,j)) - pk2(i,k) = pk(i,j,k) - enddo - enddo - -!---- GFDL modification - if( ptop < ptop_min ) then - do i=ifirst, ilast - peln(i,1,j) = peln(i,2,j) - ak1 - enddo - else - lnp = log( ptop ) - do i=ifirst, ilast - peln(i,1,j) = lnp - enddo - endif -!---- GFDL modification - - do k=1,km - do i=ifirst, ilast - pkz(i,j,k) = (pk2(i,k+1) - pk2(i,k) ) / & - (akap*(peln(i,k+1,j) - peln(i,k,j)) ) - enddo - enddo - - end subroutine pkez - - - - subroutine remap_z(km, pe1, q1, kn, pe2, q2, i1, i2, iv, kord) - -! INPUT PARAMETERS: - integer, intent(in) :: i1 !< Starting longitude - integer, intent(in) :: i2 !< Finishing longitude - integer, intent(in) :: kord !< Method order - integer, intent(in) :: km !< Original vertical dimension - integer, intent(in) :: kn !< Target vertical dimension - integer, intent(in) :: iv - - real, intent(in) :: pe1(i1:i2,km+1) !< height at layer edges from model top to bottom surface - real, intent(in) :: pe2(i1:i2,kn+1) !< height at layer edges from model top to bottom surface - real, intent(in) :: q1(i1:i2,km) !< Field input - -! INPUT/OUTPUT PARAMETERS: - real, intent(inout):: q2(i1:i2,kn) !< Field output - -! LOCAL VARIABLES: - real qs(i1:i2) - real dp1( i1:i2,km) - real q4(4,i1:i2,km) - real pl, pr, qsum, delp, esl - integer i, k, l, m, k0 - - do k=1,km - do i=i1,i2 - dp1(i,k) = pe1(i,k+1) - pe1(i,k) ! negative - q4(1,i,k) = q1(i,k) - enddo - enddo - -! Compute vertical subgrid distribution - if ( kord >7 ) then - call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) - else - call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) - endif - -! Mapping - do 3000 i=i1,i2 - k0 = 1 - do 555 k=1,kn - do 100 l=k0,km -! locate the top edge: pe2(i,k) - if(pe2(i,k) <= pe1(i,l) .and. pe2(i,k) >= pe1(i,l+1)) then - pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) - if(pe2(i,k+1) >= pe1(i,l+1)) then -! entire new grid is within the original grid - pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) - q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & - *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) - k0 = l - goto 555 - else -! Fractional area... - qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & - q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & - (r3*(1.+pl*(1.+pl)))) - do m=l+1,km -! locate the bottom edge: pe2(i,k+1) - if(pe2(i,k+1) < pe1(i,m+1) ) then -! Whole layer.. - qsum = qsum + dp1(i,m)*q4(1,i,m) - else - delp = pe2(i,k+1)-pe1(i,m) - esl = delp / dp1(i,m) - qsum = qsum + delp*(q4(2,i,m)+0.5*esl* & - (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) - k0 = m - goto 123 - endif - enddo - goto 123 - endif - endif -100 continue -123 q2(i,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) -555 continue -3000 continue - - end subroutine remap_z - - subroutine map_scalar( km, pe1, q1, qs, & - kn, pe2, q2, i1, i2, & - j, ibeg, iend, jbeg, jend, iv, kord, q_min) -! iv=1 - integer, intent(in) :: i1 !< Starting longitude - integer, intent(in) :: i2 !< Finishing longitude - integer, intent(in) :: iv !< Mode: 0 == constituents 1 == temp 2 == remap temp with cs scheme - integer, intent(in) :: kord !< Method order - integer, intent(in) :: j !< Current latitude - integer, intent(in) :: ibeg, iend, jbeg, jend - integer, intent(in) :: km !< Original vertical dimension - integer, intent(in) :: kn !< Target vertical dimension - real, intent(in) :: qs(i1:i2) !< bottom BC - real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate - real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate - real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input -! INPUT/OUTPUT PARAMETERS: - real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) !< Field output - real, intent(in):: q_min - -! DESCRIPTION: -! IV = 0: constituents -! pe1: pressure at layer edges (from model top to bottom surface) -! in the original vertical coordinate -! pe2: pressure at layer edges (from model top to bottom surface) -! in the new vertical coordinate -! LOCAL VARIABLES: - real dp1(i1:i2,km) - real q4(4,i1:i2,km) - real pl, pr, qsum, dp, esl - integer i, k, l, m, k0, jj - integer LM1,LP0,LP1 - -!$ser verbatim real q4_1_temp(i1:i2, i1:i2,km) -!$ser verbatim real q4_2_temp(i1:i2, i1:i2,km) -!$ser verbatim real q4_3_temp(i1:i2, i1:i2,km) -!$ser verbatim real q4_4_temp(i1:i2, i1:i2,km) -!$ser verbatim real dp1_temp(i1:i2, i1:i2,km) -!$ser verbatim real pe1_temp(i1:i2, i1:i2,km+1) -!$ser verbatim real pe2_temp(i1:i2, i1:i2,kn+1) - -!$ser verbatim real LM1_INDEX(i1:i2, i1:i2, km) -!$ser verbatim real LP0_INDEX(i1:i2, i1:i2, km) - - !$ser verbatim real :: qs_2d(i1:i2, i1:i2) - -!$ser verbatim q4_1_temp = 0.0 -!$ser verbatim q4_2_temp = 0.0 -!$ser verbatim q4_3_temp = 0.0 -!$ser verbatim q4_4_temp = 0.0 -!$ser verbatim dp1_temp = 0.0 -!$ser verbatim pe1_temp = 0.0 -!$ser verbatim pe2_temp = 0.0 -!$ser verbatim LM1_INDEX = 0 -!$ser verbatim LP0_INDEX = 0 -!$ser verbatim qs_2d = 0.0 - do k=1,km - do i=i1,i2 - dp1(i,k) = pe1(i,k+1) - pe1(i,k) - q4(1,i,k) = q1(i,j,k) - enddo - enddo - -!$ser verbatim if (j == 1) then -!$ser verbatim do k = 1,kn -!$ser verbatim do i = i1,i2 -!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) -!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) -!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) -!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) -!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do i = i1,i2 -!$ser verbatim qs_2d(i,j) = qs(i) -!$ser verbatim enddo -!$ser savepoint Scalar_Profile-In -!$ser data qs_=qs_2d -!$ser data q4_1=q4_1_temp -!$ser data q4_2=q4_2_temp -!$ser data q4_3=q4_3_temp -!$ser data q4_4=q4_4_temp -!$ser data dp1_=dp1_temp -!$ser data q_min=q_min -!$ser verbatim endif - -! Compute vertical subgrid distribution - if ( kord > 7 ) then - call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) - else - call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) - endif - -!$ser verbatim if (j == 1) then -!$ser verbatim do k = 1,kn -!$ser verbatim do i = i1,i2 -!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) -!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) -!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) -!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) -!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Scalar_Profile-Out -!$ser data q4_1=q4_1_temp -!$ser data q4_2=q4_2_temp -!$ser data q4_3=q4_3_temp -!$ser data q4_4=q4_4_temp -!$ser data dp1_=dp1_temp -!$ser verbatim endif - -! NOTE : q1 and q2 fields being passed into map_scalar are identical variables -! even though q1 is declared at INTENT(IN) and q2 is declared as INTENT(IN/OUT). -!$ser verbatim if(j == 1) then -!$ser verbatim do k = 1,kn -!$ser verbatim do i = i1,i2 -!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) -!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) -!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) -!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) -!$ser verbatim do jj = i1, i2 -!$ser verbatim dp1_temp(i,jj,k) = dp1(i,k) -!$ser verbatim pe1_temp(i,jj,k) = pe1(i,k) -!$ser verbatim pe2_temp(i,jj,k) = pe2(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do jj = i1, i2 -!$ser verbatim do i = i1, i2 -!$ser verbatim pe1_temp(i,jj,km+1) = pe1(i,km+1) -!$ser verbatim pe2_temp(i,jj,km+1) = pe2(i,km+1) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Lagrangian_Contribution_Interp-In -!$ser data q1=q1 pe1_=pe1_temp pe2_=pe2_temp q4_1=q4_1_temp q4_2=q4_2_temp -!$ser data q4_3=q4_3_temp q4_4=q4_4_temp dp1_=dp1_temp -!$ser verbatim endif - -! Interpolate field onto target Pressures -! --------------------------------------- - do i=i1,i2 - k0 = 1 - do 555 k=1,kn - LM1 = 1 - LP0 = 1 - do while( LP0.le.km ) - if (pe1(i,LP0).lt.pe2(i,k)) then - LP0 = LP0+1 - else - exit - endif - enddo - LM1 = max(LP0-1,1) - LP0 = min(LP0, km) -! Extrapolate Linearly above first model level -! ---------------------------------------------------- - if( LM1.eq.1 .and. LP0.eq.1 ) then - q2(i,j,k) = q1(i,j,1) + ( q1(i,j,2)-q1(i,j,1) )*( pe2(i,k)-pe1(i,1) ) & - /( pe1(i,2)-pe1(i,1) ) -! Extrapolate Linearly below last model level -! --------------------------------------------------- - else if( LM1.eq.km .and. LP0.eq.km ) then - q2(i,j,k) = q1(i,j,km) + ( q1(i,j,km)-q1(i,j,km-1) )*( pe2(i,k )-pe1(i,km ) ) & - /( pe1(i,km)-pe1(i,km-1) ) -! Interpolate Linearly between levels 1 => 2 and km-1 => km -! ----------------------------------------------------------------- - else if( LM1.eq.1 .or. LP0.eq.km ) then - q2(i,j,k) = q1(i,j,LP0) + ( q1(i,j,LM1)-q1(i,j,LP0) )*( pe2(i,k )-pe1(i,LP0) ) & - /( pe1(i,LM1)-pe1(i,LP0) ) - else - do l=k0,km -! locate the top edge: pe2(i,k) - if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then - pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) - if( pe2(i,k+1) <= pe1(i,l+1) ) then -! entire new grid is within the original grid - pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) - q2(i,j,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & - *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) - k0 = l - goto 555 - else -! Fractional area... - qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & - q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & - (r3*(1.+pl*(1.+pl)))) - do m=l+1,km -! locate the bottom edge: pe2(i,k+1) - if( pe2(i,k+1) > pe1(i,m+1) ) then -! Whole layer - qsum = qsum + dp1(i,m)*q4(1,i,m) - else - dp = pe2(i,k+1)-pe1(i,m) - esl = dp / dp1(i,m) - qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & - (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) - k0 = m - goto 123 - endif - enddo - goto 123 - endif - endif - enddo -123 q2(i,j,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) - - endif -555 continue - enddo - -!$ser verbatim if(j == 1) then -!$ser savepoint Lagrangian_Contribution_Interp-Out -!$ser data q1=q2 -!$ser verbatim endif - - end subroutine map_scalar - - - subroutine map1_ppm( km, pe1, q1, qs, & - kn, pe2, q2, i1, i2, & - j, ibeg, iend, jbeg, jend, iv, kord) - integer, intent(in) :: i1 !< Starting longitude - integer, intent(in) :: i2 !< Finishing longitude - integer, intent(in) :: iv !< Mode: 0 == constituents 1 == ??? 2 == remap temp with cs scheme - integer, intent(in) :: kord !< Method order - integer, intent(in) :: j !< Current latitude - integer, intent(in) :: ibeg, iend, jbeg, jend - integer, intent(in) :: km !< Original vertical dimension - integer, intent(in) :: kn !< Target vertical dimension - real, intent(in) :: qs(i1:i2) !< bottom BC - real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate - real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate - real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input -! INPUT/OUTPUT PARAMETERS: - real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) !< Field output - -! DESCRIPTION: -! IV = 0: constituents -! pe1: pressure at layer edges (from model top to bottom surface) -! in the original vertical coordinate -! pe2: pressure at layer edges (from model top to bottom surface) -! in the new vertical coordinate - -! LOCAL VARIABLES: - real dp1(i1:i2,km) - real q4(4,i1:i2,km) - real pl, pr, qsum, dp, esl - integer i, k, l, m, k0 - -!$ser verbatim real q4_1_temp(i1:i2, i1:i2,km) -!$ser verbatim real q4_2_temp(i1:i2, i1:i2,km) -!$ser verbatim real q4_3_temp(i1:i2, i1:i2,km) -!$ser verbatim real q4_4_temp(i1:i2, i1:i2,km) -!$ser verbatim real dp1_temp(i1:i2, i1:i2,km) -!$ser verbatim real :: qs_2d(i1:i2, i1:i2) - -!$ser verbatim q4_1_temp = 0.0 -!$ser verbatim q4_2_temp = 0.0 -!$ser verbatim q4_3_temp = 0.0 -!$ser verbatim q4_4_temp = 0.0 -!$ser verbatim qs_2d = 0.0 - - do k=1,km - do i=i1,i2 - dp1(i,k) = pe1(i,k+1) - pe1(i,k) - q4(1,i,k) = q1(i,j,k) - enddo - enddo - -! Compute vertical subgrid distribution - if ( kord >7 ) then -!$ser verbatim if (j == 1 .and. i2 == 24) then -!$ser verbatim do k = 1,kn -!$ser verbatim do i = i1,i2 -!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) -!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) -!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) -!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) -!$ser verbatim dp1_temp(i,j,k) = dp1(i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do i = i1,i2 -!$ser verbatim qs_2d(i,j) = qs(i) -!$ser verbatim enddo -!$ser savepoint CS_Profile-In -!$ser data qs_=qs_2d -!$ser data q4_1=q4_1_temp -!$ser data q4_2=q4_2_temp -!$ser data q4_3=q4_3_temp -!$ser data q4_4=q4_4_temp -!$ser data dp1_=dp1_temp -!$ser data kord_=kord -!$ser data iv_=iv -!$ser verbatim endif - call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) -!$ser verbatim if (j == 1 .and. i2 == 24) then -!$ser verbatim do k = 1,kn -!$ser verbatim do i = i1,i2 -!$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) -!$ser verbatim q4_2_temp(i,j,k) = q4(2,i,k) -!$ser verbatim q4_3_temp(i,j,k) = q4(3,i,k) -!$ser verbatim q4_4_temp(i,j,k) = q4(4,i,k) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint CS_Profile-Out -!$ser data q4_1=q4_1_temp -!$ser data q4_2=q4_2_temp -!$ser data q4_3=q4_3_temp -!$ser data q4_4=q4_4_temp -!$ser verbatim endif - else - call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) - endif - - do i=i1,i2 - k0 = 1 - do 555 k=1,kn - do l=k0,km -! locate the top edge: pe2(i,k) - if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then - pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) - if( pe2(i,k+1) <= pe1(i,l+1) ) then -! entire new grid is within the original grid - pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) - q2(i,j,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & - *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) - k0 = l - goto 555 - else -! Fractional area... - qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & - q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & - (r3*(1.+pl*(1.+pl)))) - do m=l+1,km -! locate the bottom edge: pe2(i,k+1) - if( pe2(i,k+1) > pe1(i,m+1) ) then -! Whole layer - qsum = qsum + dp1(i,m)*q4(1,i,m) - else - dp = pe2(i,k+1)-pe1(i,m) - esl = dp / dp1(i,m) - qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & - (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) - k0 = m - goto 123 - endif - enddo - goto 123 - endif - endif - enddo -123 q2(i,j,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) -555 continue - enddo - - end subroutine map1_ppm - - - subroutine mapn_tracer(nq, km, pe1, pe2, q1, dp2, kord, j, & - i1, i2, isd, ied, jsd, jed, q_min, fill) -! INPUT PARAMETERS: - integer, intent(in):: km !< vertical dimension - integer, intent(in):: j, nq, i1, i2 - integer, intent(in):: isd, ied, jsd, jed - integer, intent(in):: kord(nq) - real, intent(in):: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate - real, intent(in):: pe2(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate - real, intent(in):: dp2(i1:i2,km) - real, intent(in):: q_min - logical, intent(in):: fill - real, intent(inout):: q1(isd:ied,jsd:jed,km,nq) ! Field input -! LOCAL VARIABLES: - real:: q4(4,i1:i2,km,nq) - real:: q2(i1:i2,km,nq) !< Field output - real:: qsum(nq) - real:: dp1(i1:i2,km) - real:: qs(i1:i2) - real:: pl, pr, dp, esl, fac1, fac2 - integer:: i, k, l, m, k0, iq - - !$ser verbatim integer:: kord_iq, iv, im, js2d - !$ser verbatim js2d=jsd+3 - - do k=1,km - do i=i1,i2 - dp1(i,k) = pe1(i,k+1) - pe1(i,k) - enddo - enddo - - do iq=1,nq - do k=1,km - do i=i1,i2 - q4(1,i,k,iq) = q1(i,j,k,iq) - enddo - enddo - call scalar_profile( qs, q4(1,i1,1,iq), dp1, km, i1, i2, 0, kord(iq), q_min ) - enddo - -! Mapping - do 4000 i=i1,i2 - k0 = 1 - do 555 k=1,km - do 100 l=k0,km -! locate the top edge: pe2(i,k) - if(pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1)) then - pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) - if(pe2(i,k+1) <= pe1(i,l+1)) then -! entire new grid is within the original grid - pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) - fac1 = pr + pl - fac2 = r3*(pr*fac1 + pl*pl) - fac1 = 0.5*fac1 - do iq=1,nq - q2(i,k,iq) = q4(2,i,l,iq) + (q4(4,i,l,iq)+q4(3,i,l,iq)-q4(2,i,l,iq))*fac1 & - - q4(4,i,l,iq)*fac2 - enddo - k0 = l - goto 555 - else -! Fractional area... - dp = pe1(i,l+1) - pe2(i,k) - fac1 = 1. + pl - fac2 = r3*(1.+pl*fac1) - fac1 = 0.5*fac1 - do iq=1,nq - qsum(iq) = dp*(q4(2,i,l,iq) + (q4(4,i,l,iq)+ & - q4(3,i,l,iq) - q4(2,i,l,iq))*fac1 - q4(4,i,l,iq)*fac2) - enddo - do m=l+1,km -! locate the bottom edge: pe2(i,k+1) - if(pe2(i,k+1) > pe1(i,m+1) ) then - ! Whole layer.. - do iq=1,nq - qsum(iq) = qsum(iq) + dp1(i,m)*q4(1,i,m,iq) - enddo - else - dp = pe2(i,k+1)-pe1(i,m) - esl = dp / dp1(i,m) - fac1 = 0.5*esl - fac2 = 1.-r23*esl - do iq=1,nq - qsum(iq) = qsum(iq) + dp*( q4(2,i,m,iq) + fac1*( & - q4(3,i,m,iq)-q4(2,i,m,iq)+q4(4,i,m,iq)*fac2 ) ) - enddo - k0 = m - goto 123 - endif - enddo - goto 123 - endif - endif -100 continue -123 continue - do iq=1,nq - q2(i,k,iq) = qsum(iq) / dp2(i,k) - enddo -555 continue -4000 continue - - !$ser verbatim if(j == js2d ) then - !$ser verbatim im = i2-i1+1 - !$ser verbatim iv = 9 - !$ser savepoint Fillz-In - ! Note : Currently, serializing nq=nq and q2tracers=q2(:,:,1:nq) will not run the translate test - ! To successfully run the translate test, serialize nq=iv and q2tracers=q2(:,:,1:iv) - !$ser data im=im km=km nq=iv dp2=dp2 q2tracers=q2(:,:,1:iv) - !$ser verbatim endif - - if (fill) call fillz(i2-i1+1, km, nq, q2, dp2) - - !$ser verbatim if(j == js2d ) then - !$ser savepoint Fillz-Out - !$ser data q2tracers=q2(:,:,1:iv) - !$ser verbatim endif - - do iq=1,nq -! if (fill) call fillz(i2-i1+1, km, 1, q2(i1,1,iq), dp2) - do k=1,km - do i=i1,i2 - q1(i,j,k,iq) = q2(i,k,iq) - enddo - enddo - enddo - - end subroutine mapn_tracer - - - subroutine map1_q2(km, pe1, q1, & - kn, pe2, q2, dp2, & - i1, i2, iv, kord, j, & - ibeg, iend, jbeg, jend, q_min ) - - -! INPUT PARAMETERS: - integer, intent(in) :: j - integer, intent(in) :: i1, i2 - integer, intent(in) :: ibeg, iend, jbeg, jend - integer, intent(in) :: iv !< Mode: 0 == constituents 1 == ??? - integer, intent(in) :: kord - integer, intent(in) :: km !< Original vertical dimension - integer, intent(in) :: kn !< Target vertical dimension - - real, intent(in) :: pe1(i1:i2,km+1) !< pressure at layer edges from model top to bottom surface in the original vertical coordinate - real, intent(in) :: pe2(i1:i2,kn+1) !< pressure at layer edges from model top to bottom surface in the new vertical coordinate - real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) !< Field input - real, intent(in) :: dp2(i1:i2,kn) - real, intent(in) :: q_min -! INPUT/OUTPUT PARAMETERS: - real, intent(inout):: q2(i1:i2,kn) !< Field output -! LOCAL VARIABLES: - real qs(i1:i2) - real dp1(i1:i2,km) - real q4(4,i1:i2,km) - real pl, pr, qsum, dp, esl - - integer i, k, l, m, k0 - - do k=1,km - do i=i1,i2 - dp1(i,k) = pe1(i,k+1) - pe1(i,k) - q4(1,i,k) = q1(i,j,k) - enddo - enddo - -! Compute vertical subgrid distribution - if ( kord >7 ) then - call scalar_profile( qs, q4, dp1, km, i1, i2, iv, kord, q_min ) - else - call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) - endif - -! Mapping - do 5000 i=i1,i2 - k0 = 1 - do 555 k=1,kn - do 100 l=k0,km -! locate the top edge: pe2(i,k) - if(pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1)) then - pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) - if(pe2(i,k+1) <= pe1(i,l+1)) then -! entire new grid is within the original grid - pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) - q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & - *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) - k0 = l - goto 555 - else -! Fractional area... - qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & - q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & - (r3*(1.+pl*(1.+pl)))) - do m=l+1,km -! locate the bottom edge: pe2(i,k+1) - if(pe2(i,k+1) > pe1(i,m+1) ) then - ! Whole layer.. - qsum = qsum + dp1(i,m)*q4(1,i,m) - else - dp = pe2(i,k+1)-pe1(i,m) - esl = dp / dp1(i,m) - qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & - (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) - k0 = m - goto 123 - endif - enddo - goto 123 - endif - endif -100 continue -123 q2(i,k) = qsum / dp2(i,k) -555 continue -5000 continue - - end subroutine map1_q2 - - - - subroutine remap_2d(km, pe1, q1, & - kn, pe2, q2, & - i1, i2, iv, kord) - integer, intent(in):: i1, i2 - integer, intent(in):: iv !< Mode: 0 == constituents 1 ==others - integer, intent(in):: kord - integer, intent(in):: km !< Original vertical dimension - integer, intent(in):: kn !< Target vertical dimension - real, intent(in):: pe1(i1:i2,km+1) !< Pressure at layer edges from model top to bottom surface in the original vertical coordinate - real, intent(in):: pe2(i1:i2,kn+1) !< Pressure at layer edges from model top to bottom surface in the new vertical coordinate - real, intent(in) :: q1(i1:i2,km) !< Field input - real, intent(out):: q2(i1:i2,kn) !< Field output -! LOCAL VARIABLES: - real qs(i1:i2) - real dp1(i1:i2,km) - real q4(4,i1:i2,km) - real pl, pr, qsum, dp, esl - integer i, k, l, m, k0 - - do k=1,km - do i=i1,i2 - dp1(i,k) = pe1(i,k+1) - pe1(i,k) - q4(1,i,k) = q1(i,k) - enddo - enddo - -! Compute vertical subgrid distribution - if ( kord >7 ) then - call cs_profile( qs, q4, dp1, km, i1, i2, iv, kord ) - else - call ppm_profile( q4, dp1, km, i1, i2, iv, kord ) - endif - - do i=i1,i2 - k0 = 1 - do 555 k=1,kn -#ifdef OLD_TOP_EDGE - if( pe2(i,k+1) <= pe1(i,1) ) then -! Entire grid above old ptop - q2(i,k) = q4(2,i,1) - elseif( pe2(i,k) < pe1(i,1) .and. pe2(i,k+1)>pe1(i,1) ) then -! Partially above old ptop: - q2(i,k) = q1(i,1) -#else - if( pe2(i,k) <= pe1(i,1) ) then -! above old ptop: - q2(i,k) = q1(i,1) -#endif - else - do l=k0,km -! locate the top edge: pe2(i,k) - if( pe2(i,k) >= pe1(i,l) .and. pe2(i,k) <= pe1(i,l+1) ) then - pl = (pe2(i,k)-pe1(i,l)) / dp1(i,l) - if(pe2(i,k+1) <= pe1(i,l+1)) then -! entire new grid is within the original grid - pr = (pe2(i,k+1)-pe1(i,l)) / dp1(i,l) - q2(i,k) = q4(2,i,l) + 0.5*(q4(4,i,l)+q4(3,i,l)-q4(2,i,l)) & - *(pr+pl)-q4(4,i,l)*r3*(pr*(pr+pl)+pl**2) - k0 = l - goto 555 - else -! Fractional area... - qsum = (pe1(i,l+1)-pe2(i,k))*(q4(2,i,l)+0.5*(q4(4,i,l)+ & - q4(3,i,l)-q4(2,i,l))*(1.+pl)-q4(4,i,l)* & - (r3*(1.+pl*(1.+pl)))) - do m=l+1,km -! locate the bottom edge: pe2(i,k+1) - if(pe2(i,k+1) > pe1(i,m+1) ) then - ! Whole layer.. - qsum = qsum + dp1(i,m)*q4(1,i,m) - else - dp = pe2(i,k+1)-pe1(i,m) - esl = dp / dp1(i,m) - qsum = qsum + dp*(q4(2,i,m)+0.5*esl* & - (q4(3,i,m)-q4(2,i,m)+q4(4,i,m)*(1.-r23*esl))) - k0 = m - goto 123 - endif - enddo - goto 123 - endif - endif - enddo -123 q2(i,k) = qsum / ( pe2(i,k+1) - pe2(i,k) ) - endif -555 continue - enddo - - end subroutine remap_2d - - -!>@brief Optimized vertical profile reconstruction: -!> Latest: Apr 2008 S.-J. Lin, NOAA/GFDL - subroutine scalar_profile(qs, a4, delp, km, i1, i2, iv, kord, qmin) -! Optimized vertical profile reconstruction: -! Latest: Apr 2008 S.-J. Lin, NOAA/GFDL - integer, intent(in):: i1, i2 - integer, intent(in):: km !< vertical dimension - integer, intent(in):: iv !< iv =-1: winds iv = 0: positive definite scalars iv = 1: others - integer, intent(in):: kord - real, intent(in) :: qs(i1:i2) - real, intent(in) :: delp(i1:i2,km) !< Layer pressure thickness - real, intent(inout):: a4(4,i1:i2,km) !< Interpolated values - real, intent(in):: qmin -!----------------------------------------------------------------------- - logical, dimension(i1:i2,km):: extm, ext5, ext6 - real gam(i1:i2,km) - real q(i1:i2,km+1) - real d4(i1:i2) - real bet, a_bot, grat - real pmp_1, lac_1, pmp_2, lac_2, x0, x1 - integer i, k, im - - if ( iv .eq. -2 ) then - do i=i1,i2 - gam(i,2) = 0.5 - q(i,1) = 1.5*a4(1,i,1) - enddo - do k=2,km-1 - do i=i1, i2 - grat = delp(i,k-1) / delp(i,k) - bet = 2. + grat + grat - gam(i,k) - q(i,k) = (3.*(a4(1,i,k-1)+a4(1,i,k)) - q(i,k-1))/bet - gam(i,k+1) = grat / bet - enddo - enddo - do i=i1,i2 - grat = delp(i,km-1) / delp(i,km) - q(i,km) = (3.*(a4(1,i,km-1)+a4(1,i,km)) - grat*qs(i) - q(i,km-1)) / & - (2. + grat + grat - gam(i,km)) - q(i,km+1) = qs(i) - enddo - do k=km-1,1,-1 - do i=i1,i2 - q(i,k) = q(i,k) - gam(i,k+1)*q(i,k+1) - enddo - enddo - else - do i=i1,i2 - grat = delp(i,2) / delp(i,1) ! grid ratio - bet = grat*(grat+0.5) - q(i,1) = ( (grat+grat)*(grat+1.)*a4(1,i,1) + a4(1,i,2) ) / bet - gam(i,1) = ( 1. + grat*(grat+1.5) ) / bet - enddo - - do k=2,km - do i=i1,i2 - d4(i) = delp(i,k-1) / delp(i,k) - bet = 2. + d4(i) + d4(i) - gam(i,k-1) - q(i,k) = ( 3.*(a4(1,i,k-1)+d4(i)*a4(1,i,k)) - q(i,k-1) )/bet - gam(i,k) = d4(i) / bet - enddo - enddo - - do i=i1,i2 - a_bot = 1. + d4(i)*(d4(i)+1.5) - q(i,km+1) = (2.*d4(i)*(d4(i)+1.)*a4(1,i,km)+a4(1,i,km-1)-a_bot*q(i,km)) & - / ( d4(i)*(d4(i)+0.5) - a_bot*gam(i,km) ) - enddo - - do k=km,1,-1 - do i=i1,i2 - q(i,k) = q(i,k) - gam(i,k)*q(i,k+1) - enddo - enddo - endif - -!----- Perfectly linear scheme -------------------------------- - if ( abs(kord) > 16 ) then - do k=1,km - do i=i1,i2 - a4(2,i,k) = q(i,k ) - a4(3,i,k) = q(i,k+1) - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - enddo - return - endif -!----- Perfectly linear scheme -------------------------------- - -!------------------ -! Apply constraints -!------------------ - im = i2 - i1 + 1 - -! Apply *large-scale* constraints - do i=i1,i2 - q(i,2) = min( q(i,2), max(a4(1,i,1), a4(1,i,2)) ) - q(i,2) = max( q(i,2), min(a4(1,i,1), a4(1,i,2)) ) - enddo - - do k=2,km - do i=i1,i2 - gam(i,k) = a4(1,i,k) - a4(1,i,k-1) - enddo - enddo - -! Interior: - do k=3,km-1 - do i=i1,i2 - if ( gam(i,k-1)*gam(i,k+1)>0. ) then -! Apply large-scale constraint to ALL fields if not local max/min - q(i,k) = min( q(i,k), max(a4(1,i,k-1),a4(1,i,k)) ) - q(i,k) = max( q(i,k), min(a4(1,i,k-1),a4(1,i,k)) ) - else - if ( gam(i,k-1) > 0. ) then -! There exists a local max - q(i,k) = max(q(i,k), min(a4(1,i,k-1),a4(1,i,k))) - else -! There exists a local min - q(i,k) = min(q(i,k), max(a4(1,i,k-1),a4(1,i,k))) - if ( iv==0 ) q(i,k) = max(0., q(i,k)) - endif - endif - enddo - enddo - -! Bottom: - do i=i1,i2 - q(i,km) = min( q(i,km), max(a4(1,i,km-1), a4(1,i,km)) ) - q(i,km) = max( q(i,km), min(a4(1,i,km-1), a4(1,i,km)) ) - enddo - - do k=1,km - do i=i1,i2 - a4(2,i,k) = q(i,k ) - a4(3,i,k) = q(i,k+1) - enddo - enddo - - do k=1,km - if ( k==1 .or. k==km ) then - do i=i1,i2 - extm(i,k) = (a4(2,i,k)-a4(1,i,k)) * (a4(3,i,k)-a4(1,i,k)) > 0. - enddo - else - do i=i1,i2 - extm(i,k) = gam(i,k)*gam(i,k+1) < 0. - enddo - endif - if ( abs(kord) > 9 ) then - do i=i1,i2 - x0 = 2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k)) - x1 = abs(a4(2,i,k)-a4(3,i,k)) - a4(4,i,k) = 3.*x0 - ext5(i,k) = abs(x0) > x1 - ext6(i,k) = abs(a4(4,i,k)) > x1 - enddo - endif - enddo - -!--------------------------- -! Apply subgrid constraints: -!--------------------------- -! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) -! Top 2 and bottom 2 layers always use monotonic mapping - - if ( iv==0 ) then - do i=i1,i2 - a4(2,i,1) = max(0., a4(2,i,1)) - enddo - elseif ( iv==-1 ) then - do i=i1,i2 - if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. - enddo - elseif ( iv==2 ) then - do i=i1,i2 - a4(2,i,1) = a4(1,i,1) - a4(3,i,1) = a4(1,i,1) - a4(4,i,1) = 0. - enddo - endif - - if ( iv/=2 ) then - do i=i1,i2 - a4(4,i,1) = 3.*(2.*a4(1,i,1) - (a4(2,i,1)+a4(3,i,1))) - enddo - call cs_limiters(im, extm(i1,1), a4(1,i1,1), 1) - endif - -! k=2 - do i=i1,i2 - a4(4,i,2) = 3.*(2.*a4(1,i,2) - (a4(2,i,2)+a4(3,i,2))) - enddo - call cs_limiters(im, extm(i1,2), a4(1,i1,2), 2) - -!------------------------------------- -! Huynh's 2nd constraint for interior: -!------------------------------------- - do k=3,km-2 - if ( abs(kord)<9 ) then - do i=i1,i2 -! Left edges - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) -! Right edges - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - - elseif ( abs(kord)==9 ) then - do i=i1,i2 - if ( extm(i,k) .and. extm(i,k-1) ) then -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - a4(4,i,k) = 0. - else if ( extm(i,k) .and. extm(i,k+1) ) then -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - a4(4,i,k) = 0. - else if ( extm(i,k) .and. a4(1,i,k) abs(a4(2,i,k)-a4(3,i,k)) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - endif - endif - enddo - elseif ( abs(kord)==10 ) then - do i=i1,i2 - if( ext5(i,k) ) then - if( ext5(i,k-1) .or. ext5(i,k+1) ) then - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - endif - elseif( ext6(i,k) ) then - if( ext5(i,k-1) .or. ext5(i,k+1) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - endif - endif - enddo - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - elseif ( abs(kord)==12 ) then - do i=i1,i2 - if( extm(i,k) ) then - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - a4(4,i,k) = 0. - else ! not a local extremum - a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) -! Check within the smooth region if subgrid profile is non-monotonic - if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) - endif - endif - enddo - elseif ( abs(kord)==13 ) then - do i=i1,i2 - if( ext6(i,k) ) then - if ( ext6(i,k-1) .and. ext6(i,k+1) ) then -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - endif - endif - enddo - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - elseif ( abs(kord)==14 ) then - - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - - elseif ( abs(kord)==15 ) then ! Revised abs(kord)=9 scheme - do i=i1,i2 - if ( ext5(i,k) .and. ext5(i,k-1) ) then - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - else if ( ext5(i,k) .and. ext5(i,k+1) ) then - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - else if ( ext5(i,k) .and. a4(1,i,k)@brief The subroutine 'cs_profile' performs the optimized vertical profile reconstruction: -!>@date April 2008 -!>@author S. J. Lin, NOAA/GFDL - subroutine cs_profile(qs, a4, delp, km, i1, i2, iv, kord) -! Optimized vertical profile reconstruction: -! Latest: Apr 2008 S.-J. Lin, NOAA/GFDL - integer, intent(in):: i1, i2 - integer, intent(in):: km !< vertical dimension - integer, intent(in):: iv !< iv =-1: winds - !< iv = 0: positive definite scalars - !< iv = 1: others - integer, intent(in):: kord - real, intent(in) :: qs(i1:i2) - real, intent(in) :: delp(i1:i2,km) !< layer pressure thickness - real, intent(inout):: a4(4,i1:i2,km) !< Interpolated values -!----------------------------------------------------------------------- - logical, dimension(i1:i2,km):: extm, ext5, ext6 - real gam(i1:i2,km) - real q(i1:i2,km+1) - real d4(i1:i2) - real bet, a_bot, grat - real pmp_1, lac_1, pmp_2, lac_2, x0, x1 - integer i, k, im - - if ( iv .eq. -2 ) then - do i=i1,i2 - gam(i,2) = 0.5 - q(i,1) = 1.5*a4(1,i,1) - enddo - do k=2,km-1 - do i=i1, i2 - grat = delp(i,k-1) / delp(i,k) - bet = 2. + grat + grat - gam(i,k) - q(i,k) = (3.*(a4(1,i,k-1)+a4(1,i,k)) - q(i,k-1))/bet - gam(i,k+1) = grat / bet - enddo - enddo - do i=i1,i2 - grat = delp(i,km-1) / delp(i,km) - q(i,km) = (3.*(a4(1,i,km-1)+a4(1,i,km)) - grat*qs(i) - q(i,km-1)) / & - (2. + grat + grat - gam(i,km)) - q(i,km+1) = qs(i) - enddo - do k=km-1,1,-1 - do i=i1,i2 - q(i,k) = q(i,k) - gam(i,k+1)*q(i,k+1) - enddo - enddo - else - do i=i1,i2 - grat = delp(i,2) / delp(i,1) ! grid ratio - bet = grat*(grat+0.5) - q(i,1) = ( (grat+grat)*(grat+1.)*a4(1,i,1) + a4(1,i,2) ) / bet - gam(i,1) = ( 1. + grat*(grat+1.5) ) / bet - enddo - - do k=2,km - do i=i1,i2 - d4(i) = delp(i,k-1) / delp(i,k) - bet = 2. + d4(i) + d4(i) - gam(i,k-1) - q(i,k) = ( 3.*(a4(1,i,k-1)+d4(i)*a4(1,i,k)) - q(i,k-1) )/bet - gam(i,k) = d4(i) / bet - enddo - enddo - - do i=i1,i2 - a_bot = 1. + d4(i)*(d4(i)+1.5) - q(i,km+1) = (2.*d4(i)*(d4(i)+1.)*a4(1,i,km)+a4(1,i,km-1)-a_bot*q(i,km)) & - / ( d4(i)*(d4(i)+0.5) - a_bot*gam(i,km) ) - enddo - - do k=km,1,-1 - do i=i1,i2 - q(i,k) = q(i,k) - gam(i,k)*q(i,k+1) - enddo - enddo - endif -!----- Perfectly linear scheme -------------------------------- - if ( abs(kord) > 16 ) then - do k=1,km - do i=i1,i2 - a4(2,i,k) = q(i,k ) - a4(3,i,k) = q(i,k+1) - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - enddo - return - endif -!----- Perfectly linear scheme -------------------------------- - -!------------------ -! Apply constraints -!------------------ - im = i2 - i1 + 1 - -! Apply *large-scale* constraints - do i=i1,i2 - q(i,2) = min( q(i,2), max(a4(1,i,1), a4(1,i,2)) ) - q(i,2) = max( q(i,2), min(a4(1,i,1), a4(1,i,2)) ) - enddo - - do k=2,km - do i=i1,i2 - gam(i,k) = a4(1,i,k) - a4(1,i,k-1) - enddo - enddo - -! Interior: - do k=3,km-1 - do i=i1,i2 - if ( gam(i,k-1)*gam(i,k+1)>0. ) then -! Apply large-scale constraint to ALL fields if not local max/min - q(i,k) = min( q(i,k), max(a4(1,i,k-1),a4(1,i,k)) ) - q(i,k) = max( q(i,k), min(a4(1,i,k-1),a4(1,i,k)) ) - else - if ( gam(i,k-1) > 0. ) then -! There exists a local max - q(i,k) = max(q(i,k), min(a4(1,i,k-1),a4(1,i,k))) - else -! There exists a local min - q(i,k) = min(q(i,k), max(a4(1,i,k-1),a4(1,i,k))) - if ( iv==0 ) q(i,k) = max(0., q(i,k)) - endif - endif - enddo - enddo - -! Bottom: - do i=i1,i2 - q(i,km) = min( q(i,km), max(a4(1,i,km-1), a4(1,i,km)) ) - q(i,km) = max( q(i,km), min(a4(1,i,km-1), a4(1,i,km)) ) - enddo - - do k=1,km - do i=i1,i2 - a4(2,i,k) = q(i,k ) - a4(3,i,k) = q(i,k+1) - enddo - enddo - - do k=1,km - if ( k==1 .or. k==km ) then - do i=i1,i2 - extm(i,k) = (a4(2,i,k)-a4(1,i,k)) * (a4(3,i,k)-a4(1,i,k)) > 0. - enddo - else - do i=i1,i2 - extm(i,k) = gam(i,k)*gam(i,k+1) < 0. - enddo - endif - if ( abs(kord) > 9 ) then - do i=i1,i2 - x0 = 2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k)) - x1 = abs(a4(2,i,k)-a4(3,i,k)) - a4(4,i,k) = 3.*x0 - ext5(i,k) = abs(x0) > x1 - ext6(i,k) = abs(a4(4,i,k)) > x1 - enddo - endif - enddo - -!--------------------------- -! Apply subgrid constraints: -!--------------------------- -! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) -! Top 2 and bottom 2 layers always use monotonic mapping - - if ( iv==0 ) then - do i=i1,i2 - a4(2,i,1) = max(0., a4(2,i,1)) - enddo - elseif ( iv==-1 ) then - do i=i1,i2 - if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. - enddo - elseif ( iv==2 ) then - do i=i1,i2 - a4(2,i,1) = a4(1,i,1) - a4(3,i,1) = a4(1,i,1) - a4(4,i,1) = 0. - enddo - endif - - if ( iv/=2 ) then - do i=i1,i2 - a4(4,i,1) = 3.*(2.*a4(1,i,1) - (a4(2,i,1)+a4(3,i,1))) - enddo - call cs_limiters(im, extm(i1,1), a4(1,i1,1), 1) - endif - -! k=2 - do i=i1,i2 - a4(4,i,2) = 3.*(2.*a4(1,i,2) - (a4(2,i,2)+a4(3,i,2))) - enddo - call cs_limiters(im, extm(i1,2), a4(1,i1,2), 2) - -!------------------------------------- -! Huynh's 2nd constraint for interior: -!------------------------------------- - do k=3,km-2 - if ( abs(kord)<9 ) then - do i=i1,i2 -! Left edges - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) -! Right edges - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - - elseif ( abs(kord)==9 ) then - do i=i1,i2 - if ( extm(i,k) .and. extm(i,k-1) ) then ! c90_mp122 -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - a4(4,i,k) = 0. - else if ( extm(i,k) .and. extm(i,k+1) ) then ! c90_mp122 -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - a4(4,i,k) = 0. - else - a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) -! Check within the smooth region if subgrid profile is non-monotonic - if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) - endif - endif - enddo - elseif ( abs(kord)==10 ) then - do i=i1,i2 - if( ext5(i,k) ) then - if( ext5(i,k-1) .or. ext5(i,k+1) ) then - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - endif - elseif( ext6(i,k) ) then - if( ext5(i,k-1) .or. ext5(i,k+1) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - endif - endif - enddo - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - elseif ( abs(kord)==12 ) then - do i=i1,i2 - if( extm(i,k) ) then -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - a4(4,i,k) = 0. - else ! not a local extremum - a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) -! Check within the smooth region if subgrid profile is non-monotonic - if( abs(a4(4,i,k)) > abs(a4(2,i,k)-a4(3,i,k)) ) then - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - a4(4,i,k) = 6.*a4(1,i,k) - 3.*(a4(2,i,k)+a4(3,i,k)) - endif - endif - enddo - elseif ( abs(kord)==13 ) then - do i=i1,i2 - if( ext6(i,k) ) then - if ( ext6(i,k-1) .and. ext6(i,k+1) ) then -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - endif - endif - enddo - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - elseif ( abs(kord)==14 ) then - - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - - elseif ( abs(kord)==15 ) then ! revised kord=9 scehem - do i=i1,i2 - if ( ext5(i,k) ) then ! c90_mp122 - if ( ext5(i,k-1) .or. ext5(i,k+1) ) then ! c90_mp122 -! grid-scale 2-delta-z wave detected - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - endif - elseif( ext6(i,k) ) then -! Check within the smooth region if subgrid profile is non-monotonic - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - endif - enddo - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - elseif ( abs(kord)==16 ) then - do i=i1,i2 - if( ext5(i,k) ) then - if ( ext5(i,k-1) .or. ext5(i,k+1) ) then - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - elseif ( ext6(i,k-1) .or. ext6(i,k+1) ) then - ! Left edges - pmp_1 = a4(1,i,k) - 2.*gam(i,k+1) - lac_1 = pmp_1 + 1.5*gam(i,k+2) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), pmp_1, lac_1)), & - max(a4(1,i,k), pmp_1, lac_1) ) - ! Right edges - pmp_2 = a4(1,i,k) + 2.*gam(i,k) - lac_2 = pmp_2 - 1.5*gam(i,k-1) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), pmp_2, lac_2)), & - max(a4(1,i,k), pmp_2, lac_2) ) - endif - endif - enddo - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - else ! kord = 11 - do i=i1,i2 - if ( ext5(i,k) .and. (ext5(i,k-1) .or. ext5(i,k+1)) ) then -! Noisy region: - a4(2,i,k) = a4(1,i,k) - a4(3,i,k) = a4(1,i,k) - a4(4,i,k) = 0. - else - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - endif - enddo - endif - -! Additional constraint to ensure positivity - if ( iv==0 ) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 0) - - enddo ! k-loop - -!---------------------------------- -! Bottom layer subgrid constraints: -!---------------------------------- - if ( iv==0 ) then - do i=i1,i2 - a4(3,i,km) = max(0., a4(3,i,km)) - enddo - elseif ( iv .eq. -1 ) then - do i=i1,i2 - if ( a4(3,i,km)*a4(1,i,km) <= 0. ) a4(3,i,km) = 0. - enddo - endif - - do k=km-1,km - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - if(k==(km-1)) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 2) - if(k== km ) call cs_limiters(im, extm(i1,k), a4(1,i1,k), 1) - enddo - - end subroutine cs_profile - - - subroutine cs_limiters(im, extm, a4, iv) - integer, intent(in) :: im - integer, intent(in) :: iv - logical, intent(in) :: extm(im) - real , intent(inout) :: a4(4,im) !< PPM array -! LOCAL VARIABLES: - real da1, da2, a6da - integer i - - if ( iv==0 ) then -! Positive definite constraint - do i=1,im - if( a4(1,i)<=0.) then - a4(2,i) = a4(1,i) - a4(3,i) = a4(1,i) - a4(4,i) = 0. - else - if( abs(a4(3,i)-a4(2,i)) < -a4(4,i) ) then - if( (a4(1,i)+0.25*(a4(3,i)-a4(2,i))**2/a4(4,i)+a4(4,i)*r12) < 0. ) then -! local minimum is negative - if( a4(1,i) a4(2,i) ) then - a4(4,i) = 3.*(a4(2,i)-a4(1,i)) - a4(3,i) = a4(2,i) - a4(4,i) - else - a4(4,i) = 3.*(a4(3,i)-a4(1,i)) - a4(2,i) = a4(3,i) - a4(4,i) - endif - endif - endif - endif - enddo - elseif ( iv==1 ) then - do i=1,im - if( (a4(1,i)-a4(2,i))*(a4(1,i)-a4(3,i))>=0. ) then - a4(2,i) = a4(1,i) - a4(3,i) = a4(1,i) - a4(4,i) = 0. - else - da1 = a4(3,i) - a4(2,i) - da2 = da1**2 - a6da = a4(4,i)*da1 - if(a6da < -da2) then - a4(4,i) = 3.*(a4(2,i)-a4(1,i)) - a4(3,i) = a4(2,i) - a4(4,i) - elseif(a6da > da2) then - a4(4,i) = 3.*(a4(3,i)-a4(1,i)) - a4(2,i) = a4(3,i) - a4(4,i) - endif - endif - enddo - else -! Standard PPM constraint - do i=1,im - if( extm(i) ) then - a4(2,i) = a4(1,i) - a4(3,i) = a4(1,i) - a4(4,i) = 0. - else - da1 = a4(3,i) - a4(2,i) - da2 = da1**2 - a6da = a4(4,i)*da1 - if(a6da < -da2) then - a4(4,i) = 3.*(a4(2,i)-a4(1,i)) - a4(3,i) = a4(2,i) - a4(4,i) - elseif(a6da > da2) then - a4(4,i) = 3.*(a4(3,i)-a4(1,i)) - a4(2,i) = a4(3,i) - a4(4,i) - endif - endif - enddo - endif - end subroutine cs_limiters - - - - subroutine ppm_profile(a4, delp, km, i1, i2, iv, kord) - -! INPUT PARAMETERS: - integer, intent(in):: iv !< iv =-1: winds iv = 0: positive definite scalars iv = 1: others iv = 2: temp (if remap_t) and w (iv=-2) - integer, intent(in):: i1 !< Starting longitude - integer, intent(in):: i2 !< Finishing longitude - integer, intent(in):: km !< Vertical dimension - integer, intent(in):: kord !< Order (or more accurately method no.): - ! - real , intent(in):: delp(i1:i2,km) !< Layer pressure thickness - -! !INPUT/OUTPUT PARAMETERS: - real , intent(inout):: a4(4,i1:i2,km) !< Interpolated values - -! DESCRIPTION: -! -! Perform the piecewise parabolic reconstruction -! -! !REVISION HISTORY: -! S.-J. Lin revised at GFDL 2007 -!----------------------------------------------------------------------- -! local arrays: - real dc(i1:i2,km) - real h2(i1:i2,km) - real delq(i1:i2,km) - real df2(i1:i2,km) - real d4(i1:i2,km) - -! local scalars: - integer i, k, km1, lmt, it - real fac - real a1, a2, c1, c2, c3, d1, d2 - real qm, dq, lac, qmp, pmp - - km1 = km - 1 - it = i2 - i1 + 1 - - do k=2,km - do i=i1,i2 - delq(i,k-1) = a4(1,i,k) - a4(1,i,k-1) - d4(i,k ) = delp(i,k-1) + delp(i,k) - enddo - enddo - - do k=2,km1 - do i=i1,i2 - c1 = (delp(i,k-1)+0.5*delp(i,k))/d4(i,k+1) - c2 = (delp(i,k+1)+0.5*delp(i,k))/d4(i,k) - df2(i,k) = delp(i,k)*(c1*delq(i,k) + c2*delq(i,k-1)) / & - (d4(i,k)+delp(i,k+1)) - dc(i,k) = sign( min(abs(df2(i,k)), & - max(a4(1,i,k-1),a4(1,i,k),a4(1,i,k+1))-a4(1,i,k), & - a4(1,i,k)-min(a4(1,i,k-1),a4(1,i,k),a4(1,i,k+1))), df2(i,k) ) - enddo - enddo - -!----------------------------------------------------------- -! 4th order interpolation of the provisional cell edge value -!----------------------------------------------------------- - - do k=3,km1 - do i=i1,i2 - c1 = delq(i,k-1)*delp(i,k-1) / d4(i,k) - a1 = d4(i,k-1) / (d4(i,k) + delp(i,k-1)) - a2 = d4(i,k+1) / (d4(i,k) + delp(i,k)) - a4(2,i,k) = a4(1,i,k-1) + c1 + 2./(d4(i,k-1)+d4(i,k+1)) * & - ( delp(i,k)*(c1*(a1 - a2)+a2*dc(i,k-1)) - & - delp(i,k-1)*a1*dc(i,k ) ) - enddo - enddo - -! if(km>8 .and. kord>4) call steepz(i1, i2, km, a4, df2, dc, delq, delp, d4) - -! Area preserving cubic with 2nd deriv. = 0 at the boundaries -! Top - do i=i1,i2 - d1 = delp(i,1) - d2 = delp(i,2) - qm = (d2*a4(1,i,1)+d1*a4(1,i,2)) / (d1+d2) - dq = 2.*(a4(1,i,2)-a4(1,i,1)) / (d1+d2) - c1 = 4.*(a4(2,i,3)-qm-d2*dq) / ( d2*(2.*d2*d2+d1*(d2+3.*d1)) ) - c3 = dq - 0.5*c1*(d2*(5.*d1+d2)-3.*d1*d1) - a4(2,i,2) = qm - 0.25*c1*d1*d2*(d2+3.*d1) -! Top edge: -!------------------------------------------------------- - a4(2,i,1) = d1*(2.*c1*d1**2-c3) + a4(2,i,2) -!------------------------------------------------------- -! a4(2,i,1) = (12./7.)*a4(1,i,1)-(13./14.)*a4(1,i,2)+(3./14.)*a4(1,i,3) -!------------------------------------------------------- -! No over- and undershoot condition - a4(2,i,2) = max( a4(2,i,2), min(a4(1,i,1), a4(1,i,2)) ) - a4(2,i,2) = min( a4(2,i,2), max(a4(1,i,1), a4(1,i,2)) ) - dc(i,1) = 0.5*(a4(2,i,2) - a4(1,i,1)) - enddo - -! Enforce monotonicity within the top layer - - if( iv==0 ) then - do i=i1,i2 - a4(2,i,1) = max(0., a4(2,i,1)) - a4(2,i,2) = max(0., a4(2,i,2)) - enddo - elseif( iv==-1 ) then - do i=i1,i2 - if ( a4(2,i,1)*a4(1,i,1) <= 0. ) a4(2,i,1) = 0. - enddo - elseif( abs(iv)==2 ) then - do i=i1,i2 - a4(2,i,1) = a4(1,i,1) - a4(3,i,1) = a4(1,i,1) - enddo - endif - -! Bottom -! Area preserving cubic with 2nd deriv. = 0 at the surface - do i=i1,i2 - d1 = delp(i,km) - d2 = delp(i,km1) - qm = (d2*a4(1,i,km)+d1*a4(1,i,km1)) / (d1+d2) - dq = 2.*(a4(1,i,km1)-a4(1,i,km)) / (d1+d2) - c1 = (a4(2,i,km1)-qm-d2*dq) / (d2*(2.*d2*d2+d1*(d2+3.*d1))) - c3 = dq - 2.0*c1*(d2*(5.*d1+d2)-3.*d1*d1) - a4(2,i,km) = qm - c1*d1*d2*(d2+3.*d1) -! Bottom edge: -!----------------------------------------------------- - a4(3,i,km) = d1*(8.*c1*d1**2-c3) + a4(2,i,km) -! dc(i,km) = 0.5*(a4(3,i,km) - a4(1,i,km)) -!----------------------------------------------------- -! a4(3,i,km) = (12./7.)*a4(1,i,km)-(13./14.)*a4(1,i,km-1)+(3./14.)*a4(1,i,km-2) -! No over- and under-shoot condition - a4(2,i,km) = max( a4(2,i,km), min(a4(1,i,km), a4(1,i,km1)) ) - a4(2,i,km) = min( a4(2,i,km), max(a4(1,i,km), a4(1,i,km1)) ) - dc(i,km) = 0.5*(a4(1,i,km) - a4(2,i,km)) - enddo - - -! Enforce constraint on the "slope" at the surface - -#ifdef BOT_MONO - do i=i1,i2 - a4(4,i,km) = 0 - if( a4(3,i,km) * a4(1,i,km) <= 0. ) a4(3,i,km) = 0. - d1 = a4(1,i,km) - a4(2,i,km) - d2 = a4(3,i,km) - a4(1,i,km) - if ( d1*d2 < 0. ) then - a4(2,i,km) = a4(1,i,km) - a4(3,i,km) = a4(1,i,km) - else - dq = sign(min(abs(d1),abs(d2),0.5*abs(delq(i,km-1))), d1) - a4(2,i,km) = a4(1,i,km) - dq - a4(3,i,km) = a4(1,i,km) + dq - endif - enddo -#else - if( iv==0 ) then - do i=i1,i2 - a4(2,i,km) = max(0.,a4(2,i,km)) - a4(3,i,km) = max(0.,a4(3,i,km)) - enddo - elseif( iv<0 ) then - do i=i1,i2 - if( a4(1,i,km)*a4(3,i,km) <= 0. ) a4(3,i,km) = 0. - enddo - endif -#endif - - do k=1,km1 - do i=i1,i2 - a4(3,i,k) = a4(2,i,k+1) - enddo - enddo - -!----------------------------------------------------------- -! f(s) = AL + s*[(AR-AL) + A6*(1-s)] ( 0 <= s <= 1 ) -!----------------------------------------------------------- -! Top 2 and bottom 2 layers always use monotonic mapping - do k=1,2 - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 0) - enddo - - if(kord >= 7) then -!----------------------- -! Huynh's 2nd constraint -!----------------------- - do k=2,km1 - do i=i1,i2 -! Method#1 -! h2(i,k) = delq(i,k) - delq(i,k-1) -! Method#2 - better - h2(i,k) = 2.*(dc(i,k+1)/delp(i,k+1) - dc(i,k-1)/delp(i,k-1)) & - / ( delp(i,k)+0.5*(delp(i,k-1)+delp(i,k+1)) ) & - * delp(i,k)**2 -! Method#3 -!!! h2(i,k) = dc(i,k+1) - dc(i,k-1) - enddo - enddo - - fac = 1.5 ! original quasi-monotone - - do k=3,km-2 - do i=i1,i2 -! Right edges -! qmp = a4(1,i,k) + 2.0*delq(i,k-1) -! lac = a4(1,i,k) + fac*h2(i,k-1) + 0.5*delq(i,k-1) -! - pmp = 2.*dc(i,k) - qmp = a4(1,i,k) + pmp - lac = a4(1,i,k) + fac*h2(i,k-1) + dc(i,k) - a4(3,i,k) = min(max(a4(3,i,k), min(a4(1,i,k), qmp, lac)), & - max(a4(1,i,k), qmp, lac) ) -! Left edges -! qmp = a4(1,i,k) - 2.0*delq(i,k) -! lac = a4(1,i,k) + fac*h2(i,k+1) - 0.5*delq(i,k) -! - qmp = a4(1,i,k) - pmp - lac = a4(1,i,k) + fac*h2(i,k+1) - dc(i,k) - a4(2,i,k) = min(max(a4(2,i,k), min(a4(1,i,k), qmp, lac)), & - max(a4(1,i,k), qmp, lac)) -!------------- -! Recompute A6 -!------------- - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo -! Additional constraint to ensure positivity when kord=7 - if (iv == 0 .and. kord >= 6 ) & - call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 2) - enddo - - else - - lmt = kord - 3 - lmt = max(0, lmt) - if (iv == 0) lmt = min(2, lmt) - - do k=3,km-2 - if( kord /= 4) then - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - endif - if(kord/=6) call ppm_limiters(dc(i1,k), a4(1,i1,k), it, lmt) - enddo - endif - - do k=km1,km - do i=i1,i2 - a4(4,i,k) = 3.*(2.*a4(1,i,k) - (a4(2,i,k)+a4(3,i,k))) - enddo - call ppm_limiters(dc(i1,k), a4(1,i1,k), it, 0) - enddo - - end subroutine ppm_profile - - - subroutine ppm_limiters(dm, a4, itot, lmt) - -! INPUT PARAMETERS: - real , intent(in):: dm(*) !< Linear slope - integer, intent(in) :: itot !< Total Longitudes - integer, intent(in) :: lmt !< 0: Standard PPM constraint 1: Improved full monotonicity constraint - !< (Lin) 2: Positive definite constraint - !< 3: do nothing (return immediately) -! INPUT/OUTPUT PARAMETERS: - real , intent(inout) :: a4(4,*) !< PPM array AA <-- a4(1,i) AL <-- a4(2,i) AR <-- a4(3,i) A6 <-- a4(4,i) -! LOCAL VARIABLES: - real qmp - real da1, da2, a6da - real fmin - integer i - -! Developer: S.-J. Lin - - if ( lmt == 3 ) return - - if(lmt == 0) then -! Standard PPM constraint - do i=1,itot - if(dm(i) == 0.) then - a4(2,i) = a4(1,i) - a4(3,i) = a4(1,i) - a4(4,i) = 0. - else - da1 = a4(3,i) - a4(2,i) - da2 = da1**2 - a6da = a4(4,i)*da1 - if(a6da < -da2) then - a4(4,i) = 3.*(a4(2,i)-a4(1,i)) - a4(3,i) = a4(2,i) - a4(4,i) - elseif(a6da > da2) then - a4(4,i) = 3.*(a4(3,i)-a4(1,i)) - a4(2,i) = a4(3,i) - a4(4,i) - endif - endif - enddo - - elseif (lmt == 1) then - -! Improved full monotonicity constraint (Lin 2004) -! Note: no need to provide first guess of A6 <-- a4(4,i) - do i=1, itot - qmp = 2.*dm(i) - a4(2,i) = a4(1,i)-sign(min(abs(qmp),abs(a4(2,i)-a4(1,i))), qmp) - a4(3,i) = a4(1,i)+sign(min(abs(qmp),abs(a4(3,i)-a4(1,i))), qmp) - a4(4,i) = 3.*( 2.*a4(1,i) - (a4(2,i)+a4(3,i)) ) - enddo - - elseif (lmt == 2) then - -! Positive definite constraint - do i=1,itot - if( abs(a4(3,i)-a4(2,i)) < -a4(4,i) ) then - fmin = a4(1,i)+0.25*(a4(3,i)-a4(2,i))**2/a4(4,i)+a4(4,i)*r12 - if( fmin < 0. ) then - if(a4(1,i) a4(2,i)) then - a4(4,i) = 3.*(a4(2,i)-a4(1,i)) - a4(3,i) = a4(2,i) - a4(4,i) - else - a4(4,i) = 3.*(a4(3,i)-a4(1,i)) - a4(2,i) = a4(3,i) - a4(4,i) - endif - endif - endif - enddo - - endif - - end subroutine ppm_limiters - - - - subroutine steepz(i1, i2, km, a4, df2, dm, dq, dp, d4) - integer, intent(in) :: km, i1, i2 - real , intent(in) :: dp(i1:i2,km) !< Grid size - real , intent(in) :: dq(i1:i2,km) !< Backward diff of q - real , intent(in) :: d4(i1:i2,km) !< Backward sum: dp(k)+ dp(k-1) - real , intent(in) :: df2(i1:i2,km) !< First guess mismatch - real , intent(in) :: dm(i1:i2,km) !< Monotonic mismatch -! INPUT/OUTPUT PARAMETERS: - real , intent(inout) :: a4(4,i1:i2,km) !@brief The subroutine 'rst_remap' remaps all variables required for a restart. -!>@details npz_restart /= npz (i.e., when the number of vertical levels is -!! changed at restart) - subroutine rst_remap(km, kn, is,ie,js,je, isd,ied,jsd,jed, nq, ntp, & - delp_r, u_r, v_r, w_r, delz_r, pt_r, q_r, qdiag_r, & - delp, u, v, w, delz, pt, q, qdiag, & - ak_r, bk_r, ptop, ak, bk, hydrostatic, make_nh, & - domain, square_domain) -!------------------------------------ -! Assuming hybrid sigma-P coordinate: -!------------------------------------ -! INPUT PARAMETERS: - integer, intent(in):: km !< Restart z-dimension - integer, intent(in):: kn !< Run time dimension - integer, intent(in):: nq, ntp !< Number of tracers (including H2O) - integer, intent(in):: is,ie,isd,ied !< Starting & ending X-Dir index - integer, intent(in):: js,je,jsd,jed !< Starting & ending Y-Dir index - logical, intent(in):: hydrostatic, make_nh, square_domain - real, intent(IN) :: ptop - real, intent(in) :: ak_r(km+1) - real, intent(in) :: bk_r(km+1) - real, intent(in) :: ak(kn+1) - real, intent(in) :: bk(kn+1) - real, intent(in):: delp_r(is:ie,js:je,km) !< Pressure thickness - real, intent(in):: u_r(is:ie, js:je+1,km) !< u-wind (m/s) - real, intent(in):: v_r(is:ie+1,js:je ,km) !< v-wind (m/s) - real, intent(inout):: pt_r(is:ie,js:je,km) - real, intent(in):: w_r(is:ie,js:je,km) - real, intent(in):: q_r(is:ie,js:je,km,1:ntp) - real, intent(in):: qdiag_r(is:ie,js:je,km,ntp+1:nq) - real, intent(inout)::delz_r(is:ie,js:je,km) - type(domain2d), intent(INOUT) :: domain -! Output: - real, intent(out):: delp(isd:ied,jsd:jed,kn) !< Pressure thickness - real, intent(out):: u(isd:ied ,jsd:jed+1,kn) !< u-wind (m/s) - real, intent(out):: v(isd:ied+1,jsd:jed ,kn) !< v-wind (m/s) - real, intent(out):: w(isd: ,jsd: ,1:) !< Vertical velocity (m/s) - real, intent(out):: pt(isd:ied ,jsd:jed ,kn) !< Temperature - real, intent(out):: q(isd:ied,jsd:jed,kn,1:ntp) - real, intent(out):: qdiag(isd:ied,jsd:jed,kn,ntp+1:nq) - real, intent(out):: delz(isd:,jsd:,1:) !< Delta-height (m) -!----------------------------------------------------------------------- - real r_vir, rgrav - real ps(isd:ied,jsd:jed) !< Surface pressure - real pe1(is:ie,km+1) - real pe2(is:ie,kn+1) - real pv1(is:ie+1,km+1) - real pv2(is:ie+1,kn+1) - - integer i,j,k , iq - integer, parameter:: kord=4 - -#ifdef HYDRO_DELZ_REMAP - if (is_master() .and. .not. hydrostatic) then - print*, '' - print*, ' REMAPPING IC: INITIALIZING DELZ WITH HYDROSTATIC STATE ' - print*, '' - endif -#endif - -#ifdef HYDRO_DELZ_EXTRAP - if (is_master() .and. .not. hydrostatic) then - print*, '' - print*, ' REMAPPING IC: INITIALIZING DELZ WITH HYDROSTATIC STATE ABOVE INPUT MODEL TOP ' - print*, '' - endif -#endif - -#ifdef ZERO_W_EXTRAP - if (is_master() .and. .not. hydrostatic) then - print*, '' - print*, ' REMAPPING IC: INITIALIZING W TO ZERO ABOVE INPUT MODEL TOP ' - print*, '' - endif -#endif - - r_vir = rvgas/rdgas - 1. - rgrav = 1./grav - -!$OMP parallel do default(none) shared(is,ie,js,je,ps,ak_r) - do j=js,je - do i=is,ie - ps(i,j) = ak_r(1) - enddo - enddo - -! this OpenMP do-loop setup cannot work in it's current form.... -!$OMP parallel do default(none) shared(is,ie,js,je,km,ps,delp_r) - do j=js,je - do k=1,km - do i=is,ie - ps(i,j) = ps(i,j) + delp_r(i,j,k) - enddo - enddo - enddo - -! only one cell is needed - if ( square_domain ) then - call mpp_update_domains(ps, domain, whalo=1, ehalo=1, shalo=1, nhalo=1, complete=.true.) - else - call mpp_update_domains(ps, domain, complete=.true.) - endif - -! Compute virtual Temp -!$OMP parallel do default(none) shared(is,ie,js,je,km,pt_r,r_vir,q_r) - do k=1,km - do j=js,je - do i=is,ie - pt_r(i,j,k) = pt_r(i,j,k) * (1.+r_vir*q_r(i,j,k,1)) - enddo - enddo - enddo - -!$OMP parallel do default(none) shared(is,ie,js,je,km,ak_r,bk_r,ps,kn,ak,bk,u_r,u,delp, & -!$OMP ntp,nq,hydrostatic,make_nh,w_r,w,delz_r,delp_r,delz, & -!$OMP pt_r,pt,v_r,v,q,q_r,qdiag,qdiag_r) & -!$OMP private(pe1, pe2, pv1, pv2) - do 6000 j=js,je+1 -!------ -! map u -!------ - do k=1,km+1 - do i=is,ie - pe1(i,k) = ak_r(k) + 0.5*bk_r(k)*(ps(i,j-1)+ps(i,j)) - enddo - enddo - - do k=1,kn+1 - do i=is,ie - pe2(i,k) = ak(k) + 0.5*bk(k)*(ps(i,j-1)+ps(i,j)) - enddo - enddo - - call remap_2d(km, pe1, u_r(is:ie,j:j,1:km), & - kn, pe2, u(is:ie,j:j,1:kn), & - is, ie, -1, kord) - - if ( j /= (je+1) ) then - -!--------------- -! Hybrid sigma-p -!--------------- - do k=1,km+1 - do i=is,ie - pe1(i,k) = ak_r(k) + bk_r(k)*ps(i,j) - enddo - enddo - - do k=1,kn+1 - do i=is,ie - pe2(i,k) = ak(k) + bk(k)*ps(i,j) - enddo - enddo - -!------------- -! Compute delp -!------------- - do k=1,kn - do i=is,ie - delp(i,j,k) = pe2(i,k+1) - pe2(i,k) - enddo - enddo - -!---------------- -! Map constituents -!---------------- - if( nq /= 0 ) then - do iq=1,ntp - call remap_2d(km, pe1, q_r(is:ie,j:j,1:km,iq:iq), & - kn, pe2, q(is:ie,j:j,1:kn,iq:iq), & - is, ie, 0, kord) - enddo - do iq=ntp+1,nq - call remap_2d(km, pe1, qdiag_r(is:ie,j:j,1:km,iq:iq), & - kn, pe2, qdiag(is:ie,j:j,1:kn,iq:iq), & - is, ie, 0, kord) - enddo - endif - - if ( .not. hydrostatic .and. .not. make_nh) then -! Remap vertical wind: - call remap_2d(km, pe1, w_r(is:ie,j:j,1:km), & - kn, pe2, w(is:ie,j:j,1:kn), & - is, ie, -1, kord) - -#ifdef ZERO_W_EXTRAP - do k=1,kn - do i=is,ie - if (pe2(i,k) < pe1(i,1)) then - w(i,j,k) = 0. - endif - enddo - enddo -#endif - -#ifndef HYDRO_DELZ_REMAP -! Remap delz for hybrid sigma-p coordinate - do k=1,km - do i=is,ie - delz_r(i,j,k) = -delz_r(i,j,k)/delp_r(i,j,k) ! ="specific volume"/grav - enddo - enddo - call remap_2d(km, pe1, delz_r(is:ie,j:j,1:km), & - kn, pe2, delz(is:ie,j:j,1:kn), & - is, ie, 1, kord) - do k=1,kn - do i=is,ie - delz(i,j,k) = -delz(i,j,k)*delp(i,j,k) - enddo - enddo -#endif - endif - -! Geopotential conserving remap of virtual temperature: - do k=1,km+1 - do i=is,ie - pe1(i,k) = log(pe1(i,k)) - enddo - enddo - do k=1,kn+1 - do i=is,ie - pe2(i,k) = log(pe2(i,k)) - enddo - enddo - - call remap_2d(km, pe1, pt_r(is:ie,j:j,1:km), & - kn, pe2, pt(is:ie,j:j,1:kn), & - is, ie, 1, kord) - -#ifdef HYDRO_DELZ_REMAP - !initialize delz from the hydrostatic state - do k=1,kn - do i=is,ie - delz(i,j,k) = (rdgas*rgrav)*pt(i,j,k)*(pe2(i,k)-pe2(i,k+1)) - enddo - enddo -#endif -#ifdef HYDRO_DELZ_EXTRAP - !initialize delz from the hydrostatic state - do k=1,kn - do i=is,ie - if (pe2(i,k) < pe1(i,1)) then - delz(i,j,k) = (rdgas*rgrav)*pt(i,j,k)*(pe2(i,k)-pe2(i,k+1)) - endif - enddo - enddo -#endif -!------ -! map v -!------ - do k=1,km+1 - do i=is,ie+1 - pv1(i,k) = ak_r(k) + 0.5*bk_r(k)*(ps(i-1,j)+ps(i,j)) - enddo - enddo - do k=1,kn+1 - do i=is,ie+1 - pv2(i,k) = ak(k) + 0.5*bk(k)*(ps(i-1,j)+ps(i,j)) - enddo - enddo - - call remap_2d(km, pv1, v_r(is:ie+1,j:j,1:km), & - kn, pv2, v(is:ie+1,j:j,1:kn), & - is, ie+1, -1, kord) - - endif !(j < je+1) -6000 continue - -!$OMP parallel do default(none) shared(is,ie,js,je,kn,pt,r_vir,q) - do k=1,kn - do j=js,je - do i=is,ie - pt(i,j,k) = pt(i,j,k) / (1.+r_vir*q(i,j,k,1)) - enddo - enddo - enddo - - end subroutine rst_remap - -!>@brief The subroutine 'mappm' is a general-purpose routine for remapping -!! one set of vertical levels to another. - subroutine mappm(km, pe1, q1, kn, pe2, q2, i1, i2, iv, kord, ptop) - -! IV = 0: constituents -! IV = 1: potential temp -! IV =-1: winds - -! Mass flux preserving mapping: q1(im,km) -> q2(im,kn) - -! pe1: pressure at layer edges (from model top to bottom surface) -! in the original vertical coordinate -! pe2: pressure at layer edges (from model top to bottom surface) -! in the new vertical coordinate - - integer, intent(in):: i1, i2, km, kn, kord, iv - real, intent(in ):: pe1(i1:i2,km+1), pe2(i1:i2,kn+1) !< pe1: pressure at layer edges from model top to bottom - !! surface in the ORIGINAL vertical coordinate - !< pe2: pressure at layer edges from model top to bottom - !! surface in the NEW vertical coordinate -! Mass flux preserving mapping: q1(im,km) -> q2(im,kn) - real, intent(in ):: q1(i1:i2,km) - real, intent(out):: q2(i1:i2,kn) - real, intent(IN) :: ptop -! local - real qs(i1:i2) - real dp1(i1:i2,km) - real a4(4,i1:i2,km) - integer i, k, l - integer k0, k1 - real pl, pr, tt, delp, qsum, dpsum, esl - - do k=1,km - do i=i1,i2 - dp1(i,k) = pe1(i,k+1) - pe1(i,k) - a4(1,i,k) = q1(i,k) - enddo - enddo - - if ( kord >7 ) then - call cs_profile( qs, a4, dp1, km, i1, i2, iv, kord ) - else - call ppm_profile( a4, dp1, km, i1, i2, iv, kord ) - endif - -!------------------------------------ -! Lowest layer: constant distribution -!------------------------------------ -#ifdef NGGPS_SUBMITTED - do i=i1,i2 - a4(2,i,km) = q1(i,km) - a4(3,i,km) = q1(i,km) - a4(4,i,km) = 0. - enddo -#endif - - do 5555 i=i1,i2 - k0 = 1 - do 555 k=1,kn - - if(pe2(i,k) .le. pe1(i,1)) then -! above old ptop - q2(i,k) = q1(i,1) - elseif(pe2(i,k) .ge. pe1(i,km+1)) then -! Entire grid below old ps -#ifdef NGGPS_SUBMITTED - q2(i,k) = a4(3,i,km) ! this is not good. -#else - q2(i,k) = q1(i,km) -#endif - else - - do 45 L=k0,km -! locate the top edge at pe2(i,k) - if( pe2(i,k) .ge. pe1(i,L) .and. & - pe2(i,k) .le. pe1(i,L+1) ) then - k0 = L - PL = (pe2(i,k)-pe1(i,L)) / dp1(i,L) - if(pe2(i,k+1) .le. pe1(i,L+1)) then - -! entire new grid is within the original grid - PR = (pe2(i,k+1)-pe1(i,L)) / dp1(i,L) - TT = r3*(PR*(PR+PL)+PL**2) - q2(i,k) = a4(2,i,L) + 0.5*(a4(4,i,L)+a4(3,i,L) & - - a4(2,i,L))*(PR+PL) - a4(4,i,L)*TT - goto 555 - else -! Fractional area... - delp = pe1(i,L+1) - pe2(i,k) - TT = r3*(1.+PL*(1.+PL)) - qsum = delp*(a4(2,i,L)+0.5*(a4(4,i,L)+ & - a4(3,i,L)-a4(2,i,L))*(1.+PL)-a4(4,i,L)*TT) - dpsum = delp - k1 = L + 1 - goto 111 - endif - endif -45 continue - -111 continue - do 55 L=k1,km - if( pe2(i,k+1) .gt. pe1(i,L+1) ) then - -! Whole layer.. - - qsum = qsum + dp1(i,L)*q1(i,L) - dpsum = dpsum + dp1(i,L) - else - delp = pe2(i,k+1)-pe1(i,L) - esl = delp / dp1(i,L) - qsum = qsum + delp * (a4(2,i,L)+0.5*esl* & - (a4(3,i,L)-a4(2,i,L)+a4(4,i,L)*(1.-r23*esl)) ) - dpsum = dpsum + delp - k0 = L - goto 123 - endif -55 continue - delp = pe2(i,k+1) - pe1(i,km+1) - if(delp > 0.) then -! Extended below old ps -#ifdef NGGPS_SUBMITTED - qsum = qsum + delp * a4(3,i,km) ! not good. -#else - qsum = qsum + delp * q1(i,km) -#endif - dpsum = dpsum + delp - endif -123 q2(i,k) = qsum / dpsum - endif -555 continue -5555 continue - - end subroutine mappm - -!>@brief The subroutine 'moist_cv' computes the FV3-consistent moist heat capacity under constant volume, -!! including the heating capacity of water vapor and condensates. -!>@details See \cite emanuel1994atmospheric for information on variable heat capacities. - subroutine moist_cv(is,ie, isd,ied, jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, qd, cvm, t1) - integer, intent(in):: is, ie, isd,ied, jsd,jed, km, nwat, j, k - integer, intent(in):: sphum, liq_wat, rainwat, ice_wat, snowwat, graupel - real, intent(in), dimension(isd:ied,jsd:jed,km,nwat):: q - real, intent(out), dimension(is:ie):: cvm, qd - real, intent(in), optional:: t1(is:ie) -! - real, parameter:: t_i0 = 15. - real, dimension(is:ie):: qv, ql, qs - integer:: i - - select case (nwat) - - case(1) - do i=is,ie - qv(i) = max(q(i,j,k,sphum) ,0.0) - cvm(i) = (1.-qv(i))*cv_air + qv(i)*cv_vap - enddo - case(2) - if ( present(t1) ) then ! Special case for GFS physics - do i=is,ie - qd(i) = max(0., q(i,j,k,liq_wat)) - if ( t1(i) > tice ) then - qs(i) = 0. - elseif ( t1(i) < tice-t_i0 ) then - qs(i) = qd(i) - else - qs(i) = qd(i)*(tice-t1(i))/t_i0 - endif - ql(i) = qd(i) - qs(i) - qv(i) = max(0.,q(i,j,k,sphum)) - cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice - enddo - else - do i=is,ie - qv(i) = max(q(i,j,k,sphum) ,0.0) - qs(i) = max(q(i,j,k,liq_wat),0.0) - qd(i) = qs(i) - cvm(i) = (1.-qv(i))*cv_air + qv(i)*cv_vap - enddo - endif - case (3) - do i=is,ie - qv(i) = max(q(i,j,k,sphum) ,0.0) - ql(i) = max(q(i,j,k,liq_wat),0.0) - qs(i) = max(q(i,j,k,ice_wat),0.0) - qd(i) = ql(i) + qs(i) - cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice - enddo - case(4) ! K_warm_rain with fake ice - do i=is,ie - qv(i) = max(q(i,j,k,sphum) ,0.0) - ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) - cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + qd(i)*c_liq - enddo - case(5) - do i=is,ie - qv(i) = max(q(i,j,k,sphum) ,0.0) - ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) - qs(i) = max(q(i,j,k,ice_wat),0.0) + max(q(i,j,k,snowwat),0.0) - qd(i) = ql(i) + qs(i) - cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice - enddo - case(6:7) - do i=is,ie - qv(i) = max(q(i,j,k,sphum) ,0.0) - ql(i) = max(q(i,j,k,liq_wat),0.0) + max(q(i,j,k,rainwat),0.0) - qs(i) = max(q(i,j,k,ice_wat),0.0) + max(q(i,j,k,snowwat),0.0) + max(q(i,j,k,graupel),0.0) - qd(i) = ql(i) + qs(i) - cvm(i) = (1.-(qv(i)+qd(i)))*cv_air + qv(i)*cv_vap + ql(i)*c_liq + qs(i)*c_ice - enddo - case default - !call mpp_error (NOTE, 'fv_mapz::moist_cv - using default cv_air') - do i=is,ie - qd(i) = 0. - cvm(i) = cv_air - enddo - end select - - end subroutine moist_cv - -!>@brief The subroutine 'moist_cp' computes the FV3-consistent moist heat capacity under constant pressure, -!! including the heating capacity of water vapor and condensates. - subroutine moist_cp(is,ie, isd,ied, jsd,jed, km, j, k, nwat, sphum, liq_wat, rainwat, & - ice_wat, snowwat, graupel, q, qd, cpm, t1) - - integer, intent(in):: is, ie, isd,ied, jsd,jed, km, nwat, j, k - integer, intent(in):: sphum, liq_wat, rainwat, ice_wat, snowwat, graupel - real, intent(in), dimension(isd:ied,jsd:jed,km,nwat):: q - real, intent(out), dimension(is:ie):: cpm, qd - real, intent(in), optional:: t1(is:ie) -! - real, parameter:: t_i0 = 15. - real, dimension(is:ie):: qv, ql, qs - integer:: i - - select case (nwat) - - case(2) - if ( present(t1) ) then ! Special case for GFS physics - do i=is,ie - qd(i) = max(0., q(i,j,k,liq_wat)) - if ( t1(i) > tice ) then - qs(i) = 0. - elseif ( t1(i) < tice-t_i0 ) then - qs(i) = qd(i) - else - qs(i) = qd(i)*(tice-t1(i))/t_i0 - endif - ql(i) = qd(i) - qs(i) - qv(i) = max(0.,q(i,j,k,sphum)) - cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice - enddo - else - do i=is,ie - qv(i) = max(0.,q(i,j,k,sphum)) - qs(i) = max(0.,q(i,j,k,liq_wat)) - qd(i) = qs(i) - cpm(i) = (1.-qv(i))*cp_air + qv(i)*cp_vapor - enddo - endif - - case(3) - do i=is,ie - qv(i) = q(i,j,k,sphum) - ql(i) = q(i,j,k,liq_wat) - qs(i) = q(i,j,k,ice_wat) - qd(i) = ql(i) + qs(i) - cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice - enddo - case(4) ! K_warm_rain scheme with fake ice - do i=is,ie - qv(i) = q(i,j,k,sphum) - qd(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) - cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + qd(i)*c_liq - enddo - case(5) - do i=is,ie - qv(i) = q(i,j,k,sphum) - ql(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) - qs(i) = q(i,j,k,ice_wat) + q(i,j,k,snowwat) - qd(i) = ql(i) + qs(i) - cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice - enddo - case(6:7) - do i=is,ie - qv(i) = q(i,j,k,sphum) - ql(i) = q(i,j,k,liq_wat) + q(i,j,k,rainwat) - qs(i) = q(i,j,k,ice_wat) + q(i,j,k,snowwat) + q(i,j,k,graupel) - qd(i) = ql(i) + qs(i) - cpm(i) = (1.-(qv(i)+qd(i)))*cp_air + qv(i)*cp_vapor + ql(i)*c_liq + qs(i)*c_ice - enddo - case default - ! call mpp_error (NOTE, 'fv_mapz::moist_cp - using default cp_air') - do i=is,ie - qd(i) = 0. - cpm(i) = cp_air - enddo - end select - - end subroutine moist_cp - -!----------------------------------------------------------------------- -!BOP -! !ROUTINE: map1_gmao --- GMAO Interpolation for vertical re-mapping -! -! !INTERFACE: - subroutine map1_gmao( km, pe1, q1, & - kn, pe2, q2, i1, i2, & - j, ibeg, iend, jbeg, jend, akap, gmao_remap, P_MAP, conserv) - implicit none - -! !INPUT PARAMETERS: - integer, intent(in) :: i1 ! Starting longitude - integer, intent(in) :: i2 ! Finishing longitude - real, intent(in) :: akap - integer, intent(in) :: P_MAP ! Thermodynamic variable to remap - ! 1:TE-logP 2:PT-PK 3:T-logP - integer, intent(in) :: gmao_remap ! 3:cubic 2:quadratic 1:linear - logical, intent(in) :: conserv - integer, intent(in) :: j ! Current latitude - integer, intent(in) :: ibeg, iend, jbeg, jend - integer, intent(in) :: km ! Original vertical dimension - integer, intent(in) :: kn ! Target vertical dimension - - real, intent(in) :: pe1(i1:i2,km+1) ! pressure at layer edges - ! (from model top to bottom surface) - ! in the original vertical coordinate - real, intent(in) :: pe2(i1:i2,kn+1) ! pressure at layer edges - ! (from model top to bottom surface) - ! in the new vertical coordinate - - real, intent(in) :: q1(ibeg:iend,jbeg:jend,km) ! Field input -! !INPUT/OUTPUT PARAMETERS: - real, intent(inout):: q2(ibeg:iend,jbeg:jend,kn) ! Field output - -! !DESCRIPTION: -! -! Perform Cubic Interpolation a given latitude -! pe1: pressure at layer edges (from model top to bottom surface) -! in the original vertical coordinate -! pe2: pressure at layer edges (from model top to bottom surface) -! in the new vertical coordinate -! -! !REVISION HISTORY: -! 2005.11.14 Takacs Initial Code -! 2016.07.20 Putman Modified to make genaric for any thermodynamic variable -! -!EOP -!----------------------------------------------------------------------- -!BOC -! -! !LOCAL VARIABLES: - real qx(i1:i2,km) - real logpl1(i1:i2,km) - real logpl2(i1:i2,kn) - real dlogp1(i1:i2,km) - real vsum1(i1:i2) - real vsum2(i1:i2) - real am2,am1,ap0,ap1,P,PLP1,PLP0,PLM1,PLM2,DLP0,DLM1,DLM2 - - integer i, k, LM2,LM1,LP0,LP1 - -! Initialization -! -------------- - - select case (P_MAP) - case(1) - ! Remapping in Log(P) - do k=1,km - qx(:,k) = q1(i1:i2,j,k) - logpl1(:,k) = log( r2*(pe1(:,k)+pe1(:,k+1)) ) - enddo - do k=1,kn - logpl2(:,k) = log( r2*(pe2(:,k)+pe2(:,k+1)) ) - enddo - do k=1,km-1 - dlogp1(:,k) = logpl1(:,k+1)-logpl1(:,k) - enddo - case(2) - ! Remapping in P^KAPPA - do k=1,km - qx(:,k) = q1(i1:i2,j,k) - logpl1(:,k) = exp( akap*log( r2*(pe1(:,k)+pe1(:,k+1))) ) - enddo - do k=1,kn - logpl2(:,k) = exp( akap*log( r2*(pe2(:,k)+pe2(:,k+1))) ) - enddo - do k=1,km-1 - dlogp1(:,k) = logpl1(:,k+1)-logpl1(:,k) - enddo - end select - - if (conserv) then -! Compute vertical integral of Input field -! ---------------------------------------- - vsum1(:) = r0 - do i=i1,i2 - do k=1,km - vsum1(i) = vsum1(i) + qx(i,k)*( pe1(i,k+1)-pe1(i,k) ) - enddo - vsum1(i) = vsum1(i) / ( pe1(i,km+1)-pe1(i,1) ) - enddo - endif - -! Interpolate field onto target Pressures -! --------------------------------------- - do i=i1,i2 - do k=1,kn - LM1 = 1 - LP0 = 1 - do while( LP0.le.km ) - if (logpl1(i,LP0).lt.logpl2(i,k)) then - LP0 = LP0+1 - else - exit - endif - enddo - LM1 = max(LP0-1,1) - LP0 = min(LP0, km) - -! Extrapolate Linearly above first model level -! ---------------------------------------------------- - if( LM1.eq.1 .and. LP0.eq.1 ) then - q2(i,j,k) = qx(i,1) + ( qx(i,2)-qx(i,1) )*( logpl2(i,k)-logpl1(i,1) ) & - /( logpl1(i,2)-logpl1(i,1) ) - -! Extrapolate Linearly below last model level -! --------------------------------------------------- - else if( LM1.eq.km .and. LP0.eq.km ) then - q2(i,j,k) = qx(i,km) + ( qx(i,km)-qx(i,km-1) )*( logpl2(i,k )-logpl1(i,km ) ) & - /( logpl1(i,km)-logpl1(i,km-1) ) - -! Interpolate Linearly between levels 1 => 2 and km-1 => km -! ----------------------------------------------------------------- - else if( LM1.eq.1 .or. LP0.eq.km ) then - q2(i,j,k) = qx(i,LP0) + ( qx(i,LM1)-qx(i,LP0) )*( logpl2(i,k )-logpl1(i,LP0) ) & - /( logpl1(i,LM1)-logpl1(i,LP0) ) -! Interpolate Cubicly between other model levels -! ------------------------------------------------------ - else - LP1 = LP0+1 - LM2 = LM1-1 - P = logpl2(i,k) - PLP1 = logpl1(i,LP1) - PLP0 = logpl1(i,LP0) - PLM1 = logpl1(i,LM1) - PLM2 = logpl1(i,LM2) - DLP0 = dlogp1(i,LP0) - DLM1 = dlogp1(i,LM1) - DLM2 = dlogp1(i,LM2) - - ! Cubic Coefficients - ! ------------------ - if( gmao_remap .eq. 3 ) then - ap1 = (P-PLP0)*(P-PLM1)*(P-PLM2)/( DLP0*(DLP0+DLM1)*(DLP0+DLM1+DLM2) ) - ap0 = (PLP1-P)*(P-PLM1)*(P-PLM2)/( DLP0* DLM1 *( DLM1+DLM2) ) - am1 = (PLP1-P)*(PLP0-P)*(P-PLM2)/( DLM1* DLM2 *(DLP0+DLM1 ) ) - am2 = (PLP1-P)*(PLP0-P)*(PLM1-P)/( DLM2*(DLM1+DLM2)*(DLP0+DLM1+DLM2) ) - q2(i,j,k) = ap1*qx(i,LP1) + ap0*qx(i,LP0) + am1*qx(i,LM1) + am2*qx(i,LM2) - endif - - ! Quadratic Coefficients - ! ---------------------- - if( gmao_remap .eq. 2 ) then - ap1 = (P-PLP0)*(P-PLM1)/( (PLP1-PLP0)*(PLP1-PLM1) ) - ap0 = (PLP1-P)*(P-PLM1)/( (PLP1-PLP0)*(PLP0-PLM1) ) - am1 = (PLP1-P)*(PLP0-P)/( (PLP1-PLM1)*(PLP0-PLM1) ) - q2(i,j,k) = ap1*qx(i,LP1) + ap0*qx(i,LP0) + am1*qx(i,LM1) - endif - - ! Linear Coefficients - ! ------------------- - if( gmao_remap .eq. 1 ) then - q2(i,j,k) = qx(i,LP0) + ( qx(i,LM1)-qx(i,LP0) )*( logpl2(i,k )-logpl1(i,LP0) ) & - /( logpl1(i,LM1)-logpl1(i,LP0) ) - endif - - endif - - enddo - enddo - - if (conserv) then -! Compute vertical integral of Output field -! ----------------------------------------- - vsum2(:) = r0 - do i=i1,i2 - do k=1,kn - vsum2(i) = vsum2(i) + q2(i,j,k)*( pe2(i,k+1)-pe2(i,k) ) - enddo - vsum2(i) = vsum2(i) / ( pe2(i,kn+1)-pe2(i,1) ) - enddo -! Adjust Final field to conserve -! ------------------------------ - do i=i1,i2 - do k=1,kn - q2(i,j,k) = q2(i,j,k) + vsum1(i)-vsum2(i) - enddo - enddo - endif - - return -!EOC - end subroutine map1_gmao -!----------------------------------------------------------------------- - - -end module fv_mapz_mod From 2ebc330b1791bf37426b491defc725ee429dcbcb Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 11:52:22 -0500 Subject: [PATCH 123/252] Update README --- README.md | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 87c0edaa..c66162b1 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,9 @@ Branches: - ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit -- ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit +- ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit and `dpx` calculation + - MERGE ORDER: after `fix/GEOS/D_SW` - ⚙️ `feature/tracer_rework_part1` @Florian: Allow for update of N Tracers -- ⚙️ `fix/GEOS/TracerAdvection` @Florian [BASED ON `tracer_rework_part1`]: Allow for non-update of mass fluxes and courant number, f32 fixes, correct computation of `cmax` and `nsplit`, overcomputation into the algorithm +- ⚙️ `fix/GEOS/TracerAdvection` @Florian: Allow for non-update of mass fluxes and courant number, f32 fixes, correct computation of `cmax` and `nsplit`, overcomputation into the algorithm + - BASED ON `tracer_rework_part1` + - REQUIRES: `ndsl` with tracer rework From f34fc221a951e51ac8754e9361488e1691c5982a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 11:53:34 -0500 Subject: [PATCH 124/252] Adapt XPPM translate test to serialization tweaks in Fortran --- tests/savepoint/translate/translate_xppm.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/savepoint/translate/translate_xppm.py b/tests/savepoint/translate/translate_xppm.py index 44d8b368..2789c747 100644 --- a/tests/savepoint/translate/translate_xppm.py +++ b/tests/savepoint/translate/translate_xppm.py @@ -14,12 +14,12 @@ def __init__( ): super().__init__(grid, namelist, stencil_factory) self.in_vars["data_vars"] = { - "q": {"serialname": "qx", "jstart": "jfirst"}, - "c": {"serialname": "cx", "istart": grid.is_}, + "q": {"serialname": "xppm_q", "jstart": "jfirst"}, + "c": {"serialname": "xppm_c", "istart": grid.is_}, } self.in_vars["parameters"] = ["iord", "jfirst", "jlast"] self.out_vars = { - "xflux": { + "xppm_flux": { "istart": grid.is_, "iend": grid.ie + 1, "jstart": "jfirst", @@ -40,7 +40,7 @@ def process_inputs(self, inputs): def compute(self, inputs): self.process_inputs(inputs) - inputs["xflux"] = utils.make_storage_from_shape( + inputs["xppm_flux"] = utils.make_storage_from_shape( inputs["q"].shape, backend=self.stencil_factory.backend ) origin = self.grid.grid_indexing.origin_compute() @@ -53,7 +53,7 @@ def compute(self, inputs): origin=(origin[0], int(inputs["jfirst"]), origin[2]), domain=(domain[0], int(inputs["jlast"] - inputs["jfirst"] + 1), domain[2]), ) - self.compute_func(inputs["q"], inputs["c"], inputs["xflux"]) + self.compute_func(inputs["q"], inputs["c"], inputs["xppm_flux"]) return self.slice_output(inputs) @@ -65,5 +65,5 @@ def __init__( stencil_factory: StencilFactory, ): super().__init__(grid, namelist, stencil_factory) - self.in_vars["data_vars"]["q"]["serialname"] = "q" - self.out_vars["xflux"]["serialname"] = "xflux_2" + self.in_vars["data_vars"]["q"]["serialname"] = "xppm_q2" + self.out_vars["xppm_flux"]["serialname"] = "xppm_flux_2" From 21f4d349af334a6880fb517c74821b15f1093c12 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 11:56:49 -0500 Subject: [PATCH 125/252] Move `dpx` to be passed in the acoustics instead of temp --- pyFV3/stencils/dyn_core.py | 57 ++++++++++++++---------------- pyFV3/testing/translate_dyncore.py | 21 +++++++++-- 2 files changed, 44 insertions(+), 34 deletions(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 8a693e36..10312d7e 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -239,13 +239,6 @@ def dyncore_temporaries( units="unknown", dtype=Float, ) - for name in ["dpx"]: - temporaries[name] = quantity_factory.zeros( - dims=[X_DIM, Y_DIM, Z_DIM], - units="unknown", - dtype=np.float64, - allow_mismatch_float_precision=True, - ) return temporaries @@ -737,6 +730,7 @@ def dt2(self, dt_acoustic_substep: Float) -> Float: def __call__( self, state: DycoreState, + dpx, timestep: Float, # time to step forward by in seconds n_map=1, # [DaCe] replaces state.n_map ): @@ -888,30 +882,31 @@ def __call__( # by 1 timestep self._checkpoint_dsw_in(state) self.dgrid_shallow_water_lagrangian_dynamics( - self._vt, - state.delp, - state.pt, - state.u, - state.v, - state.w, - state.uc, - state.vc, - state.ua, - state.va, - self._divgd, - state.mfxd, - state.mfyd, - state.cxd, - state.cyd, - self._crx, - self._cry, - self._xfx, - self._yfx, - state.q_con, - self._zh, - self._heat_source, - state.diss_estd, - dt_acoustic_substep, + delpc=self._vt, + delp=state.delp, + pt=state.pt, + u=state.u, + v=state.v, + w=state.w, + uc=state.uc, + vc=state.vc, + ua=state.ua, + va=state.va, + divgd=self._divgd, + mfx=state.mfxd, + mfy=state.mfyd, + cx=state.cxd, + cy=state.cyd, + dpx=dpx, + crx=self._crx, + cry=self._cry, + xfx=self._xfx, + yfx=self._yfx, + q_con=state.q_con, + zh=self._zh, + heat_source=self._heat_source, + diss_est=state.diss_estd, + dt=dt_acoustic_substep, ) self._checkpoint_dsw_out(state) # note that uc and vc are not needed at all past this point. diff --git a/pyFV3/testing/translate_dyncore.py b/pyFV3/testing/translate_dyncore.py index 3c398f4a..17718bce 100644 --- a/pyFV3/testing/translate_dyncore.py +++ b/pyFV3/testing/translate_dyncore.py @@ -2,6 +2,7 @@ from ndsl import Namelist, Quantity, StencilFactory from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.stencils.testing import ParallelTranslate2PyState +from numpy import dtype from pyFV3._config import DynamicalCoreConfig from pyFV3.dycore_state import DycoreState from pyFV3.stencils import dyn_core @@ -101,6 +102,7 @@ def __init__( "ak": {}, "bk": {}, "diss_estd": {}, + "dpx": grid.compute_dict(), } self._base.in_vars["data_vars"]["wsd"]["kstart"] = grid.npz self._base.in_vars["data_vars"]["wsd"]["kend"] = None @@ -149,7 +151,7 @@ def compute_parallel(self, inputs, communicator): dtype_dict=inputs_dtypes, allow_mismatch_float_precision=True, ) - wsd: Quantity = self.grid.quantity_factory.zeros( + wsd = self.grid.quantity_factory.zeros( dims=[X_DIM, Y_DIM], units="unknown", ) @@ -161,11 +163,18 @@ def compute_parallel(self, inputs, communicator): state[name].data[selection] = value else: setattr(state, name, value) - phis: Quantity = self.grid.quantity_factory.zeros( + phis = self.grid.quantity_factory.zeros( dims=[X_DIM, Y_DIM], units="m", ) phis.data[:] = phis.np.asarray(inputs["phis"]) + dpx = self.grid.quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=inputs_dtypes["dpx"], + allow_mismatch_float_precision=True, + ) + dpx.data[:] = dpx.np.asarray(inputs["dpx"]) acoustic_dynamics = dyn_core.AcousticDynamics( comm=communicator, stencil_factory=self.stencil_factory, @@ -182,7 +191,12 @@ def compute_parallel(self, inputs, communicator): ) acoustic_dynamics.cappa.data[:] = inputs["cappa"][:] - acoustic_dynamics(state, timestep=inputs["mdt"], n_map=state.n_map) + acoustic_dynamics( + state, + dpx=dpx, + timestep=inputs["mdt"], + n_map=state.n_map, + ) # the "inputs" dict is not used to return, we construct a new dict based # on variables attached to `state` storages_only = {} @@ -193,4 +207,5 @@ def compute_parallel(self, inputs, communicator): storages_only[name] = value storages_only["wsd"] = wsd.data storages_only["cappa"] = acoustic_dynamics.cappa.data + storages_only["dpx"] = dpx.data return self._base.slice_output(storages_only) From 036fc2ff3e392442e9550497a511d58574ce4fdc Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 11:59:23 -0500 Subject: [PATCH 126/252] `dpx` calculated on cell center --- pyFV3/stencils/d_sw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/d_sw.py b/pyFV3/stencils/d_sw.py index 9e32de6f..754bdd54 100644 --- a/pyFV3/stencils/d_sw.py +++ b/pyFV3/stencils/d_sw.py @@ -1037,7 +1037,7 @@ def make_quantity(): ) self._accumulate_delp = stencil_factory.from_dims_halo( func=delp_increment_accumulation, - compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], + compute_dims=[X_DIM, Y_DIM, Z_DIM], ) def __call__( From 6d28e5b0798a8e8f72e53909bf695a21d6214f59 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 12:22:17 -0500 Subject: [PATCH 127/252] Add remapping reference to README --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index c66162b1..7a1838ab 100644 --- a/README.md +++ b/README.md @@ -41,3 +41,4 @@ Branches: - ⚙️ `fix/GEOS/TracerAdvection` @Florian: Allow for non-update of mass fluxes and courant number, f32 fixes, correct computation of `cmax` and `nsplit`, overcomputation into the algorithm - BASED ON `tracer_rework_part1` - REQUIRES: `ndsl` with tracer rework +- ⚙️ `feature/fv_mapz/GEOS` @ Chris K: Remapping for GEOS From b3c79414b1e0726ee20b3e8ef9a76e11292a34c1 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 15:15:50 -0500 Subject: [PATCH 128/252] Lint --- pyFV3/stencils/moist_cv.py | 78 +++++++++++++++++----------- pyFV3/stencils/mpp_global_sum.py | 89 +++++++++++++++++--------------- 2 files changed, 95 insertions(+), 72 deletions(-) diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index f975009d..a7bfba94 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -24,8 +24,8 @@ def set_cappa(qvapor, cvm, r_vir): def moist_cvm(qvapor, gz, ql, qs): # CK : GEOS applies the "max" function to tracer values cvm = ( - (1.0 - (max(qvapor,0.0) + gz)) * constants.CV_AIR - + max(qvapor,0.0) * constants.CV_VAP + (1.0 - (max(qvapor, 0.0) + gz)) * constants.CV_AIR + + max(qvapor, 0.0) * constants.CV_VAP + ql * constants.C_LIQ + qs * constants.C_ICE ) @@ -43,7 +43,7 @@ def moist_cv_nwat6_fn( ): # CK : GEOS applies the "max" function to tracer values ql = max(qliquid, 0.0) + max(qrain, 0.0) - qs = max(qice,0.0) + max(qsnow,0.0) + max(qgraupel,0.0) + qs = max(qice, 0.0) + max(qsnow, 0.0) + max(qgraupel, 0.0) gz = ql + qs cvm = moist_cvm(qvapor, gz, ql, qs) return cvm, gz @@ -169,6 +169,7 @@ def moist_pkz( cappa = set_cappa(qvapor, cvm, r_vir) pkz = compute_pkz_func(delp, delz, pt, cappa) + def moist_te( qvapor: FloatField, qliquid: FloatField, @@ -187,7 +188,7 @@ def moist_te( cosa_s: FloatFieldIJ, hs: FloatFieldIJ, delz: FloatField, - grav: Float + grav: Float, ): """ Args: @@ -208,46 +209,63 @@ def moist_te( cosa_s (in): hs (in): """ - with computation(FORWARD), interval(-1,None): + with computation(FORWARD), interval(-1, None): te = 0.0 phis = hs - with computation(BACKWARD), interval(0,-1): - phis = phis[0,0,1] - grav*delz - with computation(FORWARD), interval(0,-1): - cvm, gz = moist_cv_nwat6_fn( - qvapor, qliquid, qrain, qsnow, qice, qgraupel + with computation(BACKWARD), interval(0, -1): + phis = phis[0, 0, 1] - grav * delz + with computation(FORWARD), interval(0, -1): + cvm, gz = moist_cv_nwat6_fn(qvapor, qliquid, qrain, qsnow, qice, qgraupel) + + te = te + delp * ( + cvm * pt + + 0.5 + * ( + phis + + phis[0, 0, 1] + + w**2.0 + + 0.5 + * rsin2 + * ( + u**2.0 + + u[0, 1, 0] ** 2.0 + + v**2.0 + + v[1, 0, 0] ** 2.0 + - (u + u[0, 1, 0]) * (v + v[1, 0, 0]) * cosa_s + ) + ) ) - - te = te + delp * (cvm * pt + 0.5 * (phis + phis[0,0,1] + \ - w**2.0 + 0.5*rsin2 * (u**2.0 + u[0,1,0]**2.0 \ - + v**2.0 + v[1,0,0]**2.0 - - (u + u[0,1,0]) * (v + v[1,0,0]) * cosa_s))) -def te_zsum(te_2d: FloatFieldIJ, - te0_2d: FloatFieldIJ, - delp: FloatField, - pkz: FloatField, - zsum1: FloatFieldIJ, - ): + +def te_zsum( + te_2d: FloatFieldIJ, + te0_2d: FloatFieldIJ, + delp: FloatField, + pkz: FloatField, + zsum1: FloatFieldIJ, +): with computation(FORWARD): - with interval(0,1): + with interval(0, 1): te_2d = te0_2d - te_2d zsum1 = pkz * delp - with interval(1,None): + with interval(1, None): zsum1 = zsum1 + pkz * delp -def cond_output(q_con: FloatField, - qliquid: FloatField, - qrain: FloatField, - qsnow: FloatField, - qice: FloatField, - qgraupel: FloatField, - ): + +def cond_output( + q_con: FloatField, + qliquid: FloatField, + qrain: FloatField, + qsnow: FloatField, + qice: FloatField, + qgraupel: FloatField, +): with computation(PARALLEL), interval(...): q_con = 0.0 q_con = qliquid + qice + qrain + qsnow + qgraupel + def fv_setup( qvapor: FloatField, qliquid: FloatField, diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py index 6398c6db..61cffe78 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/stencils/mpp_global_sum.py @@ -2,36 +2,38 @@ from ndsl.comm.comm_abc import ReductionOperator import numpy as np + def mpp_global_sum(inputArray, communicator, stencil_factory=None): - NUMINT = 6 NUMBIT = 46 - r_prec = 2.0 ** NUMBIT - prec = 2 ** NUMBIT - I_prec = 1.0 / (2.0 ** NUMBIT) - pr = [r_prec**2, r_prec, 1.0, 1.0/r_prec, 1.0/r_prec**2, 1.0/r_prec**3] - I_pr = [1.0/r_prec**2, 1.0/r_prec, 1.0, r_prec, r_prec**2, r_prec**3 ] + r_prec = 2.0**NUMBIT + prec = 2**NUMBIT + I_prec = 1.0 / (2.0**NUMBIT) + pr = [r_prec**2, r_prec, 1.0, 1.0 / r_prec, 1.0 / r_prec**2, 1.0 / r_prec**3] + I_pr = [1.0 / r_prec**2, 1.0 / r_prec, 1.0, r_prec, r_prec**2, r_prec**3] prec_error = (2**62 + (2**62 - 1)) / 6 mag_max_term = 0.0 ints_sum = Quantity( - data=np.zeros((NUMINT),dtype=np.float32), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - + data=np.zeros((NUMINT), dtype=np.float32), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + ints_sum_reduce = Quantity( - data=np.zeros((NUMINT),dtype=np.float32), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - + data=np.zeros((NUMINT), dtype=np.float32), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + # Note: This loop range in i and j are for the TBC test case. for j in range(inputArray.shape[1]): for i in range(inputArray.shape[0]): - increment_ints_faster(ints_sum.data, pr, I_pr, inputArray[i,j], mag_max_term) + increment_ints_faster( + ints_sum.data, pr, I_pr, inputArray[i, j], mag_max_term + ) # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' before carry_over') carry_overflow(ints_sum.data, prec, I_prec, prec_error) @@ -47,64 +49,67 @@ def mpp_global_sum(inputArray, communicator, stencil_factory=None): return sum_ - + def increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): - if((r >= 1e30) == r < 1e30): + if (r >= 1e30) == r < 1e30: print("NaN_error") return sgn = 1 - if(r < 0.0): + if r < 0.0: sgn = -1 rs = abs(r) - if(rs > abs(max_mag_term)): + if rs > abs(max_mag_term): max_mag_term = r for i in range(len(I_pr)): - ival = int(rs*I_pr[i]) - rs = rs - ival*pr[i] - int_sum[i] = int_sum[i] + sgn*ival + ival = int(rs * I_pr[i]) + rs = rs - ival * pr[i] + int_sum[i] = int_sum[i] + sgn * ival + def carry_overflow(int_sum, prec, I_prec, prec_error): - for i in range(len(int_sum)-1,1,-1): + for i in range(len(int_sum) - 1, 1, -1): if abs(int_sum[i]) > prec: num_carry = int(int_sum[i] * I_prec) - int_sum[i] = int_sum[i] - num_carry*prec - int_sum[i-1] = int_sum[i-1] + num_carry + int_sum[i] = int_sum[i] - num_carry * prec + int_sum[i - 1] = int_sum[i - 1] + num_carry if abs(int_sum[0]) > prec_error: overflow_error = True + def regularize_ints(int_sum, prec, I_prec): - for i in range(len(int_sum)-1, 0, -1): + for i in range(len(int_sum) - 1, 0, -1): if abs(int_sum[i]) > prec: num_carry = int(int_sum[i] * I_prec) - int_sum[i] = int_sum[i] - num_carry*prec - int_sum[i-1] = int_sum[i-1] + num_carry - + int_sum[i] = int_sum[i] - num_carry * prec + int_sum[i - 1] = int_sum[i - 1] + num_carry + positive = True for i in range(len(int_sum)): if abs(int_sum[i]) > 0: if int_sum[i] < 0: - positive = False + positive = False break - + if positive: - for i in range(len(int_sum)-1, 1, -1): + for i in range(len(int_sum) - 1, 1, -1): if int_sum[i] < 0: int_sum[i] = int_sum[i] + prec - int_sum[i-1] = int_sum[i-1] - 1 - + int_sum[i - 1] = int_sum[i - 1] - 1 + else: - for i in range(len(int_sum)-1, 1, -1): + for i in range(len(int_sum) - 1, 1, -1): if int_sum[i] > 0: int_sum[i] = int_sum[i] - prec - int_sum[i-1] = int_sum[i-1] + 1 + int_sum[i - 1] = int_sum[i - 1] + 1 + def ints_to_real(ints, pr): r = 0.0 for i in range(len(ints)): - r = r + pr[i]*ints[i] + r = r + pr[i] * ints[i] - return r \ No newline at end of file + return r From 6c8b424b0b853d603cd28ddb77c5bd4cd646b10a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 15:16:11 -0500 Subject: [PATCH 129/252] Remapping GEOS obejct (untested) --- pyFV3/stencils/remapping_GEOS.py | 610 +++++++++++++++++++++++++++++++ 1 file changed, 610 insertions(+) create mode 100644 pyFV3/stencils/remapping_GEOS.py diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py new file mode 100644 index 00000000..271bc1e2 --- /dev/null +++ b/pyFV3/stencils/remapping_GEOS.py @@ -0,0 +1,610 @@ +from typing import Dict +from ndsl import ( + Quantity, + QuantityFactory, + StencilFactory, + orchestrate, +) +from ndsl.constants import ( + X_DIM, + X_INTERFACE_DIM, + Y_DIM, + Y_INTERFACE_DIM, + Z_DIM, + Z_INTERFACE_DIM, + GRAV, + CV_AIR, +) +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK +from ndsl.stencils.basic_operations import adjust_divide_stencil +from ndsl.grid import GridData +from ndsl.comm.communicator import Communicator +from pyFV3._config import RemappingConfig +from pyFV3.stencils import moist_cv +from pyFV3.stencils.map_single import MapSingle +from pyFV3.stencils.mapn_tracer import MapNTracer +from pyFV3.stencils.moist_cv import moist_pt_last_step +from pyFV3.stencils.saturation_adjustment import SatAdjust3d +from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 +from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment +from pyFV3.stencils.remapping import ( + init_pe, + moist_cv_pt_pressure, + pn2_pk_delp, + pressures_mapu, + pressures_mapv, + pe0_ptop_xmax, + pe_pk_delp_peln, + CONSV_MIN, +) +from pyFV3.stencils.mpp_global_sum import mpp_global_sum + + +class LagrangianToEulerian_GEOS: + """ + GEOS v11.4.2 remapping - derived from original fvcore. + + Fortran name is Lagrangian_to_Eulerian + """ + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + config: RemappingConfig, + comm: Communicator, + grid_data: GridData, + nq, + pfull, + tracers: Dict[str, Quantity], + ): + orchestrate( + obj=self, + config=stencil_factory.config.dace_config, + dace_compiletime_args=["tracers"], + ) + self._comm = comm + self._stencil_factory = stencil_factory + grid_indexing = stencil_factory.grid_indexing + + # Configuration + if config.kord_tm >= 0: + raise NotImplementedError("map ppm, untested mode where kord_tm >= 0") + hydrostatic = config.hydrostatic + if hydrostatic: + raise NotImplementedError("Hydrostatic is not implemented") + + self._t_min = Float(184.0) + self._nq = nq + self._w_max = Float(90.0) + self._w_min = Float(-60.0) + self._area_64 = grid_data.area_64 + self._cosa = grid_data.cosa_s + self._rsin2 = grid_data.rsin2 + self._kord_tm = abs(config.kord_tm) + self._kord_wz = config.kord_wz + self._kord_mt = config.kord_mt + self._do_sat_adjust = config.do_sat_adj + self.kmp = grid_indexing.domain[2] - 1 + for k in range(pfull.shape[0]): + if pfull.view[k] > 10.0e2: + self.kmp = k + break + # do_omega = hydrostatic and last_step # TODO pull into inputs + + # Quantities + self._domain_jextra = ( + grid_indexing.domain[0], + grid_indexing.domain[1] + 1, + grid_indexing.domain[2] + 1, + ) + + self._pe1 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", + dtype=Float, + ) + self._pe2 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", + dtype=Float, + ) + self._pe3 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", + dtype=Float, + ) + self._dp2 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", + dtype=Float, + ) + self._pn2 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", + dtype=Float, + ) + self._pe0 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", + dtype=Float, + ) + self._pe3 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", + dtype=Float, + ) + + self._gz = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="m^2 s^-2", + dtype=Float, + ) + self._cvm = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._compute_performed = quantity_factory.zeros( + [X_DIM, Y_DIM], + units="mask", + dtype=bool, + ) + self._w2 = quantity_factory._numpy.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="temp W", + dtype=Float, + ) + self._pk2 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", + dtype=Float, + ) + + self._te_2d = quantity_factory.zeros( + [X_DIM, Y_DIM], + units="Pa", + dtype=Float, + ) + + self._te0_2d = quantity_factory.zeros( + [X_DIM, Y_DIM], + units="Pa", + dtype=Float, + ) + + self._zsum1 = quantity_factory.zeros( + [X_DIM, Y_DIM], + units="Pa", + dtype=Float, + ) + self._phis = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="n/a", + dtype=Float, + ) + + # Stencils + + self._init_pe = stencil_factory.from_origin_domain( + init_pe, origin=grid_indexing.origin_compute(), domain=self._domain_jextra + ) + + self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( + moist_cv_pt_pressure, + # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, + externals={"hydrostatic": hydrostatic}, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(add=(0, 0, 1)), + ) + + self._pn2_pk_delp = stencil_factory.from_origin_domain( + pn2_pk_delp, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + self._map_single_pt = MapSingle( + stencil_factory, + quantity_factory, + self._kord_tm, + mode=1, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._mapn_tracer = MapNTracer( + stencil_factory, + quantity_factory, + abs(config.kord_tr), + nq, + fill=config.fill, + tracers=tracers, + ) + + self._map_single_w = MapSingle( + stencil_factory, + quantity_factory, + self._kord_wz, + mode=-2, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._map_single_delz = MapSingle( + stencil_factory, + quantity_factory, + self._kord_wz, + mode=1, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._moist_cv_pkz = stencil_factory.from_origin_domain( + moist_cv.moist_pkz, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + self._pressures_mapu = stencil_factory.from_origin_domain( + pressures_mapu, + origin=grid_indexing.origin_compute(), + domain=self._domain_jextra, + ) + + self._map_single_u = MapSingle( + stencil_factory, + quantity_factory, + self._kord_mt, + mode=-1, + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + ) + + self._pressures_mapv = stencil_factory.from_origin_domain( + pressures_mapv, + origin=grid_indexing.origin_compute(), + domain=( + grid_indexing.domain[0] + 1, + grid_indexing.domain[1], + grid_indexing.domain[2] + 1, + ), + ) + + self._map_single_v = MapSingle( + stencil_factory, + quantity_factory, + self._kord_mt, + mode=-1, + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + ) + + self._saturation_adjustment = SatAdjust3d( + stencil_factory, config.sat_adjust, self._area_64, self.kmp + ) + + self._moist_cv_last_step_stencil = stencil_factory.from_origin_domain( + moist_pt_last_step, + origin=(grid_indexing.isc, grid_indexing.jsc, 0), + domain=( + grid_indexing.domain[0], + grid_indexing.domain[1], + grid_indexing.domain[2] + 1, + ), + ) + + self._basic_adjust_divide_stencil = stencil_factory.from_origin_domain( + adjust_divide_stencil, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + self._rescale_delz_1 = stencil_factory.from_origin_domain( + rescale_delz_1, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + self._rescale_delz_2 = stencil_factory.from_origin_domain( + rescale_delz_2, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + self._w_fix_consrv_moment = stencil_factory.from_origin_domain( + func=W_fix_consrv_moment, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + self._pe0_ptop_xmax = stencil_factory.from_origin_domain( + pe0_ptop_xmax, + origin=( + grid_indexing.n_halo + grid_indexing.domain[0], + grid_indexing.n_halo, + 0, + ), + domain=(1, grid_indexing.domain[1], 1), + ) + self._pe_pk_delp_peln = stencil_factory.from_origin_domain( + pe_pk_delp_peln, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(add=(0, 0, 1)), + ) + self._moist_cv_te = stencil_factory.from_origin_domain( + moist_cv.moist_te, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(add=(0, 0, 1)), + ) + + self._te_zsum = stencil_factory.from_origin_domain( + moist_cv.te_zsum, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + def __call__( + self, + tracers: Dict[str, Quantity], + pt: FloatField, # type: ignore + delp: FloatField, # type: ignore + delz: FloatField, # type: ignore + peln: FloatField, # type: ignore + u: FloatField, # type: ignore + v: FloatField, # type: ignore + w: FloatField, # type: ignore + mfx: FloatField, # type: ignore + mfy: FloatField, # type: ignore + cx: FloatField, # type: ignore + cy: FloatField, # type: ignore + cappa: FloatField, # type: ignore + q_con: FloatField, # type: ignore + q_cld: FloatField, # type: ignore + pkz: FloatField, # type: ignore + pk: FloatField, # type: ignore + pe: FloatField, # type: ignore + hs: FloatFieldIJ, # type: ignore + ps: FloatFieldIJ, # type: ignore + wsd: FloatFieldIJ, # type: ignore + ak: FloatFieldK, # type: ignore + bk: FloatFieldK, # type: ignore + dp1: FloatField, # type: ignore + ptop: Float, # type: ignore + akap: Float, # type: ignore + zvir: Float, # type: ignore + last_step: bool, + consv_te: Float, # type: ignore + mdt: Float, # type: ignore + ): + """ + Remap the deformed Lagrangian surfaces onto the reference, or "Eulerian", + coordinate levels. + + Args: + tracers (inout): Tracer species tracked across + pt (inout): D-grid potential temperature + delp (inout): Pressure Thickness + delz (in): Vertical thickness of atmosphere layers + peln (inout): Logarithm of interface pressure + u (inout): D-grid x-velocity + v (inout): D-grid y-velocity + w (inout): Vertical velocity + ua (inout): A-grid x-velocity + va (inout): A-grid y-velocity + cappa (inout): Power to raise pressure to + q_con (out): Total condensate mixing ratio + q_cld (out): Cloud fraction + pkz (in): Layer mean pressure raised to the power of Kappa + pk (out): Interface pressure raised to power of kappa, final acoustic value + pe (in): Pressure at layer edges + hs (in): Surface geopotential + te0_2d (unused): Atmosphere total energy in columns + ps (out): Surface pressure + wsd (in): Vertical velocity of the lowest level + omga (unused): Vertical pressure velocity + ak (in): Atmosphere hybrid a coordinate (Pa) + bk (in): Atmosphere hybrid b coordinate (dimensionless) + pfull (in): Pressure full levels + dp1 (out): Pressure thickness before dyn_core (only written + if do_sat_adjust=True) + ptop (in): The pressure level at the top of atmosphere + akap (in): Poisson constant (KAPPA) + zvir (in): Constant (Rv/Rd-1) + last_step (in): Flag for the last step of k-split remapping + consv_te (in): If True, conserve total energy + mdt (in) : Remap time step + bdt (in): Timestep + """ + # Global structure: + # pe1 is initial lagrangian edge pressures + # pe2 is final Eulerian edge pressures + + # Build remapping profiles + self._init_pe(pe, self._pe1, self._pe2, ptop) + self._moist_cv_pt_pressure( + tracers["qvapor"], + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], + q_con, + pt, + cappa, + delp, + delz, + pe, + self._pe2, + ak, + bk, + self._dp2, + ps, + self._pn2, + peln, + remap_t=True, + r_vir=zvir, + ) + self._pn2_pk_delp(self._dp2, delp, self._pe2, self._pn2, pk, akap) + + # Now that we have the pressure profiles, we can start remapping + + # Map pressure + self._map_single_pt(pt, peln, self._pn2, qmin=self._t_min) + + # Map all tracers + self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) + + # Map vertical wind + self._map_single_w(w, self._pe1, self._pe2, qs=wsd) + self._rescale_delz_1(delz, delp) + self._map_single_delz(delz, self._pe1, self._pe2) + self._rescale_delz_2(delz, self._dp2) + self._w_fix_consrv_moment( + w=w, + w2=self._w2, + dp2=self._dp2, + gz=self._gz, + w_max=self._w_max, + w_min=self._w_min, + compute_performed=self._compute_performed, + ) + + # Map horizontal winds, fluxes and courant number + self._pressures_mapu(pe, self._pe1, ak, bk, self._pe0, self._pe3) + self._pe0_ptop_xmax(self._pe0, ptop) + self._map_single_u(u, self._pe0, self._pe3) + self._map_single_u(mfy, self._pe0, self._pe3) + self._map_single_u(cy, self._pe0, self._pe3) + + self._pressures_mapv(pe, ak, bk, self._pe0, self._pe3) + self._map_single_v(v, self._pe0, self._pe3) + self._map_single_v(mfx, self._pe0, self._pe3) + self._map_single_v(cx, self._pe0, self._pe3) + + self._pe_pk_delp_peln( + pe=pe, + pk=pk, + delp=delp, + peln=peln, + pe2=self._pe2, + pk2=self._pk2, + pn2=self._pn2, + ak=ak, + bk=bk, + akap=akap, + ptop=ptop, + ) + + self._moist_cv_pkz( + tracers["qvapor"], + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], + q_con, + self._gz, + self._cvm, + pkz, + pt, + cappa, + delp, + delz, + zvir, + ) + + dtmp = 0.0 + if last_step: + if consv_te > CONSV_MIN: + self._moist_cv_te( + qvapor=tracers["qvapor"], + qliquid=tracers["qliquid"], + qrain=tracers["qrain"], + qsnow=tracers["qsnow"], + qice=tracers["qice"], + qgraupel=tracers["qgraupel"], + u=u, + v=v, + w=w, + te=self._te_2d, + pt=pt, + phis=self._phis, + delp=delp, + rsin2=self._rsin2, + cosa_s=self._cosa, + hs=hs, + delz=delz, + grav=GRAV, + ) + + self._te_zsum( + te_2d=self._te_2d, + te0_2d=self._te0_2d, + delp=delp, + pkz=pkz, + zsum1=self._zsum1, + ) + tesum = mpp_global_sum( + inputArray=self._te_2d.data * self._area_64, + communicator=self._comm, + stencil_factory=self._stencil_factory, + ) + zsum = mpp_global_sum( + inputArray=self._zsum1.data * self._area_64, + communicator=self._comm, + stencil_factory=self._stencil_factory, + ) + dtmp = tesum / (CV_AIR * zsum) + + elif consv_te < -CONSV_MIN: + raise NotImplementedError( + "Unimplemented/untested case consv(" + + str(consv_te) + + ") < -CONSV_MIN(" + + str(-CONSV_MIN) + + ")" + ) + + if self._do_sat_adjust: + fast_mp_consv = consv_te > CONSV_MIN + self._saturation_adjustment( + dp1, + tracers["qvapor"], + tracers["qliquid"], + tracers["qice"], + tracers["qrain"], + tracers["qsnow"], + tracers["qgraupel"], + q_cld, + hs, + peln, + delp, + delz, + q_con, + pt, + pkz, + cappa, + zvir, + mdt, + fast_mp_consv, + last_step, + akap, + self.kmp, + ) + + if last_step: + # on the last step, we need the regular temperature to send + # to the physics, but if we're staying in dynamics we need + # to keep it as the virtual potential temperature + self._moist_cv_last_step_stencil( + tracers["qvapor"], + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], + self._gz, + pt, + pkz, + dtmp, + zvir, + ) + else: + # converts virtual temperature back to virtual potential temperature + self._basic_adjust_divide_stencil(pkz, pt) From e90d0e7ae3de4b393313e70fa73e0bb6534e9245 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 15:16:45 -0500 Subject: [PATCH 130/252] Update dynamics with custom Remapping --- pyFV3/stencils/fv_dynamics.py | 114 ++++++++++++++++++++++++---------- 1 file changed, 80 insertions(+), 34 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 50ffba27..f2d344d0 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -6,7 +6,7 @@ import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv -from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater +from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater, is from ndsl.checkpointer import NullCheckpointer from ndsl.comm.mpi import MPI from ndsl.constants import KAPPA, NQ, X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, ZVIR @@ -25,6 +25,8 @@ from pyFV3.stencils.dyn_core import AcousticDynamics from pyFV3.stencils.neg_adj3 import AdjustNegativeTracerMixingRatio from pyFV3.stencils.remapping import LagrangianToEulerian +from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyFV3.version import IS_GEOS def pt_to_potential_density_pt( @@ -317,6 +319,16 @@ def __init__( tracers=self.tracers, checkpointer=checkpointer, ) + self._lagrangian_to_eulerian_GEOS = LagrangianToEulerian_GEOS( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config.remapping, + comm=comm, + grid_data=grid_data, + nq=NQ, + pfull=self._pfull, + tracers=self.tracers, + ) full_xyz_spec = quantity_factory.get_quantity_halo_spec( dims=[X_DIM, Y_DIM, Z_DIM], @@ -567,39 +579,73 @@ def _compute(self, state: DycoreState, timer: Timer): with timer.clock("Remapping"): self._checkpoint_remapping_in(state) - # TODO: When NQ=9, we shouldn't need to pass qcld explicitly - # since it's in self.tracers. It should not be an issue since - # we don't have self.tracers & qcld computation at the same - # time - # When NQ=8, we do need qcld passed explicitely - self._lagrangian_to_eulerian_obj( - self.tracers, - state.pt, - state.delp, - state.delz, - state.peln, - state.u, - state.v, - state.w, - self._cappa, - state.q_con, - state.qcld, - state.pkz, - state.pk, - state.pe, - state.phis, - state.ps, - self._wsd, - self._ak, - self._bk, - self._dp_initial, - self._ptop, - KAPPA, - ZVIR, - last_step, - self._conserve_total_energy, - self._timestep / self._k_split, - ) + if IS_GEOS: + self._lagrangian_to_eulerian_GEOS( + tracers=self.tracers, + pt=state.pt, + delp=state.delp, + delz=state.delz, + peln=state.peln, + u=state.u, + v=state.v, + w=state.w, + mfx=state.mfxd, + mfy=state.mfyd, + cx=state.cxd, + cy=state.cyd, + cappa=self._cappa, + q_con=state.q_con, + q_cld=state.qcld, + pkz=state.pkz, + pk=state.pk, + pe=state.pe, + hs=state.phis, + ps=state.ps, + wsd=self._wsd, + ak=self._ak, + bk=self._bk, + dp1=self._dp_initial, + ptop=self._ptop, + akap=KAPPA, + zvir=ZVIR, + last_step=last_step, + consv_te=self._conserve_total_energy, + mdt=self._timestep / self._k_split, + ) + else: + # TODO: When NQ=9, we shouldn't need to pass qcld explicitly + # since it's in self.tracers. It should not be an issue since + # we don't have self.tracers & qcld computation at the same + # time + # When NQ=8, we do need qcld passed explicitely + self._lagrangian_to_eulerian_obj( + self.tracers, + state.pt, + state.delp, + state.delz, + state.peln, + state.u, + state.v, + state.w, + self._cappa, + state.q_con, + state.qcld, + state.pkz, + state.pk, + state.pe, + state.phis, + state.ps, + self._wsd, + self._ak, + self._bk, + self._dp_initial, + self._ptop, + KAPPA, + ZVIR, + last_step, + self._conserve_total_energy, + self._timestep / self._k_split, + ) self._checkpoint_remapping_out(state) # TODO: can we pull this block out of the loop intead of # using an if-statement? From c39a2a447319a8c7f0fb5656668798402be23c1f Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 20 Jan 2025 16:31:57 -0500 Subject: [PATCH 131/252] Typo --- pyFV3/stencils/fv_dynamics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index f2d344d0..d4f11536 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -6,7 +6,7 @@ import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv -from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater, is +from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater from ndsl.checkpointer import NullCheckpointer from ndsl.comm.mpi import MPI from ndsl.constants import KAPPA, NQ, X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, ZVIR From 0d749d20933f8f540c4414c4267ddd2be6bfc566 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 21 Jan 2025 10:08:03 -0800 Subject: [PATCH 132/252] Accounting for j=je+1 calculations in map u --- pyFV3/stencils/remapping_GEOS.py | 2 +- tests/savepoint/translate/translate_remapping_GEOS.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 271bc1e2..2ad15cce 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -320,7 +320,7 @@ def __init__( grid_indexing.n_halo, 0, ), - domain=(1, grid_indexing.domain[1], 1), + domain=(1, grid_indexing.domain[1]+1, 1), ) self._pe_pk_delp_peln = stencil_factory.from_origin_domain( pe_pk_delp_peln, diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 02fe9df1..8bcbf0a3 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -876,13 +876,13 @@ def __init__( self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[0],grid_indexing.domain[1],grid_indexing.domain[2] + 1) + domain=(grid_indexing.domain[0],grid_indexing.domain[1]+1,grid_indexing.domain[2] + 1) ) self._pe0_ptop_xmax = stencil_factory.from_origin_domain( pe0_ptop_xmax, origin=(grid_indexing.n_halo+grid_indexing.domain[0],grid_indexing.n_halo,0), - domain=(1,grid_indexing.domain[1], 1) + domain=(1,grid_indexing.domain[1]+1, 1) ) self._pressures_mapv = stencil_factory.from_origin_domain( From d23f39185c54e4aa169fc749d99cdaad33177dbb Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 21 Jan 2025 14:41:00 -0800 Subject: [PATCH 133/252] Corrected issue that caused pt to incorrected read netCDF file --- .../translate/translate_remapping_GEOS.py | 203 +++++++++--------- 1 file changed, 103 insertions(+), 100 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 8bcbf0a3..e5b10ba8 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -109,7 +109,7 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "units": "No Units", }, "ak": { - "name": "pt", + "name": "ak", "dims": [X_DIM], "units": "No Units", }, @@ -159,7 +159,7 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "units": "No Units", }, "mfx_": { - "name": "pt", + "name": "mfx_", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, @@ -189,7 +189,7 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "units": "No Units", }, "te0_2d_": { - "name": "te0_2d", + "name": "te0_2d_", "dims": [X_DIM, Y_DIM], "units": "No Units", }, @@ -296,7 +296,7 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "units": "No Units", }, "mfx_": { - "name": "pt", + "name": "mfx_", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, @@ -980,6 +980,9 @@ def compute_parallel(self, inputs, communicator): interp=True, ) + print("sum(pn1) = ", sum(sum(self._pn1))) + print("sum(pn2) = ", sum(sum(self._pn2))) + tracers = { "qvapor": state_namespace.qvapor, "qliquid": state_namespace.qliquid, "qice": state_namespace.qice, @@ -1130,104 +1133,104 @@ def compute_parallel(self, inputs, communicator): interp=False, ) - self._pe_pk_delp_peln(state_namespace.pe_, - state_namespace.pk, - state_namespace.delp, - state_namespace.peln_3d, - self._pe2, - self._pk2, - self._pn2, - state_namespace.ak, - state_namespace.bk, - state_namespace.akap, - state_namespace.ptop, - ) - - self._moist_cv_pkz( - state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, - state_namespace.pkz, - state_namespace.pt, - state_namespace.cappa, - state_namespace.delp, - state_namespace.delz, - Float(state_namespace.r_vir), - ) + # self._pe_pk_delp_peln(state_namespace.pe_, + # state_namespace.pk, + # state_namespace.delp, + # state_namespace.peln_3d, + # self._pe2, + # self._pk2, + # self._pn2, + # state_namespace.ak, + # state_namespace.bk, + # state_namespace.akap, + # state_namespace.ptop, + # ) + + # self._moist_cv_pkz( + # state_namespace.qvapor, + # state_namespace.qliquid, + # state_namespace.qrain, + # state_namespace.qsnow, + # state_namespace.qice, + # state_namespace.qgraupel, + # state_namespace.pkz, + # state_namespace.pt, + # state_namespace.cappa, + # state_namespace.delp, + # state_namespace.delz, + # Float(state_namespace.r_vir), + # ) - if state_namespace.last_step and not state_namespace.do_adiabatic_init: + # if state_namespace.last_step and not state_namespace.do_adiabatic_init: - if state_namespace.consv > state_namespace.consv_min: - - self._moist_cv_te(state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, - state_namespace.u, - state_namespace.v, - state_namespace.w, - state_namespace.te_2d_, - state_namespace.pt, - self._phis, - state_namespace.delp, - state_namespace.rsin2, - state_namespace.cosa_s, - state_namespace.hs, - state_namespace.delz, - state_namespace.grav, - ) - - print("sum(te_2d): ", sum(state_namespace.te_2d_.data)) # nans here - print("sum(qliquid): ", sum(sum(state_namespace.qliquid.data))) - print("sum(qrain): ", sum(sum(state_namespace.qrain.data))) - print("sum(qsnow): ", sum(sum(state_namespace.qsnow.data))) - print("sum(qice): ", sum(sum(state_namespace.qice.data))) - print("sum(qgraupel): ", sum(sum(state_namespace.qgraupel.data))) - print("sum(u): ", sum(sum(state_namespace.u.data))) # nans here - print("sum(v): ", sum(sum(state_namespace.v.data))) - print("sum(w): ", sum(sum(state_namespace.w.data))) - print("sum(pt): ", sum(sum(state_namespace.pt.data))) - print("sum(delp): ", sum(sum(state_namespace.delp.data))) - print("sum(rsin2): ", sum(state_namespace.rsin2.data)) - print("sum(cosa_s): ", sum(state_namespace.cosa_s.data)) - print("sum(hs): ", sum(state_namespace.hs.data)) - print("sum(delz): ", sum(sum(state_namespace.delz.data))) - - # self._te_zsum(state_namespace.te_2d_, - # state_namespace.te0_2d_, - # state_namespace.delp, - # state_namespace.pkz, - # self._zsum1, - # ) - - # # Note, mpp_global_sum is currently set up for the C24 TBC setup - # print("te_2d type :", type(state_namespace.te_2d_.data[0,0])) - # print("area_64 type :", type(state_namespace.area_64_.data[0,0])) - # tesum = mpp_global_sum(state_namespace.te_2d_.data*state_namespace.area_64_.data, communicator, self.stencil_factory) - # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64_.data, communicator, self.stencil_factory) - - # dtmp = tesum / (state_namespace.cv_air.data * zsum) - # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - - - # if state_namespace.last_step and not state_namespace.adiabatic: - - # self._most_cv_pt_last_step(state_namespace.qvapor, - # state_namespace.qliquid, - # state_namespace.qrain, - # state_namespace.qsnow, - # state_namespace.qice, - # state_namespace.qgraupel, - # self._gz, - # state_namespace.pt, - # state_namespace.pkz, - # dtmp, - # state_namespace.r_vir, + # if state_namespace.consv > state_namespace.consv_min: + + # self._moist_cv_te(state_namespace.qvapor, + # state_namespace.qliquid, + # state_namespace.qrain, + # state_namespace.qsnow, + # state_namespace.qice, + # state_namespace.qgraupel, + # state_namespace.u, + # state_namespace.v, + # state_namespace.w, + # state_namespace.te_2d_, + # state_namespace.pt, + # self._phis, + # state_namespace.delp, + # state_namespace.rsin2, + # state_namespace.cosa_s, + # state_namespace.hs, + # state_namespace.delz, + # state_namespace.grav, # ) + # print("sum(te_2d): ", sum(state_namespace.te_2d_.data)) # nans here + # print("sum(qliquid): ", sum(sum(state_namespace.qliquid.data))) + # print("sum(qrain): ", sum(sum(state_namespace.qrain.data))) + # print("sum(qsnow): ", sum(sum(state_namespace.qsnow.data))) + # print("sum(qice): ", sum(sum(state_namespace.qice.data))) + # print("sum(qgraupel): ", sum(sum(state_namespace.qgraupel.data))) + # print("sum(u): ", sum(sum(state_namespace.u.data))) # nans here + # print("sum(v): ", sum(sum(state_namespace.v.data))) + # print("sum(w): ", sum(sum(state_namespace.w.data))) + # print("sum(pt): ", sum(sum(state_namespace.pt.data))) + # print("sum(delp): ", sum(sum(state_namespace.delp.data))) + # print("sum(rsin2): ", sum(state_namespace.rsin2.data)) + # print("sum(cosa_s): ", sum(state_namespace.cosa_s.data)) + # print("sum(hs): ", sum(state_namespace.hs.data)) + # print("sum(delz): ", sum(sum(state_namespace.delz.data))) + + # # self._te_zsum(state_namespace.te_2d_, + # # state_namespace.te0_2d_, + # # state_namespace.delp, + # # state_namespace.pkz, + # # self._zsum1, + # # ) + + # # # Note, mpp_global_sum is currently set up for the C24 TBC setup + # # print("te_2d type :", type(state_namespace.te_2d_.data[0,0])) + # # print("area_64 type :", type(state_namespace.area_64_.data[0,0])) + # # tesum = mpp_global_sum(state_namespace.te_2d_.data*state_namespace.area_64_.data, communicator, self.stencil_factory) + # # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64_.data, communicator, self.stencil_factory) + + # # dtmp = tesum / (state_namespace.cv_air.data * zsum) + # # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed + + + # # if state_namespace.last_step and not state_namespace.adiabatic: + + # # self._most_cv_pt_last_step(state_namespace.qvapor, + # # state_namespace.qliquid, + # # state_namespace.qrain, + # # state_namespace.qsnow, + # # state_namespace.qice, + # # state_namespace.qgraupel, + # # self._gz, + # # state_namespace.pt, + # # state_namespace.pkz, + # # dtmp, + # # state_namespace.r_vir, + # # ) + return self.outputs_from_state(state) \ No newline at end of file From 3953b29fd27c2e09248819e715acf1fff77e2160 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 22 Jan 2025 14:58:20 -0800 Subject: [PATCH 134/252] Continued verifying GEOS remapping/fv_mapz --- pyFV3/stencils/moist_cv.py | 11 +- .../translate/translate_remapping_GEOS.py | 215 +++++++++--------- 2 files changed, 114 insertions(+), 112 deletions(-) diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index a7bfba94..5a66d034 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -263,7 +263,16 @@ def cond_output( ): with computation(PARALLEL), interval(...): q_con = 0.0 - q_con = qliquid + qice + qrain + qsnow + qgraupel + if(qliquid > 0.0): + q_con = q_con + qliquid + if(qice > 0.0): + q_con = q_con + qice + if(qrain > 0.0): + q_con = q_con + qrain + if(qsnow > 0.0): + q_con = q_con + qsnow + if(qgraupel > 0.0): + q_con = q_con + qgraupel def fv_setup( diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index e5b10ba8..71a64908 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -198,11 +198,11 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "te": { - "name": "te", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, + # "te": { + # "name": "te", + # "dims": [X_DIM, Y_DIM, Z_DIM], + # "units": "No Units", + # }, } outputs = { "pt": { @@ -325,6 +325,16 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, + "te_2d_": { + "name": "te_2d_", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + # "te": { + # "name": "te", + # "dims": [X_DIM, Y_DIM, Z_DIM], + # "units": "No Units", + # }, } def __init__( self, @@ -487,7 +497,7 @@ def __init__( "jstart": grid.jsd, "jend": grid.jed, }, - "te": {} + # "te": {} } self._base.in_vars["parameters"] = [ "ptop", @@ -676,12 +686,12 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, - # "te_2d_": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # }, + "te_2d_": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, # "te": {} } @@ -980,9 +990,6 @@ def compute_parallel(self, inputs, communicator): interp=True, ) - print("sum(pn1) = ", sum(sum(self._pn1))) - print("sum(pn2) = ", sum(sum(self._pn2))) - tracers = { "qvapor": state_namespace.qvapor, "qliquid": state_namespace.qliquid, "qice": state_namespace.qice, @@ -1133,104 +1140,90 @@ def compute_parallel(self, inputs, communicator): interp=False, ) - # self._pe_pk_delp_peln(state_namespace.pe_, - # state_namespace.pk, - # state_namespace.delp, - # state_namespace.peln_3d, - # self._pe2, - # self._pk2, - # self._pn2, - # state_namespace.ak, - # state_namespace.bk, - # state_namespace.akap, - # state_namespace.ptop, - # ) - - # self._moist_cv_pkz( - # state_namespace.qvapor, - # state_namespace.qliquid, - # state_namespace.qrain, - # state_namespace.qsnow, - # state_namespace.qice, - # state_namespace.qgraupel, - # state_namespace.pkz, - # state_namespace.pt, - # state_namespace.cappa, - # state_namespace.delp, - # state_namespace.delz, - # Float(state_namespace.r_vir), - # ) + self._pe_pk_delp_peln(state_namespace.pe_, + state_namespace.pk, + state_namespace.delp, + state_namespace.peln_3d, + self._pe2, + self._pk2, + self._pn2, + state_namespace.ak, + state_namespace.bk, + state_namespace.akap, + state_namespace.ptop, + ) + + self._moist_cv_pkz( + state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.pkz, + state_namespace.pt, + state_namespace.cappa, + state_namespace.delp, + state_namespace.delz, + Float(state_namespace.r_vir), + ) - # if state_namespace.last_step and not state_namespace.do_adiabatic_init: + if state_namespace.last_step and not state_namespace.do_adiabatic_init: - # if state_namespace.consv > state_namespace.consv_min: - - # self._moist_cv_te(state_namespace.qvapor, - # state_namespace.qliquid, - # state_namespace.qrain, - # state_namespace.qsnow, - # state_namespace.qice, - # state_namespace.qgraupel, - # state_namespace.u, - # state_namespace.v, - # state_namespace.w, - # state_namespace.te_2d_, - # state_namespace.pt, - # self._phis, - # state_namespace.delp, - # state_namespace.rsin2, - # state_namespace.cosa_s, - # state_namespace.hs, - # state_namespace.delz, - # state_namespace.grav, - # ) - - # print("sum(te_2d): ", sum(state_namespace.te_2d_.data)) # nans here - # print("sum(qliquid): ", sum(sum(state_namespace.qliquid.data))) - # print("sum(qrain): ", sum(sum(state_namespace.qrain.data))) - # print("sum(qsnow): ", sum(sum(state_namespace.qsnow.data))) - # print("sum(qice): ", sum(sum(state_namespace.qice.data))) - # print("sum(qgraupel): ", sum(sum(state_namespace.qgraupel.data))) - # print("sum(u): ", sum(sum(state_namespace.u.data))) # nans here - # print("sum(v): ", sum(sum(state_namespace.v.data))) - # print("sum(w): ", sum(sum(state_namespace.w.data))) - # print("sum(pt): ", sum(sum(state_namespace.pt.data))) - # print("sum(delp): ", sum(sum(state_namespace.delp.data))) - # print("sum(rsin2): ", sum(state_namespace.rsin2.data)) - # print("sum(cosa_s): ", sum(state_namespace.cosa_s.data)) - # print("sum(hs): ", sum(state_namespace.hs.data)) - # print("sum(delz): ", sum(sum(state_namespace.delz.data))) - - # # self._te_zsum(state_namespace.te_2d_, - # # state_namespace.te0_2d_, - # # state_namespace.delp, - # # state_namespace.pkz, - # # self._zsum1, - # # ) + if state_namespace.consv > state_namespace.consv_min: + + self._moist_cv_te(state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.u, + state_namespace.v, + state_namespace.w, + state_namespace.te_2d_, + state_namespace.pt, + self._phis, + state_namespace.delp, + state_namespace.rsin2, + state_namespace.cosa_s, + state_namespace.hs, + state_namespace.delz, + state_namespace.grav, + ) + + self._te_zsum(state_namespace.te_2d_, + state_namespace.te0_2d_, + state_namespace.delp, + state_namespace.pkz, + self._zsum1, + ) - # # # Note, mpp_global_sum is currently set up for the C24 TBC setup - # # print("te_2d type :", type(state_namespace.te_2d_.data[0,0])) - # # print("area_64 type :", type(state_namespace.area_64_.data[0,0])) - # # tesum = mpp_global_sum(state_namespace.te_2d_.data*state_namespace.area_64_.data, communicator, self.stencil_factory) - # # zsum = mpp_global_sum(self._zsum1*state_namespace.area_64_.data, communicator, self.stencil_factory) - - # # dtmp = tesum / (state_namespace.cv_air.data * zsum) - # # # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - - - # # if state_namespace.last_step and not state_namespace.adiabatic: - - # # self._most_cv_pt_last_step(state_namespace.qvapor, - # # state_namespace.qliquid, - # # state_namespace.qrain, - # # state_namespace.qsnow, - # # state_namespace.qice, - # # state_namespace.qgraupel, - # # self._gz, - # # state_namespace.pt, - # # state_namespace.pkz, - # # dtmp, - # # state_namespace.r_vir, - # # ) + # Note, mpp_global_sum is currently set up for the C24 TBC setup + inputArray = state_namespace.te_2d_.data*state_namespace.area_64_.data + tesum = mpp_global_sum(inputArray[3:27,3:27], communicator, self.stencil_factory) + # print("tesum: ", tesum) + inputArray = self._zsum1*state_namespace.area_64_.data[0:-1,0:-1] + zsum = mpp_global_sum(inputArray[3:27,3:27], communicator, self.stencil_factory) + # print("zsum: ", zsum) + dtmp = tesum / (state_namespace.cv_air.data * zsum) + # print("dtmp: ", dtmp) + # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed + + + # if state_namespace.last_step and not state_namespace.adiabatic: + + # self._most_cv_pt_last_step(state_namespace.qvapor, + # state_namespace.qliquid, + # state_namespace.qrain, + # state_namespace.qsnow, + # state_namespace.qice, + # state_namespace.qgraupel, + # self._gz, + # state_namespace.pt, + # state_namespace.pkz, + # dtmp, + # state_namespace.r_vir, + # ) return self.outputs_from_state(state) \ No newline at end of file From 5567bc979d433cc0623654bc50190efbc7978934 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 22 Jan 2025 15:00:13 -0800 Subject: [PATCH 135/252] Adding fv_mapz serialization file --- fv_mapz.F90.SER | 219 +++++++++++++++++++++++++++++------------------- 1 file changed, 135 insertions(+), 84 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index cb56d14a..dc7e8e54 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -649,24 +649,9 @@ contains km, pe2, te, & is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.true.) else - !$ser verbatim if(j == js2d) then - !$ser verbatim do k = 1, km+1 - !$ser verbatim do i = is, ie - !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) - !$ser verbatim pn2_3d(i,j,k) = pn2(i,k) - !$ser verbatim enddo - !$ser verbatim enddo - !$ser savepoint MapScalar_2d-In - !$ser verbatim mode=1 - !$ser data pe1_=pn1_3d pt=pt q1=te mode=mode j_2d=js2d q_min=-1.e25 - !$ser verbatim endif call map_scalar(km, pn1, te, gz, & km, pn2, te, & is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), te_min) - !$ser verbatim if(j == js2d) then - !$ser savepoint MapScalar_2d-Out - !$ser data q1=te - !$ser verbatim endif endif else !---------------------------------- @@ -1371,72 +1356,72 @@ contains 2000 continue -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js, je -!$ser verbatim do i = is, ie -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js-1, je+1 -!$ser verbatim do i = is-1, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser savepoint Remapping_GEOS-Out -!$ser data pe_=pe_3d -!$ser data qvapor=q(:,:,:,sphum) -!$ser data qliquid=q(:,:,:,liq_wat) -!$ser data qice=q(:,:,:,ice_wat) -!$ser data qrain=q(:,:,:,rainwat) -!$ser data qsnow=q(:,:,:,snowwat) -!$ser data qgraupel=q(:,:,:,graupel) -!$ser data qcld=q(:,:,:,cld_amt) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data peln_3d=peln_3d -!$ser data r_vir=r_vir -!$ser data cvm=cvm -!$ser data pk=pk -!$ser data pkz=pkz -!$ser data akap=akap -!$ser data t_min=qmin -!$ser data kord_wz=kord_wz -!$ser data ws_=ws -!$ser data w=w -!$ser data u=u -!$ser data v=v -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser data w_max=w_max_ -!$ser data w_min=w_min_ -!$ser data kord_mt=kord_mt -!$ser data cosa_s=cosa_s -!$ser data rsin2=rsin2 -!$ser data hs=hs -!$ser data grav=grav_ -!$ser data te_2d_=te_2d_f32 -!$ser data te0_2d_=te0_2d_f32 -!$ser data area_64_=area_64_ -!$ser data te=te +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do j = js, je +!!$ser verbatim do i = is, ie +!!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim do k = 1, km+1 +!!$ser verbatim do j = js-1, je+1 +!!$ser verbatim do i = is-1, ie+1 +!!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser verbatim enddo +!!$ser savepoint Remapping_GEOS-Out +!!$ser data pe_=pe_3d +!!$ser data qvapor=q(:,:,:,sphum) +!!$ser data qliquid=q(:,:,:,liq_wat) +!!$ser data qice=q(:,:,:,ice_wat) +!!$ser data qrain=q(:,:,:,rainwat) +!!$ser data qsnow=q(:,:,:,snowwat) +!!$ser data qgraupel=q(:,:,:,graupel) +!!$ser data qcld=q(:,:,:,cld_amt) +!!$ser data qo3mr=q(:,:,:,8) +!!$ser data qsgs_tke=q(:,:,:,9) +!!$ser data q_con=q_con +!!$ser data pt=pt +!!$ser data cappa=cappa +!!$ser data delp=delp +!!$ser data delz=delz +!!$ser data ps=ps +!!$ser data peln_3d=peln_3d +!!$ser data r_vir=r_vir +!!$ser data cvm=cvm +!!$ser data pk=pk +!!$ser data pkz=pkz +!!$ser data akap=akap +!!$ser data t_min=qmin +!!$ser data kord_wz=kord_wz +!!$ser data ws_=ws +!!$ser data w=w +!!$ser data u=u +!!$ser data v=v +!!$ser verbatim if (present(mfy)) then +!!$ser data mfy=mfy +!!$ser verbatim endif +!!$ser verbatim if (present(cy)) then +!!$ser data cy=cy +!!$ser verbatim endif +!!$ser verbatim if (present(mfx)) then +!!$ser data mfx_=mfx +!!$ser verbatim endif +!!$ser verbatim if (present(cx)) then +!!$ser data cx_=cx +!!$ser verbatim endif +!!$ser data w_max=w_max_ +!!$ser data w_min=w_min_ +!!$ser data kord_mt=kord_mt +!!$ser data cosa_s=cosa_s +!!$ser data rsin2=rsin2 +!!$ser data hs=hs +!!$ser data grav=grav_ +!!$ser data te_2d_=te_2d_f32 +!!$ser data te0_2d_=te0_2d_f32 +!!$ser data area_64_=area_64_ +!!$ser data te=te ! Do total energy conservation and fast saturation adjustment as requested ! and fill new PT (Theta_V) for next k_split step or export dry T @@ -1549,9 +1534,7 @@ contains #endif enddo ! k-loop !$ser verbatim if(j == js) then - !$ser verbatim do i = is, ie - !$ser verbatim te_2d_f32(i,:) = te_2d(i,:) - !$ser verbatim enddo + !$ser verbatim te_2d_f32 = real(te_2d) !$ser savepoint MoistCVPlusTe_2d-Out !$ser data te_2d_=te_2d_f32 !$ser verbatim endif @@ -1670,6 +1653,74 @@ contains endif ! end consv check endif ! end last_step check +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js, je +!$ser verbatim do i = is, ie +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js-1, je+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim te_2d_f32 = real(te_2d) +!$ser verbatim if(last_step .and. (.not.do_adiabatic_init)) then +!$ser verbatim print*, 'tesum = ', tesum +!$ser verbatim print*, 'zsum = ', zsum +!$ser verbatim print*, 'dtmp = ', dtmp +!$ser verbatim endif +!$ser savepoint Remapping_GEOS-Out +!$ser data pe_=pe_3d +!$ser data qvapor=q(:,:,:,sphum) +!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qice=q(:,:,:,ice_wat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qgraupel=q(:,:,:,graupel) +!$ser data qcld=q(:,:,:,cld_amt) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data peln_3d=peln_3d +!$ser data r_vir=r_vir +!$ser data cvm=cvm +!$ser data pk=pk +!$ser data pkz=pkz +!$ser data akap=akap +!$ser data w=w +!$ser data u=u +!$ser data v=v +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser data kord_mt=kord_mt +!$ser data cosa_s=cosa_s +!$ser data rsin2=rsin2 +!$ser data hs=hs +!$ser data grav=grav_ +!$ser data te_2d_=te_2d_f32 +!$ser data te0_2d_=te0_2d_f32 +!!$ser data area_64_=area_64_ +!!$ser data te=te + ! Note: pt at this stage is T_v if ( do_sat_adj .and. nwat>=6 ) then print*,'CODE NOT TESTED HERE 15' From a53c3febd35d82bab0ff837e006b9e88526ab4d9 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 22 Jan 2025 16:16:25 -0800 Subject: [PATCH 136/252] Removed extraneous variables from GEOS remapping translate test and added fill condensate stencil to translate test --- pyFV3/stencils/moist_cv.py | 2 - .../translate/translate_remapping_GEOS.py | 149 +++++------------- 2 files changed, 41 insertions(+), 110 deletions(-) diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index 5a66d034..f6ddc49c 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -92,7 +92,6 @@ def moist_pt_last_step( qsnow: FloatField, qice: FloatField, qgraupel: FloatField, - gz: FloatField, pt: FloatField, pkz: FloatField, dtmp: Float, @@ -106,7 +105,6 @@ def moist_pt_last_step( qsnow (in): qice (in): qgraupel (in): - gz (out): pt (inout): pkz (in): dtmp (in): diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 71a64908..d2d5a4d5 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -178,11 +178,6 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "te_2d_": { - "name": "te_2d_", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, "hs": { "name": "hs", "dims": [X_DIM, Y_DIM], @@ -198,11 +193,6 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - # "te": { - # "name": "te", - # "dims": [X_DIM, Y_DIM, Z_DIM], - # "units": "No Units", - # }, } outputs = { "pt": { @@ -325,16 +315,11 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "te_2d_": { - "name": "te_2d_", - "dims": [X_DIM, Y_DIM], + "q_con": { + "name": "q_con", + "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - # "te": { - # "name": "te", - # "dims": [X_DIM, Y_DIM, Z_DIM], - # "units": "No Units", - # }, } def __init__( self, @@ -473,12 +458,6 @@ def __init__( "jstart": grid.jsd, "jend": grid.jed, }, - "te_2d_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - }, "hs": { "istart": grid.isd, "iend": grid.ied, @@ -497,7 +476,6 @@ def __init__( "jstart": grid.jsd, "jend": grid.jed, }, - # "te": {} } self._base.in_vars["parameters"] = [ "ptop", @@ -518,69 +496,11 @@ def __init__( "adiabatic", ] self._base.out_vars = { - # "pe1_": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "pe2_": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "pe0_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz+1 - # }, - # "pe3_": { - # "istart": grid.is_, - # "iend": grid.ie+1, - # "jstart": grid.js, - # "jend": grid.je+1, - # "kend": grid.npz+1 - # }, "pt": {}, "cappa": {}, - # "q_con": {}, + "q_con": {}, "delp": {}, "delz": {}, - # "ps": {}, - # "dp2_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz-1, - # }, - # "pn1_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "pn2_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - # "pk2_3d": { - # "istart": grid.is_, - # "iend": grid.ie, - # "jstart": grid.js, - # "jend": grid.je, - # "kend": grid.npz + 1, - # }, - "qvapor": { "kend": grid.npz-1, }, @@ -686,13 +606,6 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, - "te_2d_": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - }, - # "te": {} } self.stencil_factory = stencil_factory @@ -835,6 +748,13 @@ def __init__( ), dtype=Float, ) + self._te_2d = self.quantity_factory._numpy.zeros( + ( + grid.nid, + grid.njd, + ), dtype=Float, + ) + self._init_pe = stencil_factory.from_origin_domain( init_pe, @@ -936,6 +856,12 @@ def __init__( domain=(grid.nic, grid.njc, grid.npz), ) + self._fill_cond = stencil_factory.from_origin_domain( + moist_cv.cond_output, + origin=grid.compute_origin(), + domain=(grid.nic, grid.njc, grid.npz), + ) + def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") @@ -1181,7 +1107,7 @@ def compute_parallel(self, inputs, communicator): state_namespace.u, state_namespace.v, state_namespace.w, - state_namespace.te_2d_, + self._te_2d, state_namespace.pt, self._phis, state_namespace.delp, @@ -1192,7 +1118,7 @@ def compute_parallel(self, inputs, communicator): state_namespace.grav, ) - self._te_zsum(state_namespace.te_2d_, + self._te_zsum(self._te_2d, state_namespace.te0_2d_, state_namespace.delp, state_namespace.pkz, @@ -1200,7 +1126,7 @@ def compute_parallel(self, inputs, communicator): ) # Note, mpp_global_sum is currently set up for the C24 TBC setup - inputArray = state_namespace.te_2d_.data*state_namespace.area_64_.data + inputArray = self._te_2d.data*state_namespace.area_64_.data[0:-1,0:-1] tesum = mpp_global_sum(inputArray[3:27,3:27], communicator, self.stencil_factory) # print("tesum: ", tesum) inputArray = self._zsum1*state_namespace.area_64_.data[0:-1,0:-1] @@ -1211,19 +1137,26 @@ def compute_parallel(self, inputs, communicator): # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - # if state_namespace.last_step and not state_namespace.adiabatic: - - # self._most_cv_pt_last_step(state_namespace.qvapor, - # state_namespace.qliquid, - # state_namespace.qrain, - # state_namespace.qsnow, - # state_namespace.qice, - # state_namespace.qgraupel, - # self._gz, - # state_namespace.pt, - # state_namespace.pkz, - # dtmp, - # state_namespace.r_vir, - # ) + if state_namespace.last_step and not state_namespace.adiabatic: + + self._most_cv_pt_last_step(state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.pt, + state_namespace.pkz, + Float(dtmp), + state_namespace.r_vir, + ) + + self._fill_cond(state_namespace.q_con, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + ) return self.outputs_from_state(state) \ No newline at end of file From e61b6bbcdf0dce915a3a7206b72919d2c228b126 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 22 Jan 2025 16:17:03 -0800 Subject: [PATCH 137/252] Added fv_mapz serialization file --- fv_mapz.F90.SER | 203 ++++++++++++++++-------------------------------- 1 file changed, 65 insertions(+), 138 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index dc7e8e54..455c0f07 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -402,7 +402,6 @@ contains !$ser verbatim rsin2=gridstruct%rsin2 !$ser verbatim cosa_s=gridstruct%cosa_s !$ser verbatim grav_=grav -!$ser verbatim te_2d_f32 = real(te_2d) !$ser verbatim area_64_ = real(gridstruct%area_64) !$ser verbatim te0_2d_f32 = real(te0_2d) !$ser savepoint Remapping_GEOS-In @@ -456,10 +455,8 @@ contains !$ser data rsin2=rsin2 !$ser data hs=hs !$ser data grav=grav_ -!$ser data te_2d_=te_2d_f32 !$ser data te0_2d_=te0_2d_f32 !$ser data area_64_=area_64_ -!$ser data te=te !$ser data last_step=last_step !$ser data do_adiabatic_init=do_adiabatic_init !$ser data consv=consv @@ -1356,73 +1353,6 @@ contains 2000 continue -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do j = js, je -!!$ser verbatim do i = is, ie -!!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim do k = 1, km+1 -!!$ser verbatim do j = js-1, je+1 -!!$ser verbatim do i = is-1, ie+1 -!!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser verbatim enddo -!!$ser savepoint Remapping_GEOS-Out -!!$ser data pe_=pe_3d -!!$ser data qvapor=q(:,:,:,sphum) -!!$ser data qliquid=q(:,:,:,liq_wat) -!!$ser data qice=q(:,:,:,ice_wat) -!!$ser data qrain=q(:,:,:,rainwat) -!!$ser data qsnow=q(:,:,:,snowwat) -!!$ser data qgraupel=q(:,:,:,graupel) -!!$ser data qcld=q(:,:,:,cld_amt) -!!$ser data qo3mr=q(:,:,:,8) -!!$ser data qsgs_tke=q(:,:,:,9) -!!$ser data q_con=q_con -!!$ser data pt=pt -!!$ser data cappa=cappa -!!$ser data delp=delp -!!$ser data delz=delz -!!$ser data ps=ps -!!$ser data peln_3d=peln_3d -!!$ser data r_vir=r_vir -!!$ser data cvm=cvm -!!$ser data pk=pk -!!$ser data pkz=pkz -!!$ser data akap=akap -!!$ser data t_min=qmin -!!$ser data kord_wz=kord_wz -!!$ser data ws_=ws -!!$ser data w=w -!!$ser data u=u -!!$ser data v=v -!!$ser verbatim if (present(mfy)) then -!!$ser data mfy=mfy -!!$ser verbatim endif -!!$ser verbatim if (present(cy)) then -!!$ser data cy=cy -!!$ser verbatim endif -!!$ser verbatim if (present(mfx)) then -!!$ser data mfx_=mfx -!!$ser verbatim endif -!!$ser verbatim if (present(cx)) then -!!$ser data cx_=cx -!!$ser verbatim endif -!!$ser data w_max=w_max_ -!!$ser data w_min=w_min_ -!!$ser data kord_mt=kord_mt -!!$ser data cosa_s=cosa_s -!!$ser data rsin2=rsin2 -!!$ser data hs=hs -!!$ser data grav=grav_ -!!$ser data te_2d_=te_2d_f32 -!!$ser data te0_2d_=te0_2d_f32 -!!$ser data area_64_=area_64_ -!!$ser data te=te - ! Do total energy conservation and fast saturation adjustment as requested ! and fill new PT (Theta_V) for next k_split step or export dry T @@ -1653,74 +1583,6 @@ contains endif ! end consv check endif ! end last_step check -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js, je -!$ser verbatim do i = is, ie -!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js-1, je+1 -!$ser verbatim do i = is-1, ie+1 -!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim enddo -!$ser verbatim te_2d_f32 = real(te_2d) -!$ser verbatim if(last_step .and. (.not.do_adiabatic_init)) then -!$ser verbatim print*, 'tesum = ', tesum -!$ser verbatim print*, 'zsum = ', zsum -!$ser verbatim print*, 'dtmp = ', dtmp -!$ser verbatim endif -!$ser savepoint Remapping_GEOS-Out -!$ser data pe_=pe_3d -!$ser data qvapor=q(:,:,:,sphum) -!$ser data qliquid=q(:,:,:,liq_wat) -!$ser data qice=q(:,:,:,ice_wat) -!$ser data qrain=q(:,:,:,rainwat) -!$ser data qsnow=q(:,:,:,snowwat) -!$ser data qgraupel=q(:,:,:,graupel) -!$ser data qcld=q(:,:,:,cld_amt) -!$ser data qo3mr=q(:,:,:,8) -!$ser data qsgs_tke=q(:,:,:,9) -!$ser data q_con=q_con -!$ser data pt=pt -!$ser data cappa=cappa -!$ser data delp=delp -!$ser data delz=delz -!$ser data ps=ps -!$ser data peln_3d=peln_3d -!$ser data r_vir=r_vir -!$ser data cvm=cvm -!$ser data pk=pk -!$ser data pkz=pkz -!$ser data akap=akap -!$ser data w=w -!$ser data u=u -!$ser data v=v -!$ser verbatim if (present(mfy)) then -!$ser data mfy=mfy -!$ser verbatim endif -!$ser verbatim if (present(cy)) then -!$ser data cy=cy -!$ser verbatim endif -!$ser verbatim if (present(mfx)) then -!$ser data mfx_=mfx -!$ser verbatim endif -!$ser verbatim if (present(cx)) then -!$ser data cx_=cx -!$ser verbatim endif -!$ser data kord_mt=kord_mt -!$ser data cosa_s=cosa_s -!$ser data rsin2=rsin2 -!$ser data hs=hs -!$ser data grav=grav_ -!$ser data te_2d_=te_2d_f32 -!$ser data te0_2d_=te0_2d_f32 -!!$ser data area_64_=area_64_ -!!$ser data te=te - ! Note: pt at this stage is T_v if ( do_sat_adj .and. nwat>=6 ) then print*,'CODE NOT TESTED HERE 15' @@ -1863,6 +1725,71 @@ contains endif !$OMP end parallel +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js, je +!$ser verbatim do i = is, ie +!$ser verbatim peln_3d(i,j,k) = peln(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim do k = 1, km+1 +!$ser verbatim do j = js-1, je+1 +!$ser verbatim do i = is-1, ie+1 +!$ser verbatim pe_3d(i,j,k) = pe(i,k,j) +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim enddo +!$ser verbatim te_2d_f32 = real(te_2d) +!$ser verbatim if(last_step .and. (.not.do_adiabatic_init)) then +!$ser verbatim print*, 'tesum = ', tesum +!$ser verbatim print*, 'zsum = ', zsum +!$ser verbatim print*, 'dtmp = ', dtmp +!$ser verbatim endif +!$ser savepoint Remapping_GEOS-Out +!$ser data pe_=pe_3d +!$ser data qvapor=q(:,:,:,sphum) +!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qice=q(:,:,:,ice_wat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qgraupel=q(:,:,:,graupel) +!$ser data qcld=q(:,:,:,cld_amt) +!$ser data qo3mr=q(:,:,:,8) +!$ser data qsgs_tke=q(:,:,:,9) +!$ser data q_con=q_con +!$ser data pt=pt +!$ser data cappa=cappa +!$ser data delp=delp +!$ser data delz=delz +!$ser data ps=ps +!$ser data peln_3d=peln_3d +!$ser data r_vir=r_vir +!$ser data cvm=cvm +!$ser data pk=pk +!$ser data pkz=pkz +!$ser data akap=akap +!$ser data w=w +!$ser data u=u +!$ser data v=v +!$ser verbatim if (present(mfy)) then +!$ser data mfy=mfy +!$ser verbatim endif +!$ser verbatim if (present(cy)) then +!$ser data cy=cy +!$ser verbatim endif +!$ser verbatim if (present(mfx)) then +!$ser data mfx_=mfx +!$ser verbatim endif +!$ser verbatim if (present(cx)) then +!$ser data cx_=cx +!$ser verbatim endif +!$ser data kord_mt=kord_mt +!$ser data cosa_s=cosa_s +!$ser data rsin2=rsin2 +!$ser data hs=hs +!$ser data grav=grav_ +!$ser data te0_2d_=te0_2d_f32 + !$ser savepoint GetMPIProp-In !$ser data delz=delz !$ser savepoint GetMPIProp-Out From 6c138790a1795b171c28aed291a612d02f0d5c82 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 24 Jan 2025 08:31:56 -0800 Subject: [PATCH 138/252] Small correction in mpp_global_sum.py --- pyFV3/stencils/mpp_global_sum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py index 61cffe78..f6af8b98 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/stencils/mpp_global_sum.py @@ -69,7 +69,7 @@ def increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): def carry_overflow(int_sum, prec, I_prec, prec_error): - for i in range(len(int_sum) - 1, 1, -1): + for i in range(len(int_sum) - 1, 0, -1): if abs(int_sum[i]) > prec: num_carry = int(int_sum[i] * I_prec) int_sum[i] = int_sum[i] - num_carry * prec From 8e59317d823938674bae3d9d5cd3974f43f77a9f Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 24 Jan 2025 13:13:33 -0500 Subject: [PATCH 139/252] `remapper`: align with modified dynamics `moist_cv`: lint --- pyFV3/stencils/moist_cv.py | 12 +- pyFV3/stencils/remapping_GEOS.py | 234 +++++++++++++++++-------------- 2 files changed, 131 insertions(+), 115 deletions(-) diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index f6ddc49c..d8dbeaf4 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -213,7 +213,7 @@ def moist_te( with computation(BACKWARD), interval(0, -1): phis = phis[0, 0, 1] - grav * delz with computation(FORWARD), interval(0, -1): - cvm, gz = moist_cv_nwat6_fn(qvapor, qliquid, qrain, qsnow, qice, qgraupel) + cvm, _gz = moist_cv_nwat6_fn(qvapor, qliquid, qrain, qsnow, qice, qgraupel) te = te + delp * ( cvm * pt @@ -261,15 +261,15 @@ def cond_output( ): with computation(PARALLEL), interval(...): q_con = 0.0 - if(qliquid > 0.0): + if qliquid > 0.0: q_con = q_con + qliquid - if(qice > 0.0): + if qice > 0.0: q_con = q_con + qice - if(qrain > 0.0): + if qrain > 0.0: q_con = q_con + qrain - if(qsnow > 0.0): + if qsnow > 0.0: q_con = q_con + qsnow - if(qgraupel > 0.0): + if qgraupel > 0.0: q_con = q_con + qgraupel diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 2ad15cce..3bd30fcb 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -38,6 +38,7 @@ CONSV_MIN, ) from pyFV3.stencils.mpp_global_sum import mpp_global_sum +from pyFV3.tracers import Tracers class LagrangianToEulerian_GEOS: @@ -56,7 +57,8 @@ def __init__( grid_data: GridData, nq, pfull, - tracers: Dict[str, Quantity], + tracers: Tracers, + adiabatic: bool, ): orchestrate( obj=self, @@ -68,23 +70,25 @@ def __init__( grid_indexing = stencil_factory.grid_indexing # Configuration - if config.kord_tm >= 0: - raise NotImplementedError("map ppm, untested mode where kord_tm >= 0") hydrostatic = config.hydrostatic if hydrostatic: raise NotImplementedError("Hydrostatic is not implemented") + if adiabatic: + raise NotImplementedError("Adiabatic is not implemented") + self._t_min = Float(184.0) self._nq = nq self._w_max = Float(90.0) self._w_min = Float(-60.0) self._area_64 = grid_data.area_64 - self._cosa = grid_data.cosa_s + self._cosa_s = grid_data.cosa_s self._rsin2 = grid_data.rsin2 self._kord_tm = abs(config.kord_tm) self._kord_wz = config.kord_wz self._kord_mt = config.kord_mt self._do_sat_adjust = config.do_sat_adj + self._adiabatic = adiabatic self.kmp = grid_indexing.domain[2] - 1 for k in range(pfull.shape[0]): if pfull.view[k] > 10.0e2: @@ -93,12 +97,6 @@ def __init__( # do_omega = hydrostatic and last_step # TODO pull into inputs # Quantities - self._domain_jextra = ( - grid_indexing.domain[0], - grid_indexing.domain[1] + 1, - grid_indexing.domain[2] + 1, - ) - self._pe1 = quantity_factory.zeros( [X_DIM, Y_DIM, Z_INTERFACE_DIM], units="Pa", @@ -119,6 +117,11 @@ def __init__( units="Pa", dtype=Float, ) + self._pn1 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", + dtype=Float, + ) self._pn2 = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="Pa", @@ -136,7 +139,7 @@ def __init__( ) self._gz = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [X_DIM, Y_DIM], units="m^2 s^-2", dtype=Float, ) @@ -150,7 +153,7 @@ def __init__( units="mask", dtype=bool, ) - self._w2 = quantity_factory._numpy.zeros( + self._w2 = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="temp W", dtype=Float, @@ -167,12 +170,6 @@ def __init__( dtype=Float, ) - self._te0_2d = quantity_factory.zeros( - [X_DIM, Y_DIM], - units="Pa", - dtype=Float, - ) - self._zsum1 = quantity_factory.zeros( [X_DIM, Y_DIM], units="Pa", @@ -187,12 +184,13 @@ def __init__( # Stencils self._init_pe = stencil_factory.from_origin_domain( - init_pe, origin=grid_indexing.origin_compute(), domain=self._domain_jextra + init_pe, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(add=(0, 1, 1)), ) self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( moist_cv_pt_pressure, - # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, externals={"hydrostatic": hydrostatic}, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(0, 0, 1)), @@ -200,8 +198,8 @@ def __init__( self._pn2_pk_delp = stencil_factory.from_origin_domain( pn2_pk_delp, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), + origin=grid_indexing.origin_compute(add=(0, 0, 1)), + domain=grid_indexing.domain_compute(add=(0, 0, -1)), ) self._map_single_pt = MapSingle( @@ -216,9 +214,9 @@ def __init__( stencil_factory, quantity_factory, abs(config.kord_tr), - nq, fill=config.fill, tracers=tracers, + exclude_tracers=[], ) self._map_single_w = MapSingle( @@ -246,7 +244,7 @@ def __init__( self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=self._domain_jextra, + domain=grid_indexing.domain_compute(add=(0, 1, 1)), ) self._map_single_u = MapSingle( @@ -260,11 +258,7 @@ def __init__( self._pressures_mapv = stencil_factory.from_origin_domain( pressures_mapv, origin=grid_indexing.origin_compute(), - domain=( - grid_indexing.domain[0] + 1, - grid_indexing.domain[1], - grid_indexing.domain[2] + 1, - ), + domain=grid_indexing.domain_compute(add=(1, 0, 1)), ) self._map_single_v = MapSingle( @@ -289,7 +283,13 @@ def __init__( ), ) - self._basic_adjust_divide_stencil = stencil_factory.from_origin_domain( + self._fill_cond = stencil_factory.from_origin_domain( + moist_cv.cond_output, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + + self._adjust_divide = stencil_factory.from_origin_domain( adjust_divide_stencil, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), @@ -320,7 +320,7 @@ def __init__( grid_indexing.n_halo, 0, ), - domain=(1, grid_indexing.domain[1]+1, 1), + domain=(1, grid_indexing.domain[1] + 1, 1), ) self._pe_pk_delp_peln = stencil_factory.from_origin_domain( pe_pk_delp_peln, @@ -341,7 +341,7 @@ def __init__( def __call__( self, - tracers: Dict[str, Quantity], + tracers: Tracers, pt: FloatField, # type: ignore delp: FloatField, # type: ignore delz: FloatField, # type: ignore @@ -355,11 +355,11 @@ def __call__( cy: FloatField, # type: ignore cappa: FloatField, # type: ignore q_con: FloatField, # type: ignore - q_cld: FloatField, # type: ignore pkz: FloatField, # type: ignore pk: FloatField, # type: ignore pe: FloatField, # type: ignore hs: FloatFieldIJ, # type: ignore + te0_2d: FloatFieldIJ, # type: ignore ps: FloatFieldIJ, # type: ignore wsd: FloatFieldIJ, # type: ignore ak: FloatFieldK, # type: ignore @@ -389,12 +389,11 @@ def __call__( va (inout): A-grid y-velocity cappa (inout): Power to raise pressure to q_con (out): Total condensate mixing ratio - q_cld (out): Cloud fraction pkz (in): Layer mean pressure raised to the power of Kappa pk (out): Interface pressure raised to power of kappa, final acoustic value pe (in): Pressure at layer edges hs (in): Surface geopotential - te0_2d (unused): Atmosphere total energy in columns + te0_2d (inout): Atmosphere total energy in columns ps (out): Surface pressure wsd (in): Vertical velocity of the lowest level omga (unused): Vertical pressure velocity @@ -418,40 +417,52 @@ def __call__( # Build remapping profiles self._init_pe(pe, self._pe1, self._pe2, ptop) self._moist_cv_pt_pressure( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - q_con, - pt, - cappa, - delp, - delz, - pe, - self._pe2, - ak, - bk, - self._dp2, - ps, - self._pn2, - peln, + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], + q_con=q_con, + pt=pt, + cappa=cappa, + delp=delp, + delz=delz, + pe=pe, + pe2=self._pe2, + ak=ak, + bk=bk, + dp2=self._dp2, + ps=ps, + pn1=self._pn1, + pn2=self._pn2, + peln=peln, remap_t=True, r_vir=zvir, ) - self._pn2_pk_delp(self._dp2, delp, self._pe2, self._pn2, pk, akap) + self._pn2_pk_delp( + pe2=self._pe2, + pn2=self._pn2, + pk=self._pk2, + akap=akap, + ) # Now that we have the pressure profiles, we can start remapping # Map pressure - self._map_single_pt(pt, peln, self._pn2, qmin=self._t_min) + self._map_single_pt( + pt, + self._pn1, + self._pn2, + qmin=self._t_min, + interp=True, + ) # Map all tracers self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) # Map vertical wind - self._map_single_w(w, self._pe1, self._pe2, qs=wsd) + self._map_single_w(w, self._pe1, self._pe2, qs=wsd, interp=False) self._rescale_delz_1(delz, delp) self._map_single_delz(delz, self._pe1, self._pe2) self._rescale_delz_2(delz, self._dp2) @@ -466,16 +477,16 @@ def __call__( ) # Map horizontal winds, fluxes and courant number - self._pressures_mapu(pe, self._pe1, ak, bk, self._pe0, self._pe3) + self._pressures_mapu(pe, ak, bk, self._pe0, self._pe3, ptop) self._pe0_ptop_xmax(self._pe0, ptop) - self._map_single_u(u, self._pe0, self._pe3) - self._map_single_u(mfy, self._pe0, self._pe3) - self._map_single_u(cy, self._pe0, self._pe3) + self._map_single_u(u, self._pe0, self._pe3, interp=False) + self._map_single_u(mfy, self._pe0, self._pe3, interp=False) + self._map_single_u(cy, self._pe0, self._pe3, interp=False) self._pressures_mapv(pe, ak, bk, self._pe0, self._pe3) - self._map_single_v(v, self._pe0, self._pe3) - self._map_single_v(mfx, self._pe0, self._pe3) - self._map_single_v(cx, self._pe0, self._pe3) + self._map_single_v(v, self._pe0, self._pe3, interp=False) + self._map_single_v(mfx, self._pe0, self._pe3, interp=False) + self._map_single_v(cx, self._pe0, self._pe3, interp=False) self._pe_pk_delp_peln( pe=pe, @@ -492,33 +503,30 @@ def __call__( ) self._moist_cv_pkz( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - q_con, - self._gz, - self._cvm, - pkz, - pt, - cappa, - delp, - delz, - zvir, + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], + pkz=pkz, + pt=pt, + cappa=cappa, + delp=delp, + delz=delz, + r_vir=zvir, ) dtmp = 0.0 if last_step: if consv_te > CONSV_MIN: self._moist_cv_te( - qvapor=tracers["qvapor"], - qliquid=tracers["qliquid"], - qrain=tracers["qrain"], - qsnow=tracers["qsnow"], - qice=tracers["qice"], - qgraupel=tracers["qgraupel"], + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], u=u, v=v, w=w, @@ -527,7 +535,7 @@ def __call__( phis=self._phis, delp=delp, rsin2=self._rsin2, - cosa_s=self._cosa, + cosa_s=self._cosa_s, hs=hs, delz=delz, grav=GRAV, @@ -535,18 +543,19 @@ def __call__( self._te_zsum( te_2d=self._te_2d, - te0_2d=self._te0_2d, + te0_2d=te0_2d, delp=delp, pkz=pkz, zsum1=self._zsum1, ) + tesum = mpp_global_sum( - inputArray=self._te_2d.data * self._area_64, + inputArray=self._te_2d.view[:] * self._area_64.view[:], communicator=self._comm, stencil_factory=self._stencil_factory, ) zsum = mpp_global_sum( - inputArray=self._zsum1.data * self._area_64, + inputArray=self._zsum1.view[:] * self._area_64.view[:], communicator=self._comm, stencil_factory=self._stencil_factory, ) @@ -565,13 +574,13 @@ def __call__( fast_mp_consv = consv_te > CONSV_MIN self._saturation_adjustment( dp1, - tracers["qvapor"], - tracers["qliquid"], - tracers["qice"], - tracers["qrain"], - tracers["qsnow"], - tracers["qgraupel"], - q_cld, + tracers["vapor"], + tracers["liquid"], + tracers["ice"], + tracers["rain"], + tracers["snow"], + tracers["graupel"], + tracers["cloud"], hs, peln, delp, @@ -588,23 +597,30 @@ def __call__( self.kmp, ) - if last_step: + if last_step and not self._adiabatic: # on the last step, we need the regular temperature to send # to the physics, but if we're staying in dynamics we need # to keep it as the virtual potential temperature self._moist_cv_last_step_stencil( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - self._gz, - pt, - pkz, - dtmp, - zvir, + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], + pt=pt, + pkz=pkz, + dtmp=Float(dtmp), + r_vir=zvir, + ) + self._fill_cond( + q_con=q_con, + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], ) else: # converts virtual temperature back to virtual potential temperature - self._basic_adjust_divide_stencil(pkz, pt) + self._adjust_divide(pkz, pt) From 2a774acea200640a306351eac6d1f049140a8dff Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 24 Jan 2025 13:23:06 -0500 Subject: [PATCH 140/252] Revert "Update dynamics with custom Remapping" This reverts commit e90d0e7ae3de4b393313e70fa73e0bb6534e9245. --- pyFV3/stencils/fv_dynamics.py | 112 ++++++++++------------------------ 1 file changed, 33 insertions(+), 79 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index d4f11536..50ffba27 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -25,8 +25,6 @@ from pyFV3.stencils.dyn_core import AcousticDynamics from pyFV3.stencils.neg_adj3 import AdjustNegativeTracerMixingRatio from pyFV3.stencils.remapping import LagrangianToEulerian -from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS -from pyFV3.version import IS_GEOS def pt_to_potential_density_pt( @@ -319,16 +317,6 @@ def __init__( tracers=self.tracers, checkpointer=checkpointer, ) - self._lagrangian_to_eulerian_GEOS = LagrangianToEulerian_GEOS( - stencil_factory=stencil_factory, - quantity_factory=quantity_factory, - config=config.remapping, - comm=comm, - grid_data=grid_data, - nq=NQ, - pfull=self._pfull, - tracers=self.tracers, - ) full_xyz_spec = quantity_factory.get_quantity_halo_spec( dims=[X_DIM, Y_DIM, Z_DIM], @@ -579,73 +567,39 @@ def _compute(self, state: DycoreState, timer: Timer): with timer.clock("Remapping"): self._checkpoint_remapping_in(state) - if IS_GEOS: - self._lagrangian_to_eulerian_GEOS( - tracers=self.tracers, - pt=state.pt, - delp=state.delp, - delz=state.delz, - peln=state.peln, - u=state.u, - v=state.v, - w=state.w, - mfx=state.mfxd, - mfy=state.mfyd, - cx=state.cxd, - cy=state.cyd, - cappa=self._cappa, - q_con=state.q_con, - q_cld=state.qcld, - pkz=state.pkz, - pk=state.pk, - pe=state.pe, - hs=state.phis, - ps=state.ps, - wsd=self._wsd, - ak=self._ak, - bk=self._bk, - dp1=self._dp_initial, - ptop=self._ptop, - akap=KAPPA, - zvir=ZVIR, - last_step=last_step, - consv_te=self._conserve_total_energy, - mdt=self._timestep / self._k_split, - ) - else: - # TODO: When NQ=9, we shouldn't need to pass qcld explicitly - # since it's in self.tracers. It should not be an issue since - # we don't have self.tracers & qcld computation at the same - # time - # When NQ=8, we do need qcld passed explicitely - self._lagrangian_to_eulerian_obj( - self.tracers, - state.pt, - state.delp, - state.delz, - state.peln, - state.u, - state.v, - state.w, - self._cappa, - state.q_con, - state.qcld, - state.pkz, - state.pk, - state.pe, - state.phis, - state.ps, - self._wsd, - self._ak, - self._bk, - self._dp_initial, - self._ptop, - KAPPA, - ZVIR, - last_step, - self._conserve_total_energy, - self._timestep / self._k_split, - ) + # TODO: When NQ=9, we shouldn't need to pass qcld explicitly + # since it's in self.tracers. It should not be an issue since + # we don't have self.tracers & qcld computation at the same + # time + # When NQ=8, we do need qcld passed explicitely + self._lagrangian_to_eulerian_obj( + self.tracers, + state.pt, + state.delp, + state.delz, + state.peln, + state.u, + state.v, + state.w, + self._cappa, + state.q_con, + state.qcld, + state.pkz, + state.pk, + state.pe, + state.phis, + state.ps, + self._wsd, + self._ak, + self._bk, + self._dp_initial, + self._ptop, + KAPPA, + ZVIR, + last_step, + self._conserve_total_energy, + self._timestep / self._k_split, + ) self._checkpoint_remapping_out(state) # TODO: can we pull this block out of the loop intead of # using an if-statement? From 2645e4fe93ffc669e9a1d5f93793d6e344ff95bf Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 24 Jan 2025 13:23:10 -0500 Subject: [PATCH 141/252] Revert "Typo" This reverts commit c39a2a447319a8c7f0fb5656668798402be23c1f. --- pyFV3/stencils/fv_dynamics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 50ffba27..2c9d4f45 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -6,7 +6,7 @@ import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv -from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater +from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater, is from ndsl.checkpointer import NullCheckpointer from ndsl.comm.mpi import MPI from ndsl.constants import KAPPA, NQ, X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, ZVIR From 9934749f84a3e7f1bb13f79c6a10f6d00dc80edf Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 07:08:47 -0500 Subject: [PATCH 142/252] Fluxes and courant number as independent parameters (out of state) --- pyFV3/stencils/dyn_core.py | 119 ++++--------------------------------- 1 file changed, 13 insertions(+), 106 deletions(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 10312d7e..426f11bc 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -30,7 +30,6 @@ WrappedHaloUpdater, orchestrate, ) -from ndsl.checkpointer import NullCheckpointer from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -42,7 +41,7 @@ from ndsl.dsl.dace.orchestration import dace_inhibitor from ndsl.dsl.typing import Float, FloatField, FloatField64, FloatFieldIJ from ndsl.grid import DampingCoefficients, GridData -from ndsl.typing import Checkpointer, Communicator +from ndsl.typing import Communicator from pyFV3._config import AcousticDynamicsConfig from pyFV3.dycore_state import DycoreState from pyFV3.stencils.c_sw import CGridShallowWaterDynamics @@ -386,7 +385,6 @@ def __init__( phis: FloatFieldIJ, wsd: FloatFieldIJ, state, # [DaCe] hack to get around quantity as parameters for halo updates - checkpointer: Optional[Checkpointer] = None, ): """ Args: @@ -401,9 +399,6 @@ def __init__( config: configuration settings pfull: atmospheric Eulerian grid reference pressure (Pa) phis: surface geopotential height - checkpointer: if given, used to perform operations on model data - at specific points in model execution, such as testing against - reference data """ orchestrate( obj=self, @@ -411,32 +406,6 @@ def __init__( dace_compiletime_args=["state"], ) - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_csw", - dace_compiletime_args=["state", "tag"], - ) - - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_dsw_in", - dace_compiletime_args=["state", "tag"], - ) - - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_dsw_out", - dace_compiletime_args=["state", "tag"], - ) - - self.call_checkpointer = checkpointer is not None - if checkpointer is None: - self.checkpointer: Checkpointer = NullCheckpointer() - else: - self.checkpointer = checkpointer grid_indexing = stencil_factory.grid_indexing self.config = config if config.d_ext != 0: @@ -654,68 +623,6 @@ def __init__( def _get_da_min(self) -> Float: return Float(self._da_min) - def _checkpoint_csw(self, state: DycoreState, tag: str): - if self.call_checkpointer: - self.checkpointer( - f"C_SW-{tag}", - delpd=state.delp, - ptd=state.pt, - ud=state.u, - vd=state.v, - wd=state.w, - ucd=state.uc, - vcd=state.vc, - uad=state.ua, - vad=state.va, - utd=self._ut, - vtd=self._vt, - divgdd=self._divgd, - ) - - def _checkpoint_dsw_in(self, state: DycoreState): - if self.call_checkpointer: - self.checkpointer( - "D_SW-In", - ucd=state.uc, - vcd=state.vc, - wd=state.w, - # delpc is a temporary and not a variable in D_SW savepoint - delpcd=self._vt, - delpd=state.delp, - ud=state.u, - vd=state.v, - ptd=state.pt, - uad=state.ua, - vad=state.va, - zhd=self._zh, - divgdd=self._divgd, - xfxd=self._xfx, - yfxd=self._yfx, - mfxd=state.mfxd, - mfyd=state.mfyd, - ) - - def _checkpoint_dsw_out(self, state: DycoreState): - if self.call_checkpointer: - self.checkpointer( - "D_SW-Out", - ucd=state.uc, - vcd=state.vc, - wd=state.w, - delpcd=self._vt, - delpd=state.delp, - ud=state.u, - vd=state.v, - ptd=state.pt, - uad=state.ua, - vad=state.va, - divgdd=self._divgd, - xfxd=self._xfx, - yfxd=self._yfx, - mfxd=state.mfxd, - mfyd=state.mfyd, - ) - # TODO: fix me - we shouldn't need a function here, Dace is fudging the types # See https://github.com/GEOS-ESM/pace/issues/9 @dace_inhibitor @@ -730,6 +637,10 @@ def dt2(self, dt_acoustic_substep: Float) -> Float: def __call__( self, state: DycoreState, + mfxd, + mfyd, + cxd, + cyd, dpx, timestep: Float, # time to step forward by in seconds n_map=1, # [DaCe] replaces state.n_map @@ -749,10 +660,10 @@ def __call__( self._halo_updaters.q_con__cappa.wait() self._zero_data( - state.mfxd, - state.mfyd, - state.cxd, - state.cyd, + mfxd, + mfyd, + cxd, + cyd, self._heat_source, state.diss_estd, n_map == 1, @@ -801,7 +712,6 @@ def __call__( self._halo_updaters.w.wait() # compute the c-grid winds at t + 1/2 timestep - self._checkpoint_csw(state, tag="In") self.cgrid_shallow_water_lagrangian_dynamics( state.delp, state.pt, @@ -818,7 +728,6 @@ def __call__( state.omga, dt2, ) - self._checkpoint_csw(state, tag="Out") # TODO: Computing the pressure gradient outside of C_SW was originally done # so that we could transpose into a vertical-first memory ordering for the @@ -880,7 +789,6 @@ def __call__( self._halo_updaters.uc__vc.wait() # use the computed c-grid winds to evolve the d-grid winds forward # by 1 timestep - self._checkpoint_dsw_in(state) self.dgrid_shallow_water_lagrangian_dynamics( delpc=self._vt, delp=state.delp, @@ -893,10 +801,10 @@ def __call__( ua=state.ua, va=state.va, divgd=self._divgd, - mfx=state.mfxd, - mfy=state.mfyd, - cx=state.cxd, - cy=state.cyd, + mfx=mfxd, + mfy=mfyd, + cx=cxd, + cy=cyd, dpx=dpx, crx=self._crx, cry=self._cry, @@ -908,7 +816,6 @@ def __call__( diss_est=state.diss_estd, dt=dt_acoustic_substep, ) - self._checkpoint_dsw_out(state) # note that uc and vc are not needed at all past this point. # they will be re-computed from scratch on the next acoustic timestep. From e3b799b2db4c0de110c6a6b2afe050e7078ad4ac Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 07:10:05 -0500 Subject: [PATCH 143/252] Update translate test to independent fluxes/courant --- pyFV3/testing/translate_dyncore.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyFV3/testing/translate_dyncore.py b/pyFV3/testing/translate_dyncore.py index 17718bce..47415671 100644 --- a/pyFV3/testing/translate_dyncore.py +++ b/pyFV3/testing/translate_dyncore.py @@ -193,6 +193,10 @@ def compute_parallel(self, inputs, communicator): acoustic_dynamics( state, + mfxd=state.mfxd, + mfyd=state.mfyd, + cxd=state.cxd, + cyd=state.cyd, dpx=dpx, timestep=inputs["mdt"], n_map=state.n_map, From 754bc146be35aad921a0bf46b20fdc4fdbb111b2 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 07:14:02 -0500 Subject: [PATCH 144/252] Add compute total energy function (partial) --- pyFV3/stencils/compute_total_energy.py | 160 +++++++++++++++++++++++++ 1 file changed, 160 insertions(+) create mode 100644 pyFV3/stencils/compute_total_energy.py diff --git a/pyFV3/stencils/compute_total_energy.py b/pyFV3/stencils/compute_total_energy.py new file mode 100644 index 00000000..919193d8 --- /dev/null +++ b/pyFV3/stencils/compute_total_energy.py @@ -0,0 +1,160 @@ +from ndsl import StencilFactory, QuantityFactory +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, GRAV +from pyFV3._config import DynamicalCoreConfig +from pyFV3.tracers import Tracers +from pyFV3.stencils.moist_cv import moist_cv_nwat6_fn +from gt4py.cartesian.gtscript import ( + BACKWARD, + FORWARD, + interval, + computation, + K, +) +from ndsl.grid import GridData + + +def _compute_total_energy__stencil( + hs: FloatFieldIJ, # type: ignore + delp: FloatField, # type: ignore + delz: FloatField, # type: ignore + qc: FloatField, # type:ignore + pt: FloatField, # type: ignore + u: FloatField, # type: ignore + v: FloatField, # type: ignore + w: FloatField, # type: ignore + qvapor: FloatField, # type: ignore + qliquid: FloatField, # type: ignore + qrain: FloatField, # type: ignore + qsnow: FloatField, # type: ignore + qice: FloatField, # type: ignore + qgraupel: FloatField, # type: ignore + rsin2: FloatFieldIJ, # type: ignore + cosa_s: FloatFieldIJ, # type: ignore + phyz: FloatField, # type: ignore + te_2d: FloatFieldIJ, # type: ignore +): + """ + Dev Note: this is _very_ close to moist_cv.moist_te. The only numerical differences + is that the te/te_2d computation as an extra (1.+qc(i,j,k))*(1.-qd(i)) + + Args: + hs(in): + delp(in): + delz(in): + pt(in): + qc(in): + u(in): + v(in): + w(in): + qvapor(in): + qliquid(in): + qrain(in): + qsnow(in): + qice(in): + qgraupel(in): + rsin2(in): + cosa_s(in): + phyz(inout): + te_2d(out): + """ + + with computation(BACKWARD), interval(-1, None): + te_2d = 0.0 + phis = hs + with computation(BACKWARD), interval(0, -1): + phis = phis[K + 1] - GRAV * delz + with computation(FORWARD), interval(0, -1): + cvm, qd = moist_cv_nwat6_fn(qvapor, qliquid, qrain, qsnow, qice, qgraupel) + + te_2d = te_2d + delp * ( + cvm * pt * (1.0 + qc) * (1.0 - qd) + + 0.5 + * ( + phis + + phis[0, 0, 1] + + w**2.0 + + 0.5 + * rsin2 + * ( + u**2.0 + + u[0, 1, 0] ** 2.0 + + v**2.0 + + v[1, 0, 0] ** 2.0 + - (u + u[0, 1, 0]) * (v + v[1, 0, 0]) * cosa_s + ) + ) + ) + + +class ComputeTotalEnergy: + """Compute total energy performs the FV3-consistent + computation of the global total energy. + + It includes the potential, internal (latent and sensible heat), kinetic terms.""" + + def __init__( + self, + config: DynamicalCoreConfig, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + grid_data: GridData, + ) -> None: + if config.hydrostatic: + raise NotImplementedError( + "Dynamics (Compute Total Energy): " + " hydrostatic option is not implemented." + ) + + if not config.moist_phys: + raise NotImplementedError( + "Dynamics (Compute Total Energy): " + " moist_phys=False option is not implemented." + ) + + self._phyz = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Unknown", + dtype=Float, + ) + + self._compute_total_energy = stencil_factory.from_dims_halo( + func=_compute_total_energy__stencil, + compute_dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], + ) + self._rsin2 = grid_data.rsin2 + self._cosa_s = grid_data.cosa_s + + def __call__( + self, + hs: FloatFieldIJ, # type: ignore + delp: FloatField, # type: ignore + delz: FloatField, # type: ignore + qc: FloatField, # type:ignore + pt: FloatField, # type: ignore + u: FloatField, # type: ignore + v: FloatField, # type: ignore + w: FloatField, # type: ignore + tracers: Tracers, + te_2d: FloatFieldIJ, # type: ignore + ) -> None: + self._compute_total_energy( + hs=hs, + delp=delp, + delz=delz, + qc=qc, + pt=pt, + u=u, + v=v, + w=w, + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], + rsin2=self._rsin2, + cosa_s=self._cosa_s, + phyz=self._phyz, + te_2d=te_2d, + ) From e4acb471b2858565e919b82e7602b4992c70c88d Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 07:16:15 -0500 Subject: [PATCH 145/252] Fluxes & Courant number are local then increment global fields (f32) Add DryMassRoudnOff code (f32) Compute total energy in preamble GEOS Lagrangian code --- pyFV3/stencils/fv_dynamics.py | 463 ++++++++++++++++++++++++++++------ 1 file changed, 389 insertions(+), 74 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 50ffba27..2cfe1390 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -2,20 +2,38 @@ from typing import Mapping, Optional from dace.frontend.python.interface import nounroll as dace_no_unroll -from gt4py.cartesian.gtscript import PARALLEL, computation, interval +from gt4py.cartesian.gtscript import PARALLEL, FORWARD, computation, interval import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater from ndsl.checkpointer import NullCheckpointer from ndsl.comm.mpi import MPI -from ndsl.constants import KAPPA, NQ, X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, ZVIR +from ndsl.constants import ( + KAPPA, + NQ, + X_DIM, + Y_DIM, + Z_DIM, + Z_INTERFACE_DIM, + ZVIR, + Y_INTERFACE_DIM, + X_INTERFACE_DIM, +) from ndsl.dsl.dace.orchestration import dace_inhibitor, orchestrate -from ndsl.dsl.typing import Float, FloatField +from ndsl.dsl.typing import ( + Float, + FloatField, + FloatField64, + FloatFieldIJ64, + global_set_floating_point_precision, + NDSL_32BIT_FLOAT_TYPE, + NDSL_64BIT_FLOAT_TYPE, +) from ndsl.grid import DampingCoefficients, GridData from ndsl.logging import ndsl_log from ndsl.performance import NullTimer, Timer -from ndsl.stencils.basic_operations import copy_defn +from ndsl.stencils.basic_operations import copy_defn, set_value_defn from ndsl.stencils.c2l_ord import CubedToLatLon from ndsl.typing import Checkpointer, Communicator from pyFV3._config import DynamicalCoreConfig @@ -25,10 +43,132 @@ from pyFV3.stencils.dyn_core import AcousticDynamics from pyFV3.stencils.neg_adj3 import AdjustNegativeTracerMixingRatio from pyFV3.stencils.remapping import LagrangianToEulerian +from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyFV3.stencils.compute_total_energy import ComputeTotalEnergy +from pyFV3.version import IS_GEOS + + +class DryMassRoundOff: + def __init__( + self, + comm: Communicator, + quantity_factory: QuantityFactory, + stencil_factory: StencilFactory, + state: DycoreState, + hydrostatic: bool, + ) -> None: + self.psx_2d = quantity_factory.zeros( + dims=[X_DIM, Y_DIM], + units="unknown", + dtype=NDSL_64BIT_FLOAT_TYPE, + allow_mismatch_float_precision=True, + ) + self.dpx = quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=NDSL_64BIT_FLOAT_TYPE, + allow_mismatch_float_precision=True, + ) + self.dpx0_2d = quantity_factory.zeros( + dims=[X_DIM, Y_DIM], + units="unknown", + dtype=NDSL_64BIT_FLOAT_TYPE, + allow_mismatch_float_precision=True, + ) + + self._reset = stencil_factory.from_origin_domain( + DryMassRoundOff._reset_stencil, + origin=stencil_factory.grid_indexing.origin_compute(), + domain=stencil_factory.grid_indexing.domain_compute(), + ) + self._apply_psx_to_pe = stencil_factory.from_origin_domain( + DryMassRoundOff._apply_psx_to_pe_stencil, + origin=stencil_factory.grid_indexing.origin_compute(), + domain=stencil_factory.grid_indexing.domain_compute(), + ) + self._apply_dpx_to_psx = stencil_factory.from_origin_domain( + DryMassRoundOff._apply_dpx_to_psx_stencil, + origin=stencil_factory.grid_indexing.origin_compute(), + domain=stencil_factory.grid_indexing.domain_compute(), + ) + + halo_spec = quantity_factory.get_quantity_halo_spec( + dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], + n_halo=stencil_factory.grid_indexing.n_halo, + dtype=Float, + ) + self._pe_halo_updater = WrappedHaloUpdater( + comm.get_scalar_halo_updater([halo_spec]), + state, + ["pe"], + ) + + self._hydrostatic = hydrostatic + + @staticmethod + def _reset_stencil( + dpx: FloatField64, # type:ignore + psx_2d: FloatFieldIJ64, # type:ignore + pe: FloatField, # type:ignore + ): + with computation(PARALLEL), interval(...): + dpx = 0.0 + with computation(FORWARD), interval(-1, None): + psx_2d = pe[0, 0, 1] + + @staticmethod + def _apply_dpx_to_psx_stencil( + dpx: FloatField64, # type:ignore + dpx0_2d: FloatFieldIJ64, # type:ignore + psx_2d: FloatFieldIJ64, # type:ignore + ): + with computation(FORWARD), interval(0, 1): + dpx0_2d = dpx + + with computation(FORWARD), interval(1, None): + dpx0_2d += dpx + + with computation(FORWARD), interval(0, 1): + psx_2d += psx_2d + dpx0_2d + + @staticmethod + def _apply_psx_to_pe_stencil( + psx_2d: FloatFieldIJ64, # type:ignore + pe: FloatField, # type:ignore + ): + with computation(FORWARD), interval(-1, None): + pe[0, 0, 1] = psx_2d + + def reset(self, pe: FloatField): # type:ignore + self._reset(dpx=self.dpx, psx_2d=self.psx_2d, pe=pe) + + def apply(self, pe: FloatField): # type:ignore + self._apply_dpx_to_psx(self.dpx, self.dpx0_2d, self.psx_2d) + self._pe_halo_updater.update() + self._apply_psx_to_pe(self.psx_2d, pe) + + +def _increment_stencil( + value: FloatField, # type:ignore + increment: FloatField, # type:ignore +): + with computation(PARALLEL), interval(...): + value += increment + + +def _copy_cast_defn( + q_in_64: FloatField64, # type:ignore + q_out: FloatField, # type:ignore +): + with computation(PARALLEL), interval(...): + q_out = q_in_64 def pt_to_potential_density_pt( - pkz: FloatField, dp_initial: FloatField, q_con: FloatField, pt: FloatField + pkz: FloatField, # type: ignore + dp_initial: FloatField, # type: ignore + q_con: FloatField, # type: ignore + pt: FloatField, # type: ignore ): """ Args: @@ -43,7 +183,12 @@ def pt_to_potential_density_pt( pt = pt * (1.0 + dp_initial) * (1.0 - q_con) / pkz -def omega_from_w(delp: FloatField, delz: FloatField, w: FloatField, omega: FloatField): +def omega_from_w( + delp: FloatField, # type: ignore + delz: FloatField, # type: ignore + w: FloatField, # type: ignore + omega: FloatField, # type: ignore +): """ Args: delp (in): vertical layer thickness in Pa @@ -59,7 +204,7 @@ def fvdyn_temporaries( quantity_factory: QuantityFactory, ) -> Mapping[str, Quantity]: tmps = {} - for name in ["te_2d", "te0_2d", "wsd"]: + for name in ["te0_2d", "wsd"]: quantity = quantity_factory.zeros( dims=[X_DIM, Y_DIM], units="unknown", @@ -204,6 +349,14 @@ def __init__( self._da_min = damping_coefficients.da_min self.config = config + self.dry_mass_control = DryMassRoundOff( + comm=comm, + quantity_factory=quantity_factory, + stencil_factory=stencil_factory, + state=state, + hydrostatic=self.config.hydrostatic, + ) + tracer_transport = fvtp2d.FiniteVolumeTransport( stencil_factory=stencil_factory, quantity_factory=quantity_factory, @@ -218,7 +371,6 @@ def __init__( self.tracers[name] = state.__dict__[name] temporaries = fvdyn_temporaries(quantity_factory) - self._te_2d = temporaries["te_2d"] self._te0_2d = temporaries["te0_2d"] self._wsd = temporaries["wsd"] self._dp_initial = temporaries["dp1"] @@ -307,17 +459,38 @@ def __init__( hydrostatic=self.config.hydrostatic, ) - self._lagrangian_to_eulerian_obj = LagrangianToEulerian( + self._compute_total_energy = ComputeTotalEnergy( + config=config, stencil_factory=stencil_factory, quantity_factory=quantity_factory, - config=config.remapping, - area_64=grid_data.area_64, - nq=NQ, - pfull=self._pfull, - tracers=self.tracers, - checkpointer=checkpointer, + grid_data=grid_data, ) + if IS_GEOS: + self._lagrangian_to_eulerian_GEOS = LagrangianToEulerian_GEOS( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config.remapping, + comm=comm, + grid_data=grid_data, + nq=NQ, + pfull=self._pfull, + tracers=state.tracers, + adiabatic=config.adiabatic, + ) + + else: + self._lagrangian_to_eulerian_obj = LagrangianToEulerian( + stencil_factory=stencil_factory, + quantity_factory=quantity_factory, + config=config.remapping, + area_64=grid_data.area_64, + pfull=self._pfull, + tracers=state.tracers, + exclude_tracers=exclude_tracers, + checkpointer=checkpointer, + ) + full_xyz_spec = quantity_factory.get_quantity_halo_spec( dims=[X_DIM, Y_DIM, Z_DIM], n_halo=grid_indexing.n_halo, @@ -331,10 +504,75 @@ def __init__( self._conserve_total_energy = config.consv_te self._timestep = timestep.total_seconds() + # At 32-bit precision we still need + self._f32_correction = ( + global_set_floating_point_precision() == NDSL_32BIT_FLOAT_TYPE + ) + if self._f32_correction: + self._mfx_f64 = quantity_factory.zeros( + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=NDSL_64BIT_FLOAT_TYPE, + allow_mismatch_float_precision=True, + ) + self._mfy_f64 = quantity_factory.zeros( + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + units="unknown", + dtype=NDSL_64BIT_FLOAT_TYPE, + allow_mismatch_float_precision=True, + ) + self._cx_f64 = quantity_factory.zeros( + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=NDSL_64BIT_FLOAT_TYPE, + allow_mismatch_float_precision=True, + ) + self._cy_f64 = quantity_factory.zeros( + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + units="unknown", + dtype=NDSL_64BIT_FLOAT_TYPE, + allow_mismatch_float_precision=True, + ) + self._mfx_local = quantity_factory.zeros( + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._mfy_local = quantity_factory.zeros( + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._cx_local = quantity_factory.zeros( + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._cy_local = quantity_factory.zeros( + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._set_value = stencil_factory.from_origin_domain( + func=set_value_defn, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + self._increment = stencil_factory.from_origin_domain( + func=_increment_stencil, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + self._copy_cast = stencil_factory.from_origin_domain( + func=_copy_cast_defn, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + # See divergence_damping.py, _get_da_min for explanation of this function @dace_inhibitor - def _get_da_min(self) -> float: - return self._da_min + def _get_da_min(self) -> Float: # type: ignore + return Float(self._da_min) def _checkpoint_fvdynamics(self, state: DycoreState, tag: str): if self.call_checkpointer: @@ -461,6 +699,12 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool): if __debug__: log_on_rank_0("FV Setup") + # Reset fluxes + self._set_value(state.mfxd, Float(0.0)) + self._set_value(state.mfyd, Float(0.0)) + self._set_value(state.cxd, Float(0.0)) + self._set_value(state.cyd, Float(0.0)) + self._fv_setup_stencil( state.qvapor, state.qliquid, @@ -478,31 +722,50 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool): self._dp_initial, ) - if self._conserve_total_energy > 0: - raise NotImplementedError( - "Dynamical Core (fv_dynamics): compute total energy is not implemented" + # Compute total energy + if self.config.consv_te > 0.0: + self._compute_total_energy( + hs=state.phis, + delp=state.delp, + delz=state.delz, + qc=self._dp_initial, + pt=state.pt, + u=state.u, + v=state.v, + w=state.w, + tracers=state.tracers, + te_2d=self._te0_2d, ) - if (not self.config.rf_fast) and self.config.tau != 0: + # Rayleigh fast + if ( + not self.config.acoustic_dynamics.rf_fast + and self.config.acoustic_dynamics.tau > 0 + ): raise NotImplementedError( - "Dynamical Core (fv_dynamics): Rayleigh_Super," - " called when rf_fast=False and tau !=0, is not implemented" + "Dynamical Core (fv_dynamics): Rayleigh Friction is not implemented." ) - if self.config.adiabatic and self.config.kord_tm > 0: + # Adjust pt + if self.config.adiabatic: raise NotImplementedError( - "Dynamical Core (fv_dynamics): Adiabatic with positive kord_tm" - " is not implemented." + "Dynamical Core (fv_dynamics): Adiabatic pt adjust is not implemented." ) else: - if __debug__: - log_on_rank_0("Adjust pt") - self._pt_to_potential_density_pt( - state.pkz, - self._dp_initial, - state.q_con, - state.pt, - ) + if self.config.hydrostatic: + raise NotImplementedError( + "Dynamical Core (fv_dynamics): Hydrostatic pt adjust" + " is not implemented." + ) + else: + self._pt_to_potential_density_pt( + state.pkz, + self._dp_initial, + state.q_con, + state.pt, + ) + + self.dry_mass_control.reset(pe=state.pe) def __call__(self, *args, **kwargs): return self.step_dynamics(*args, **kwargs) @@ -526,10 +789,22 @@ def _compute(self, state: DycoreState, timer: Timer): log_on_rank_0("DynCore") with timer.clock("DynCore"): self.acoustic_dynamics( - state, + state=state, + mfxd=self._mfx_f64 if self._f32_correction else self._mfx_local, + mfyd=self._mfy_f64 if self._f32_correction else self._mfy_local, + cxd=self._cx_f64 if self._f32_correction else self._cx_local, + cyd=self._cy_f64 if self._f32_correction else self._cy_local, + dpx=self.dry_mass_control.dpx, timestep=self._timestep / self._k_split, n_map=n_map, ) + if self._f32_correction: + self._copy_cast(self._mfx_f64, self._mfx_local) + self._copy_cast(self._mfy_f64, self._mfy_local) + self._copy_cast(self._cx_f64, self._cx_local) + self._copy_cast(self._cy_f64, self._cy_local) + if last_step and self.config.hydrostatic: + self.dry_mass_control.apply(state.pe) if self.config.z_tracer: if __debug__: log_on_rank_0("TracerAdvection") @@ -538,10 +813,10 @@ def _compute(self, state: DycoreState, timer: Timer): self.tracer_advection( self.tracers, self._dp_initial, - state.mfxd, - state.mfyd, - state.cxd, - state.cyd, + x_mass_flux=self._mfx_local, + y_mass_flux=self._mfy_local, + x_courant=self._cx_local, + y_courant=self._cy_local, ) self._checkpoint_tracer_advection_out(state) else: @@ -567,44 +842,84 @@ def _compute(self, state: DycoreState, timer: Timer): with timer.clock("Remapping"): self._checkpoint_remapping_in(state) - # TODO: When NQ=9, we shouldn't need to pass qcld explicitly - # since it's in self.tracers. It should not be an issue since - # we don't have self.tracers & qcld computation at the same - # time - # When NQ=8, we do need qcld passed explicitely - self._lagrangian_to_eulerian_obj( - self.tracers, - state.pt, - state.delp, - state.delz, - state.peln, - state.u, - state.v, - state.w, - self._cappa, - state.q_con, - state.qcld, - state.pkz, - state.pk, - state.pe, - state.phis, - state.ps, - self._wsd, - self._ak, - self._bk, - self._dp_initial, - self._ptop, - KAPPA, - ZVIR, - last_step, - self._conserve_total_energy, - self._timestep / self._k_split, - ) + if IS_GEOS: + self._lagrangian_to_eulerian_GEOS( + tracers=state.tracers, + pt=state.pt, + delp=state.delp, + delz=state.delz, + peln=state.peln, + u=state.u, + v=state.v, + w=state.w, + mfx=self._mfx_local, + mfy=self._mfy_local, + cx=self._cx_local, + cy=self._cy_local, + cappa=self._cappa, + q_con=state.q_con, + pkz=state.pkz, + pk=state.pk, + pe=state.pe, + hs=state.phis, + te0_2d=self._te0_2d, + ps=state.ps, + wsd=self._wsd, + ak=self._ak, + bk=self._bk, + dp1=self._dp_initial, + ptop=self._ptop, + akap=KAPPA, + zvir=ZVIR, + last_step=last_step, + consv_te=self._conserve_total_energy, + mdt=self._timestep / self._k_split, + ) + else: + # TODO: When NQ=9, we shouldn't need to pass qcld explicitly + # since it's in self.tracers. It should not be an issue + # since we don't have self.tracers & qcld computation + # at the same time + # When NQ=8, we do need qcld passed explicitely + self._lagrangian_to_eulerian_obj( + state.tracers, + state.pt, + state.delp, + state.delz, + state.peln, + state.u, + state.v, + state.w, + self._cappa, + state.q_con, + state.pkz, + state.pk, + state.pe, + state.phis, + state.ps, + self._wsd, + self._ak, + self._bk, + self._dp_initial, + self._ptop, + KAPPA, + ZVIR, + last_step, + self._conserve_total_energy, + self._timestep / self._k_split, + ) self._checkpoint_remapping_out(state) # TODO: can we pull this block out of the loop intead of # using an if-statement? + + # Update state fluxes and courant number + self._increment(state.mfxd, self._mfx_local) + self._increment(state.mfyd, self._mfy_local) + self._increment(state.cxd, self._cx_local) + self._increment(state.cyd, self._cy_local) + if last_step: - da_min: Float = self._get_da_min() + da_min = self._get_da_min() if not self.config.hydrostatic: if __debug__: log_on_rank_0("Omega") @@ -620,7 +935,7 @@ def _compute(self, state: DycoreState, timer: Timer): if __debug__: log_on_rank_0("Del2Cubed") self._omega_halo_updater.update() - self._hyperdiffusion(state.omga, 0.18 * da_min) + self._hyperdiffusion(state.omga, Float(0.18) * da_min) if __debug__: log_on_rank_0("Neg Adj 3") From 12cdd7062d91f61c1ae5310574c611bd5f7778c6 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 07:16:38 -0500 Subject: [PATCH 146/252] Adapt translate test & allow for ULP<100 to pass test --- pyFV3/testing/translate_fvdynamics.py | 78 +++++++++++++------ .../translate/overrides/standard.yaml | 4 +- 2 files changed, 58 insertions(+), 24 deletions(-) diff --git a/pyFV3/testing/translate_fvdynamics.py b/pyFV3/testing/translate_fvdynamics.py index 5daea9d0..af4f5ac0 100644 --- a/pyFV3/testing/translate_fvdynamics.py +++ b/pyFV3/testing/translate_fvdynamics.py @@ -33,6 +33,27 @@ def __init__( self.namelist = DynamicalCoreConfig.from_namelist(namelist) +TRACERS_IN_PYFV3 = [ + "vapor", + "liquid", + "ice", + "rain", + "snow", + "graupel", + "cloud", +] + +TRACERS_IN_FORTRAN = [ + "qvapor", + "qliquid", + "qice", + "qrain", + "qsnow", + "qgraupel", + "qcld", +] + + class TranslateFVDynamics(ParallelTranslateBaseSlicing): compute_grid_option = True inputs: Dict[str, Any] = { @@ -90,25 +111,25 @@ class TranslateFVDynamics(ParallelTranslateBaseSlicing): "dims": [X_DIM, Z_INTERFACE_DIM, Y_DIM], "n_halo": 0, }, - "mfxd": { + "mfxd_FV": { "name": "accumulated_x_mass_flux", "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], "units": "unknown", "n_halo": 0, }, - "mfyd": { + "mfyd_FV": { "name": "accumulated_y_mass_flux", "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], "units": "unknown", "n_halo": 0, }, - "cxd": { + "cxd_FV": { "name": "accumulated_x_courant_number", "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], "units": "", "n_halo": (0, 3), }, - "cyd": { + "cyd_FV": { "name": "accumulated_y_courant_number", "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], "units": "", @@ -194,16 +215,6 @@ class TranslateFVDynamics(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "kg/kg", }, - "qo3mr": { - "name": "ozone_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - }, - "qsgs_tke": { - "name": "turbulent_kinetic_energy", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "m**2/s**2", - }, "qcld": { "name": "cloud_fraction", "dims": [X_DIM, Y_DIM, Z_DIM], @@ -243,8 +254,6 @@ def __init__( "qrain": grid.compute_dict(), "qsnow": grid.compute_dict(), "qgraupel": grid.compute_dict(), - "qo3mr": grid.compute_dict(), - "qsgs_tke": grid.compute_dict(), "qcld": {}, "ps": {}, "pe": { @@ -274,10 +283,10 @@ def __init__( "va": {}, "uc": grid.x3d_domain_dict(), "vc": grid.y3d_domain_dict(), - "mfxd": grid.x3d_compute_dict(), - "mfyd": grid.y3d_compute_dict(), - "cxd": grid.x3d_compute_domain_y_dict(), - "cyd": grid.y3d_compute_domain_x_dict(), + "mfxd_FV": grid.x3d_compute_dict(), + "mfyd_FV": grid.y3d_compute_dict(), + "cxd_FV": grid.x3d_compute_domain_y_dict(), + "cyd_FV": grid.y3d_compute_domain_x_dict(), "diss_estd": {}, } self._base.in_vars["data_vars"].update(fv_dynamics_vars) @@ -290,8 +299,13 @@ def __init__( self.max_error = 1e-5 self.ignore_near_zero_errors = {} - for qvar in utils.tracer_variables: - self.ignore_near_zero_errors[qvar] = True + self.ignore_near_zero_errors["qvapor"] = True + self.ignore_near_zero_errors["qliquid"] = True + self.ignore_near_zero_errors["qice"] = True + self.ignore_near_zero_errors["qrain"] = True + self.ignore_near_zero_errors["qsnow"] = True + self.ignore_near_zero_errors["qgraupel"] = True + self.ignore_near_zero_errors["qcld"] = True self.ignore_near_zero_errors["q_con"] = True self.dycore: Optional[fv_dynamics.DynamicalCore] = None self.stencil_factory = stencil_factory @@ -299,6 +313,19 @@ def __init__( def state_from_inputs(self, inputs): input_storages = super().state_from_inputs(inputs) + # extract tracers + input_storages["vapor"] = input_storages.pop("qvapor") + input_storages["liquid"] = input_storages.pop("qliquid") + input_storages["ice"] = input_storages.pop("qice") + input_storages["rain"] = input_storages.pop("qrain") + input_storages["snow"] = input_storages.pop("qsnow") + input_storages["graupel"] = input_storages.pop("qgraupel") + input_storages["cloud"] = input_storages.pop("qcld") + # Move fluxes and courant numbers + input_storages["mfxd"] = input_storages.pop("mfxd_FV") + input_storages["mfyd"] = input_storages.pop("mfyd_FV") + input_storages["cxd"] = input_storages.pop("cxd_FV") + input_storages["cyd"] = input_storages.pop("cyd_FV") # making sure we init DycoreState with the exact set of variables accepted_keys = [_field.name for _field in fields(DycoreState)] todelete = [] @@ -352,7 +379,12 @@ def outputs_from_state(self, state: dict): outputs = {} storages = {} for name, properties in self.outputs.items(): - if isinstance(state[name], Quantity): + if name in TRACERS_IN_FORTRAN: + idx = TRACERS_IN_FORTRAN.index(name) + storages[name] = state["tracers"][TRACERS_IN_PYFV3[idx]].data + elif name in ["mfxd_FV", "mfyd_FV", "cxd_FV", "cyd_FV"]: + storages[name] = state[name[:-3]].data + elif isinstance(state[name], Quantity): storages[name] = state[name].data elif len(self.outputs[name]["dims"]) > 0: storages[name] = state[name] # assume it's a storage diff --git a/tests/savepoint/translate/overrides/standard.yaml b/tests/savepoint/translate/overrides/standard.yaml index 4b9994e3..5f126a59 100644 --- a/tests/savepoint/translate/overrides/standard.yaml +++ b/tests/savepoint/translate/overrides/standard.yaml @@ -191,7 +191,9 @@ UtilVectors: - max_error: 2e-10 # 48_6ranks FVDynamics: - - max_error: 5e-5 # 48_6ranks using metric terms + - backend: numpy + multimodal: + ulp_threshold: 100 DivergenceDamping: - backend: dace:cpu From 968a535747deb95f78947cecc584939d8d68e0a8 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 07:21:42 -0500 Subject: [PATCH 147/252] Remove checkpointers --- pyFV3/stencils/fv_dynamics.py | 159 +--------------------------------- 1 file changed, 1 insertion(+), 158 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 2cfe1390..d04d7dfa 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -7,7 +7,6 @@ import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater -from ndsl.checkpointer import NullCheckpointer from ndsl.comm.mpi import MPI from ndsl.constants import ( KAPPA, @@ -35,7 +34,7 @@ from ndsl.performance import NullTimer, Timer from ndsl.stencils.basic_operations import copy_defn, set_value_defn from ndsl.stencils.c2l_ord import CubedToLatLon -from ndsl.typing import Checkpointer, Communicator +from ndsl.typing import Communicator from pyFV3._config import DynamicalCoreConfig from pyFV3.dycore_state import DycoreState from pyFV3.stencils import fvtp2d, tracer_2d_1l @@ -244,7 +243,6 @@ def __init__( phis: Quantity, state: DycoreState, timestep: timedelta, - checkpointer: Optional[Checkpointer] = None, ): """ Args: @@ -257,9 +255,6 @@ def __init__( phis: surface geopotential height state: model state timestep: model timestep - checkpointer: if given, used to perform operations on model data - at specific points in model execution, such as testing against - reference data """ orchestrate( obj=self, @@ -282,41 +277,6 @@ def __init__( dace_compiletime_args=["state", "timer"], ) - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_fvdynamics", - dace_compiletime_args=["state", "tag"], - ) - - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_remapping_in", - dace_compiletime_args=[ - "state", - ], - ) - - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_remapping_out", - dace_compiletime_args=["state"], - ) - - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_tracer_advection_in", - dace_compiletime_args=["state"], - ) - orchestrate( - obj=self, - config=stencil_factory.config.dace_config, - method_to_orchestrate="_checkpoint_tracer_advection_out", - dace_compiletime_args=["state"], - ) if timestep == timedelta(seconds=0): raise RuntimeError( "Bad dynamical core configuration:" @@ -324,11 +284,6 @@ def __init__( ) # nested and stretched_grid are options in the Fortran code which we # have not implemented, so they are hard-coded here. - self.call_checkpointer = checkpointer is not None - if checkpointer is None: - self.checkpointer: Checkpointer = NullCheckpointer() - else: - self.checkpointer = checkpointer nested = False stretched_grid = False grid_indexing = stencil_factory.grid_indexing @@ -427,7 +382,6 @@ def __init__( phis=self._phis, wsd=self._wsd, state=state, - checkpointer=checkpointer, ) self._hyperdiffusion = HyperdiffusionDamping( stencil_factory, @@ -488,7 +442,6 @@ def __init__( pfull=self._pfull, tracers=state.tracers, exclude_tracers=exclude_tracers, - checkpointer=checkpointer, ) full_xyz_spec = quantity_factory.get_quantity_halo_spec( @@ -574,109 +527,6 @@ def __init__( def _get_da_min(self) -> Float: # type: ignore return Float(self._da_min) - def _checkpoint_fvdynamics(self, state: DycoreState, tag: str): - if self.call_checkpointer: - self.checkpointer( - f"FVDynamics-{tag}", - u=state.u, - v=state.v, - w=state.w, - delz=state.delz, - # ua is not checked as its halo values differ from Fortran, - # this can be re-enabled if no longer comparing to Fortran, if the - # Fortran is updated to match the Python, or if the checkpointer - # can check only the compute domain values - # ua=state.ua, - va=state.va, - uc=state.uc, - vc=state.vc, - qvapor=state.qvapor, - ) - - def _checkpoint_remapping_in( - self, - state: DycoreState, - ): - if self.call_checkpointer: - self.checkpointer( - "Remapping-In", - pt=state.pt, - delp=state.delp, - delz=state.delz, - peln=state.peln.transpose( - [X_DIM, Z_INTERFACE_DIM, Y_DIM] - ), # [x, z, y] fortran data - u=state.u, - v=state.v, - w=state.w, - ua=state.ua, - va=state.va, - cappa=self._cappa, - pk=state.pk, - pe=state.pe.transpose( - [X_DIM, Z_INTERFACE_DIM, Y_DIM] - ), # [x, z, y] fortran data - phis=state.phis, - te_2d=self._te0_2d, - ps=state.ps, - wsd=self._wsd, - omga=state.omga, - dp1=self._dp_initial, - ) - - def _checkpoint_remapping_out( - self, - state: DycoreState, - ): - if self.call_checkpointer: - self.checkpointer( - "Remapping-Out", - pt=state.pt, - delp=state.delp, - delz=state.delz, - peln=state.peln.transpose( - [X_DIM, Z_INTERFACE_DIM, Y_DIM] - ), # [x, z, y] fortran data - u=state.u, - v=state.v, - w=state.w, - cappa=self._cappa, - pkz=state.pkz, - pk=state.pk, - pe=state.pe.transpose( - [X_DIM, Z_INTERFACE_DIM, Y_DIM] - ), # [x, z, y] fortran data - dp1=self._dp_initial, - ) - - def _checkpoint_tracer_advection_in( - self, - state: DycoreState, - ): - if self.call_checkpointer: - self.checkpointer( - "Tracer2D1L-In", - dp1=self._dp_initial, - mfxd=state.mfxd, - mfyd=state.mfyd, - cxd=state.cxd, - cyd=state.cyd, - ) - - def _checkpoint_tracer_advection_out( - self, - state: DycoreState, - ): - if self.call_checkpointer: - self.checkpointer( - "Tracer2D1L-Out", - dp1=self._dp_initial, - mfxd=state.mfxd, - mfyd=state.mfyd, - cxd=state.cxd, - cyd=state.cyd, - ) - def step_dynamics( self, state: DycoreState, @@ -689,9 +539,7 @@ def step_dynamics( timer: keep time of model sections state: model prognostic state and inputs """ - self._checkpoint_fvdynamics(state=state, tag="In") self._compute(state, timer) - self._checkpoint_fvdynamics(state=state, tag="Out") def compute_preamble(self, state: DycoreState, is_root_rank: bool): if self.config.hydrostatic: @@ -809,7 +657,6 @@ def _compute(self, state: DycoreState, timer: Timer): if __debug__: log_on_rank_0("TracerAdvection") with timer.clock("TracerAdvection"): - self._checkpoint_tracer_advection_in(state) self.tracer_advection( self.tracers, self._dp_initial, @@ -818,7 +665,6 @@ def _compute(self, state: DycoreState, timer: Timer): x_courant=self._cx_local, y_courant=self._cy_local, ) - self._checkpoint_tracer_advection_out(state) else: raise NotImplementedError("z_tracer=False is not implemented") @@ -840,8 +686,6 @@ def _compute(self, state: DycoreState, timer: Timer): if __debug__: log_on_rank_0("Remapping") with timer.clock("Remapping"): - self._checkpoint_remapping_in(state) - if IS_GEOS: self._lagrangian_to_eulerian_GEOS( tracers=state.tracers, @@ -908,7 +752,6 @@ def _compute(self, state: DycoreState, timer: Timer): self._conserve_total_energy, self._timestep / self._k_split, ) - self._checkpoint_remapping_out(state) # TODO: can we pull this block out of the loop intead of # using an if-statement? From 23b813b570311362d4cb3c9ee16bb413773fc47b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 07:23:28 -0500 Subject: [PATCH 148/252] Update README --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index 7a1838ab..c01225c4 100644 --- a/README.md +++ b/README.md @@ -42,3 +42,6 @@ Branches: - BASED ON `tracer_rework_part1` - REQUIRES: `ndsl` with tracer rework - ⚙️ `feature/fv_mapz/GEOS` @ Chris K: Remapping for GEOS +- ⚙️ `fix/GEOSv11_4_2/Dynamics`@Florian: Fix for the f32 & GEOS version of dynamics + - REQUIRES: `ndsl` with tracer rework + - REQUIRES: `tracer_rework_part1`, `fix/GEOSv11_4_2/Dyncore`, `fix/GEOS/TracerAdvection` From 0cfcd7392f7360584ca69ecb10115095c3cfc985 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Sun, 26 Jan 2025 11:52:43 -0500 Subject: [PATCH 149/252] Fix `Tracers` on GPU --- pyFV3/stencils/tracer_2d_1l.py | 14 +++++++++----- pyFV3/tracers.py | 7 +++++-- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index ce81c915..f2a0ba19 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -1,4 +1,5 @@ from typing import List +import numpy as np import gt4py.cartesian.gtscript as gtscript from gt4py.cartesian.gtscript import ( @@ -30,6 +31,8 @@ from ndsl.comm.communicator import Communicator, ReductionOperator from pyFV3.stencils.fvtp2d import FiniteVolumeTransport from pyFV3.tracers import Tracers +from ndsl.dsl.gt4py_utils import asarray +from ndsl.utils import safe_assign_array @gtscript.function @@ -368,10 +371,10 @@ def __call__( working_x_courant = x_courant working_y_courant = y_courant else: - self._tmp_mfx.data[:] = x_mass_flux.data[:] - self._tmp_mfy.data[:] = y_mass_flux.data[:] - self._tmp_cx.data[:] = x_courant.data[:] - self._tmp_cy.data[:] = y_courant.data[:] + safe_assign_array(self._tmp_mfx.data, x_mass_flux) + safe_assign_array(self._tmp_mfy.data, y_mass_flux) + safe_assign_array(self._tmp_cx.data, x_courant) + safe_assign_array(self._tmp_cy.data, y_courant) working_x_mass_flux = self._tmp_mfx working_y_mass_flux = self._tmp_mfy working_x_courant = self._tmp_cx @@ -422,7 +425,8 @@ def __call__( # a loop on the highest number of nsplit, but restraining # actual update in `apply_tracer_flux` to only the valid # K level for each tracers - max_n_split = i32(1.0 + self._cmax.view[:].max()) + cmax_on_host = asarray(self._cmax.view[:], to_type=np.ndarray) + max_n_split = i32(1.0 + cmax_on_host.max()) for current_nsplit in range(int(max_n_split)): last_call = current_nsplit == max_n_split - 1 diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py index 4b938c5f..7b7b4d08 100644 --- a/pyFV3/tracers.py +++ b/pyFV3/tracers.py @@ -6,6 +6,7 @@ from ndsl import Quantity, QuantityFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.utils import safe_assign_array # FOR REFERENCE - previous descriptive of the tracers, lining up with Pace work @@ -42,7 +43,9 @@ def copy_tracer_data( f"[pyFV3] Tracer {name} size ({data.shape}" f" is bigger than grid {qty.data.shape})" ) - qty.data[: data.shape[0], : data.shape[1], : data.shape[2]] = data + safe_assign_array( + qty.data[: data.shape[0], : data.shape[1], : data.shape[2]], data + ) self._quantities[name] = qty @property @@ -71,7 +74,7 @@ def as_4D_array(self) -> np.ndarray: # Skip the extra data point that is meant to align interface # and non interface fields for idx, q in enumerate(self.values()): - var4d[:, :, :, idx] = q.data[:-1, :-1, :-1] + safe_assign_array(var4d[:, :, :, idx], q.data[:-1, :-1, :-1]) return var4d def __getitem__(self, key): From fa9680e475b684647922f414fe962dd8fb4c1f85 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 28 Jan 2025 07:43:18 -0800 Subject: [PATCH 150/252] Small edit to bounds in mpp_global_sum --- pyFV3/stencils/mpp_global_sum.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py index f6af8b98..3651eddc 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/stencils/mpp_global_sum.py @@ -94,13 +94,13 @@ def regularize_ints(int_sum, prec, I_prec): break if positive: - for i in range(len(int_sum) - 1, 1, -1): + for i in range(len(int_sum) - 1, 0, -1): if int_sum[i] < 0: int_sum[i] = int_sum[i] + prec int_sum[i - 1] = int_sum[i - 1] - 1 else: - for i in range(len(int_sum) - 1, 1, -1): + for i in range(len(int_sum) - 1, 0, -1): if int_sum[i] > 0: int_sum[i] = int_sum[i] - prec int_sum[i - 1] = int_sum[i - 1] + 1 From 560a124ef9038c3a6f0b1b43a71c1aeaa47f9ab7 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 28 Jan 2025 07:52:44 -0800 Subject: [PATCH 151/252] Removed extraneous code preventing fv_dynamics from executing --- pyFV3/stencils/fv_dynamics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 2c9d4f45..50ffba27 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -6,7 +6,7 @@ import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv -from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater, is +from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater from ndsl.checkpointer import NullCheckpointer from ndsl.comm.mpi import MPI from ndsl.constants import KAPPA, NQ, X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, ZVIR From 8f4cd7cb70ace7d2b8eef89f3bc4a35e93294aac Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 28 Jan 2025 15:32:28 -0800 Subject: [PATCH 152/252] Linting --- fv_mapz.F90.SER | 102 +-- pyFV3/stencils/__init__.py | 2 +- pyFV3/stencils/map_single.py | 81 ++- pyFV3/stencils/moist_cv.py | 10 +- pyFV3/stencils/mpp_global_sum.py | 21 +- pyFV3/stencils/remapping.py | 29 +- pyFV3/stencils/remapping_GEOS.py | 33 +- pyFV3/stencils/scale_delz.py | 15 +- pyFV3/stencils/w_fix_consrv_moment.py | 46 +- tests/savepoint/translate/__init__.py | 39 +- .../translate/translate_MapN_Tracer_2d.py | 85 ++- .../translate/translate_Pressures_mapU.py | 142 ++--- .../translate/translate_Pressures_mapV.py | 131 ++-- .../translate/translate_cond_output.py | 40 +- .../translate/translate_cs_profile.py | 62 +- .../translate/translate_getMPIprop.py | 62 +- ...ranslate_lagrangian_contribution_interp.py | 71 ++- .../translate/translate_map1_ppm_W.py | 35 +- .../translate/translate_map1_ppm_delz.py | 55 +- .../translate/translate_map_scalar.py | 37 +- .../translate/translate_moistcvpluspt_2d.py | 2 +- .../translate_moistcvpluspt_2d_last_step.py | 50 +- .../translate/translate_moistcvpluste_2d.py | 52 +- .../translate/translate_mpp_global_sum.py | 60 +- .../translate/translate_pe_pk_delp_peln.py | 99 ++- .../translate/translate_pn2_pk_delp.py | 77 +-- .../translate/translate_remapping_GEOS.py | 585 +++++++++--------- .../translate/translate_scalar_profile.py | 66 +- .../savepoint/translate/translate_te_zsum.py | 16 +- .../translate_w_fix_consrv_moment.py | 38 +- 30 files changed, 1107 insertions(+), 1036 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 455c0f07..9aad009b 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -390,7 +390,7 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js-1, je+1 +!$ser verbatim do j = js-1, je+1 !$ser verbatim do i = is-1, ie+1 !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) !$ser verbatim enddo @@ -408,10 +408,10 @@ contains !$ser data pe_=pe_3d !$ser data ptop=ptop !$ser data qvapor=q(:,:,:,sphum) -!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qliquid=q(:,:,:,liq_wat) !$ser data qice=q(:,:,:,ice_wat) -!$ser data qrain=q(:,:,:,rainwat) -!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) !$ser data qgraupel=q(:,:,:,graupel) !$ser data qcld=q(:,:,:,cld_amt) !$ser data qo3mr=q(:,:,:,8) @@ -492,7 +492,7 @@ contains enddo else ! Transform "density pt" to "density temp" - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' !$ser savepoint MoistCVPlusPt_2d-In !$ser data qvapor_js=q(:,j,:,sphum) qliquid_js=q(:,j,:,liq_wat) qice_js=q(:,j,:,ice_wat) qrain_js=q(:,j,:,rainwat) qsnow_js=q(:,j,:,snowwat) qgraupel_js=q(:,j,:,graupel) qcld_js=q(:,j,:,cld_amt) gz1d=gz cvm=cvm r_vir=r_vir cappa=cappa rrg=rrg delp=delp delz=delz pt=pt k1k=k1k j_2d=js2d q_con=q_con @@ -512,7 +512,7 @@ contains enddo #endif enddo - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser savepoint MoistCVPlusPt_2d-Out !$ser data gz1d=gz cvm=cvm pt=pt cappa=cappa q_con=q_con !$ser verbatim endif @@ -660,7 +660,7 @@ contains km, pe2, pt, & is, ie, j, isd, ied, jsd, jed, akap, gmao_remap, P_MAP=1, conserv=.false.) else - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) @@ -674,7 +674,7 @@ contains call map_scalar(km, pn1, pt, gz, & km, pn2, pt, & is, ie, j, isd, ied, jsd, jed, 1, abs(kord_tm), t_min) - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser savepoint Map_Scalar-Out !$ser data q1=pt !$ser verbatim endif @@ -690,7 +690,7 @@ contains !$ser verbatim do i = is, ie !$ser verbatim pe1_3d(i,j,k) = pe1(i,k) !$ser verbatim pe2_3d(i,j,k) = pe2(i,k) - !$ser verbatim + !$ser verbatim !$ser verbatim if (k < km+1) dp2_3d(i,j,k) = dp2(i,k) !$ser verbatim enddo !$ser verbatim enddo @@ -708,7 +708,7 @@ contains !$ser verbatim endif call mapn_tracer(nq, km, pe1, pe2, q, dp2, kord_tr, j, & is, ie, isd, ied, jsd, jed, 0., fill) - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser savepoint MapN_Tracer_2d-Out !$ser data qvapor=q(:,:,:,1) !$ser data qliquid=q(:,:,:,2) @@ -739,7 +739,7 @@ contains if ( .not. hydrostatic ) then ! Remap vertical wind: - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) @@ -756,12 +756,12 @@ contains call map1_ppm (km, pe1, w, ws(is,j), & km, pe2, w, & is, ie, j, isd, ied, jsd, jed, -2, kord_wz) - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser savepoint Map1_PPM_W-Out !$ser data w_=w !$ser verbatim endif - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km+1 !$ser verbatim do i = is, ie !$ser verbatim pn1_3d(i,j,k) = pn1(i,k) @@ -801,7 +801,7 @@ contains delz(i,j,k) = -delz(i,j,k)*dp2(i,k) enddo enddo - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser savepoint Map1_PPM_delz-Out !$ser data delz_=delz !$ser verbatim endif @@ -809,7 +809,7 @@ contains !Fix excessive w - momentum conserving --- sjl ! gz(:) used here as a temporary array if ( w_limiter ) then -!$ser verbatim if(j == js2d) then +!$ser verbatim if(j == js2d) then !$ser verbatim do k = 1, km !$ser verbatim do i = is, ie !$ser verbatim dp2_3d(i,:,k) = dp2(i,k) @@ -888,7 +888,7 @@ contains w(i,j,k) = w2(i,k) enddo enddo -!$ser verbatim if(j == js2d) then +!$ser verbatim if(j == js2d) then !$ser savepoint W_fix_consrv_moment-Out !$ser data w=w !$ser verbatim endif @@ -914,8 +914,8 @@ contains !$ser data pe_=pe_3d !$ser data pe0_=pe0_3d !$ser data pe3_=pe3_3d -!$ser data ak=ak -!$ser data bk=bk +!$ser data ak=ak +!$ser data bk=bk !$ser data ptop=ptop !$ser data u_=u !$ser data kord_mt=kord_mt @@ -1000,7 +1000,7 @@ contains !$ser data pe_=pe_3d !$ser data pe0_=pe0_3d !$ser data pe3_=pe3_3d -!$ser data ak=ak +!$ser data ak=ak !$ser data bk=bk !$ser data v_=v !$ser data kord_mt=kord_mt @@ -1266,22 +1266,22 @@ contains if (remap_t) then ! print*,'CODE EXECUTED' ! Note: pt at this stage is T_v or T_m - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !!$ser verbatim print*,'MoistCVPlusPt_2D serialization' !$ser savepoint MoistCVPlusPkz_2d-In - !$ser data qvapor_js=q(:,j,:,sphum) - !$ser data qliquid_js=q(:,j,:,liq_wat) - !$ser data qice_js=q(:,j,:,ice_wat) - !$ser data qrain_js=q(:,j,:,rainwat) - !$ser data qsnow_js=q(:,j,:,snowwat) - !$ser data qgraupel_js=q(:,j,:,graupel) - !$ser data qcld_js=q(:,j,:,cld_amt) - !$ser data r_vir=r_vir - !$ser data cappa=cappa - !$ser data rrg=rrg - !$ser data delp=delp - !$ser data delz=delz - !$ser data pkz=pkz + !$ser data qvapor_js=q(:,j,:,sphum) + !$ser data qliquid_js=q(:,j,:,liq_wat) + !$ser data qice_js=q(:,j,:,ice_wat) + !$ser data qrain_js=q(:,j,:,rainwat) + !$ser data qsnow_js=q(:,j,:,snowwat) + !$ser data qgraupel_js=q(:,j,:,graupel) + !$ser data qcld_js=q(:,j,:,cld_amt) + !$ser data r_vir=r_vir + !$ser data cappa=cappa + !$ser data rrg=rrg + !$ser data delp=delp + !$ser data delz=delz + !$ser data pkz=pkz !$ser data pt=pt !$ser data k1k=k1k !$ser verbatim endif @@ -1299,9 +1299,9 @@ contains enddo #endif enddo - !$ser verbatim if(j == js2d) then + !$ser verbatim if(j == js2d) then !$ser savepoint MoistCVPlusPkz_2d-Out - !$ser data pkz=pkz + !$ser data pkz=pkz !$ser data cappa=cappa !$ser verbatim endif endif @@ -1415,7 +1415,7 @@ contains !$ser verbatim grav_=grav !!$ser verbatim print*,'MoistCVPlusTe_2D serialization' !$ser savepoint MoistCVPlusTe_2d-In - !$ser data qvapor_js=q(:,j,:,sphum) + !$ser data qvapor_js=q(:,j,:,sphum) !$ser data qliquid_js=q(:,j,:,liq_wat) !$ser data qice_js=q(:,j,:,ice_wat) !$ser data qrain_js=q(:,j,:,rainwat) @@ -1463,7 +1463,7 @@ contains enddo #endif enddo ! k-loop - !$ser verbatim if(j == js) then + !$ser verbatim if(j == js) then !$ser verbatim te_2d_f32 = real(te_2d) !$ser savepoint MoistCVPlusTe_2d-Out !$ser data te_2d_=te_2d_f32 @@ -1530,7 +1530,7 @@ contains !$ser verbatim te_2d=te_2d*gridstruct%area_64(is:ie,js:je) !!$ser verbatim print*,'Sum of input into mpp_global_sum = ', sum(te_2d) !$ser data inputArray=te_2d - + !$ser savepoint Mpp_global_sum-Out !$ser data tesum=tesum @@ -1632,14 +1632,14 @@ contains if ( .not. hydrostatic ) then ! print*,'CODE EXECUTED' !$ser savepoint MoistCVPlusPt_2d_last_step-In - !$ser data qvapor=q(:,:,:,sphum) + !$ser data qvapor=q(:,:,:,sphum) !$ser data qliquid=q(:,:,:,liq_wat) !$ser data qice=q(:,:,:,ice_wat) !$ser data qrain=q(:,:,:,rainwat) !$ser data qsnow=q(:,:,:,snowwat) !$ser data qgraupel=q(:,:,:,graupel) - !$ser data r_vir=r_vir - !$ser data dtmp=dtmp + !$ser data r_vir=r_vir + !$ser data dtmp=dtmp !$ser data pt=pt !$ser data pkz=pkz !$OMP do @@ -1702,7 +1702,7 @@ contains elseif ( last_step .and. adiabatic ) then print*,'CODE NOT TESTED HERE 17' !$OMP do - do k=1,km + do k=1,km do j=js,je do i=is,ie pt(i,j,k) = (pt(i,j,k)+dtmp*pkz(i,j,k)) @@ -1733,7 +1733,7 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser verbatim do k = 1, km+1 -!$ser verbatim do j = js-1, je+1 +!$ser verbatim do j = js-1, je+1 !$ser verbatim do i = is-1, ie+1 !$ser verbatim pe_3d(i,j,k) = pe(i,k,j) !$ser verbatim enddo @@ -1748,10 +1748,10 @@ contains !$ser savepoint Remapping_GEOS-Out !$ser data pe_=pe_3d !$ser data qvapor=q(:,:,:,sphum) -!$ser data qliquid=q(:,:,:,liq_wat) +!$ser data qliquid=q(:,:,:,liq_wat) !$ser data qice=q(:,:,:,ice_wat) -!$ser data qrain=q(:,:,:,rainwat) -!$ser data qsnow=q(:,:,:,snowwat) +!$ser data qrain=q(:,:,:,rainwat) +!$ser data qsnow=q(:,:,:,snowwat) !$ser data qgraupel=q(:,:,:,graupel) !$ser data qcld=q(:,:,:,cld_amt) !$ser data qo3mr=q(:,:,:,8) @@ -2104,7 +2104,7 @@ contains real q4(4,i1:i2,km) real pl, pr, qsum, dp, esl integer i, k, l, m, k0, jj - integer LM1,LP0,LP1 + integer LM1,LP0,LP1 !$ser verbatim real q4_1_temp(i1:i2, i1:i2,km) !$ser verbatim real q4_2_temp(i1:i2, i1:i2,km) @@ -2186,7 +2186,7 @@ contains ! NOTE : q1 and q2 fields being passed into map_scalar are identical variables ! even though q1 is declared at INTENT(IN) and q2 is declared as INTENT(IN/OUT). -!$ser verbatim if(j == 1) then +!$ser verbatim if(j == 1) then !$ser verbatim do k = 1,kn !$ser verbatim do i = i1,i2 !$ser verbatim q4_1_temp(i,j,k) = q4(1,i,k) @@ -2206,7 +2206,7 @@ contains !$ser verbatim pe2_temp(i,jj,km+1) = pe2(i,km+1) !$ser verbatim enddo !$ser verbatim enddo -!$ser savepoint Lagrangian_Contribution_Interp-In +!$ser savepoint Lagrangian_Contribution_Interp-In !$ser data q1=q1 pe1_=pe1_temp pe2_=pe2_temp q4_1=q4_1_temp q4_2=q4_2_temp !$ser data q4_3=q4_3_temp q4_4=q4_4_temp dp1_=dp1_temp !$ser verbatim endif @@ -2283,9 +2283,9 @@ contains 555 continue enddo -!$ser verbatim if(j == 1) then +!$ser verbatim if(j == 1) then !$ser savepoint Lagrangian_Contribution_Interp-Out -!$ser data q1=q2 +!$ser data q1=q2 !$ser verbatim endif end subroutine map_scalar diff --git a/pyFV3/stencils/__init__.py b/pyFV3/stencils/__init__.py index 40cbedd3..d5f5178d 100644 --- a/pyFV3/stencils/__init__.py +++ b/pyFV3/stencils/__init__.py @@ -13,6 +13,7 @@ from .fxadv import FiniteVolumeFluxPrep from .map_single import MapSingle from .mapn_tracer import MapNTracer +from .mpp_global_sum import mpp_global_sum from .neg_adj3 import AdjustNegativeTracerMixingRatio from .nh_p_grad import NonHydrostaticPressureGradient from .pk3_halo import PK3Halo @@ -28,7 +29,6 @@ from .updatedzd import UpdateHeightOnDGrid from .xppm import XPiecewiseParabolic from .yppm import YPiecewiseParabolic -from .mpp_global_sum import mpp_global_sum """ diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index 7bd202b7..2f4c5edf 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -4,7 +4,14 @@ from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntField, IntFieldIJ, BoolFieldIJ # noqa: F401 +from ndsl.dsl.typing import ( # noqa: F401 + BoolFieldIJ, + Float, + FloatField, + FloatFieldIJ, + IntField, + IntFieldIJ, +) from ndsl.stencils.basic_operations import copy_defn from pyFV3.stencils.remap_profile import RemapProfile @@ -15,6 +22,7 @@ def set_dp(dp1: FloatField, pe1: FloatField, lev: IntFieldIJ): with computation(FORWARD), interval(0, 1): lev = 0 + def lagrangian_contributions( q: FloatField, pe1: FloatField, @@ -77,6 +85,7 @@ def lagrangian_contributions( q = qsum / (pe2[0, 0, 1] - pe2) lev = lev - 1 + def lagrangian_contributions_interp( km: int, not_exit_loop: BoolFieldIJ, @@ -109,58 +118,65 @@ def lagrangian_contributions_interp( lev (inout): """ - # This computation creates a IntField that allows for "absolute" references in the k-dimension - # for q and pe1. + # This computation creates a IntField that allows for "absolute" references + # in the k-dimension for q and pe1. - # INDEX_LM1 and INDEX_LP0 is initialized such that if it's plugged into "q" (ex: q[0,0,INDEX_LM1]), the k level in - # q is "k = 0". + # INDEX_LM1 and INDEX_LP0 is initialized such that if it's plugged into "q" + # (ex: q[0,0,INDEX_LM1]), the k level in q is "k = 0". # For example, during the stencil computation at k = 2, INDEX_LM1[i,j,2] = -2 with computation(FORWARD): - with interval(0,1): + with interval(0, 1): INDEX_LM1 = 0 INDEX_LP0 = 0 - with interval(1,None): - INDEX_LM1 = INDEX_LM1[0,0,-1] - 1 - INDEX_LP0 = INDEX_LP0[0,0,-1] - 1 + with interval(1, None): + INDEX_LM1 = INDEX_LM1[0, 0, -1] - 1 + INDEX_LP0 = INDEX_LP0[0, 0, -1] - 1 # TODO: Can we make lev a 2D temporary? with computation(FORWARD), interval(...): LM1 = 1 LP0 = 1 not_exit_loop = True - while(LP0 <= km and not_exit_loop): - if(pe1[0,0,INDEX_LP0] < pe2): + while LP0 <= km and not_exit_loop: + if pe1[0, 0, INDEX_LP0] < pe2: LP0 = LP0 + 1 INDEX_LP0 = INDEX_LP0 + 1 else: not_exit_loop = False - LM1 = max(LP0-1,1) + LM1 = max(LP0 - 1, 1) INDEX_LM1 = INDEX_LM1 + (LM1 - 1) LP0 = min(LP0, km) - if(LP0 == 1): + if LP0 == 1: INDEX_LP0 = INDEX_LM1 - elif(LP0 <= km): - INDEX_LP0 = INDEX_LM1+1 + elif LP0 <= km: + INDEX_LP0 = INDEX_LM1 + 1 else: INDEX_LP0 = INDEX_LM1 - - if(LM1 == 1 and LP0 == 1): - q_temp = q[0,0,INDEX_LM1] + (q[0,0,INDEX_LM1+1] - q[0,0,INDEX_LM1]) * (pe2 - pe1[0,0,INDEX_LM1]) \ - / (pe1[0,0,INDEX_LM1+1] - pe1[0,0,INDEX_LM1]) - - elif(LM1 == km and LP0 == km): - q_temp = q[0,0,INDEX_LM1] + (q[0,0,INDEX_LM1] - q[0,0,INDEX_LM1-1]) * (pe2 - pe1[0,0,INDEX_LM1]) \ - / (pe1[0,0,INDEX_LM1] - pe1[0,0,INDEX_LM1-1]) - - elif(LM1 == 1 or LP0 == km): - q_temp = q[0,0,INDEX_LP0] + (q[0,0,INDEX_LM1] - q[0,0,INDEX_LP0]) * (pe2 - pe1[0,0,INDEX_LP0]) \ - / (pe1[0,0,INDEX_LM1] - pe1[0,0,INDEX_LP0]) - + + if LM1 == 1 and LP0 == 1: + q_temp = q[0, 0, INDEX_LM1] + ( + q[0, 0, INDEX_LM1 + 1] - q[0, 0, INDEX_LM1] + ) * (pe2 - pe1[0, 0, INDEX_LM1]) / ( + pe1[0, 0, INDEX_LM1 + 1] - pe1[0, 0, INDEX_LM1] + ) + + elif LM1 == km and LP0 == km: + q_temp = q[0, 0, INDEX_LM1] + ( + q[0, 0, INDEX_LM1] - q[0, 0, INDEX_LM1 - 1] + ) * (pe2 - pe1[0, 0, INDEX_LM1]) / ( + pe1[0, 0, INDEX_LM1] - pe1[0, 0, INDEX_LM1 - 1] + ) + + elif LM1 == 1 or LP0 == km: + q_temp = q[0, 0, INDEX_LP0] + (q[0, 0, INDEX_LM1] - q[0, 0, INDEX_LP0]) * ( + pe2 - pe1[0, 0, INDEX_LP0] + ) / (pe1[0, 0, INDEX_LM1] - pe1[0, 0, INDEX_LP0]) + else: - while(pe2 < pe1[0,0,lev] or pe2 > pe1[0,0,lev+1]): + while pe2 < pe1[0, 0, lev] or pe2 > pe1[0, 0, lev + 1]: lev = lev + 1 pl = (pe2 - pe1[0, 0, lev]) / dp1[0, 0, lev] if pe2[0, 0, 1] <= pe1[0, 0, lev + 1]: @@ -202,6 +218,7 @@ def lagrangian_contributions_interp( q = q_temp + class MapSingle: """ Fortran name is map_single, test classes are Map1_PPM_2d, Map_Scalar_2d @@ -282,7 +299,9 @@ def make_quantity(): dtype=int, ) self._km = grid_indexing.domain[2] - self._not_exit_loop = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=bool) + self._not_exit_loop = quantity_factory.zeros( + [X_DIM, Y_DIM], units="", dtype=bool + ) @property def i_extent(self): @@ -336,7 +355,7 @@ def __call__( Float(qmin), ) - if(interp == False): + if interp is False: self._lagrangian_contributions( q1, pe1, diff --git a/pyFV3/stencils/moist_cv.py b/pyFV3/stencils/moist_cv.py index d8dbeaf4..a8c53242 100644 --- a/pyFV3/stencils/moist_cv.py +++ b/pyFV3/stencils/moist_cv.py @@ -1,9 +1,9 @@ import gt4py.cartesian.gtscript as gtscript from gt4py.cartesian.gtscript import ( __INLINED, - PARALLEL, - FORWARD, BACKWARD, + FORWARD, + PARALLEL, computation, exp, interval, @@ -221,13 +221,13 @@ def moist_te( * ( phis + phis[0, 0, 1] - + w**2.0 + + w ** 2.0 + 0.5 * rsin2 * ( - u**2.0 + u ** 2.0 + u[0, 1, 0] ** 2.0 - + v**2.0 + + v ** 2.0 + v[1, 0, 0] ** 2.0 - (u + u[0, 1, 0]) * (v + v[1, 0, 0]) * cosa_s ) diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py index 3651eddc..39f93518 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/stencils/mpp_global_sum.py @@ -1,17 +1,18 @@ -from ndsl.quantity import Quantity -from ndsl.comm.comm_abc import ReductionOperator import numpy as np +from ndsl.comm.comm_abc import ReductionOperator +from ndsl.quantity import Quantity + def mpp_global_sum(inputArray, communicator, stencil_factory=None): NUMINT = 6 NUMBIT = 46 - r_prec = 2.0**NUMBIT - prec = 2**NUMBIT - I_prec = 1.0 / (2.0**NUMBIT) - pr = [r_prec**2, r_prec, 1.0, 1.0 / r_prec, 1.0 / r_prec**2, 1.0 / r_prec**3] - I_pr = [1.0 / r_prec**2, 1.0 / r_prec, 1.0, r_prec, r_prec**2, r_prec**3] - prec_error = (2**62 + (2**62 - 1)) / 6 + r_prec = 2.0 ** NUMBIT + prec = 2 ** NUMBIT + I_prec = 1.0 / (2.0 ** NUMBIT) + pr = [r_prec ** 2, r_prec, 1.0, 1.0 / r_prec, 1.0 / r_prec ** 2, 1.0 / r_prec ** 3] + I_pr = [1.0 / r_prec ** 2, 1.0 / r_prec, 1.0, r_prec, r_prec ** 2, r_prec ** 3] + prec_error = (2 ** 62 + (2 ** 62 - 1)) / 6 mag_max_term = 0.0 ints_sum = Quantity( @@ -35,15 +36,11 @@ def mpp_global_sum(inputArray, communicator, stencil_factory=None): ints_sum.data, pr, I_pr, inputArray[i, j], mag_max_term ) - # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' before carry_over') carry_overflow(ints_sum.data, prec, I_prec, prec_error) - # print("rank ", communicator.rank, "ints_sum = ", sum(ints_sum.data), ' after carry_over') communicator.all_reduce(ints_sum, ReductionOperator.SUM, ints_sum_reduce) - # print("rank ", communicator.rank, "sum(ints_sum_reduce) = ", sum(ints_sum_reduce.data), ' after all_reduce') regularize_ints(ints_sum_reduce.data, prec, I_prec) - # print("rank ", communicator.rank,"ints_sum_reduce = ", sum(ints_sum_reduce.data), ' after regularize_ints') sum_ = ints_to_real(ints_sum_reduce.data, pr) diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index b0d2ce4b..8e5a1dd6 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -1,7 +1,6 @@ from typing import Dict, Optional from gt4py.cartesian.gtscript import ( - __INLINED, BACKWARD, FORWARD, PARALLEL, @@ -130,12 +129,11 @@ def moist_cv_pt_pressure( remap_t (in): r_vir (in): """ - from __externals__ import hydrostatic#, kord_tm # moist_cv.moist_pt with computation(PARALLEL), interval(0, -1): # if __INLINED(kord_tm < 0): - if(remap_t): + if remap_t: cvm, gz, q_con, cappa, pt = moist_pt_func( qvapor, qliquid, @@ -154,7 +152,7 @@ def moist_cv_pt_pressure( # # delz_adjust # if __INLINED(not hydrostatic): # delz = -delz / delp - + # pressure_updates with computation(FORWARD): with interval(-1, None): @@ -204,11 +202,12 @@ def pn2_pk_delp( pn2 = log(pe2) pk = exp(akap * pn2) -def pe0_ptop_xmax(pe0: FloatField, - ptop: Float): - with computation(PARALLEL), interval(0,1): + +def pe0_ptop_xmax(pe0: FloatField, ptop: Float): + with computation(PARALLEL), interval(0, 1): pe0 = ptop + def pressures_mapu( pe: FloatField, # pe1: FloatField, @@ -302,6 +301,7 @@ def copy_from_below(a: FloatField, b: FloatField): with computation(PARALLEL), interval(1, None): b = a[0, 0, -1] + def pe_pk_delp_peln( pe: FloatField, pk: FloatField, @@ -327,26 +327,27 @@ def pe_pk_delp_peln( pe2 = ptop pn2 = peln pk2 = pk - with interval(1,-1): + with interval(1, -1): pe2 = ak + bk * pe_bottom pn2 = log(pe2) - pk2 = exp(akap*pn2) + pk2 = exp(akap * pn2) with interval(-1, None): pe2 = pe pn2 = peln pk2 = pk with computation(PARALLEL): - with interval(0,-1): + with interval(0, -1): pe = pe2 pk = pk2 - delp = pe2[0,0,1] - pe2[0,0,0] + delp = pe2[0, 0, 1] - pe2[0, 0, 0] peln = pn2 - with interval(-1,None): + with interval(-1, None): pe = pe2 pk = pk2 peln = pn2 + class LagrangianToEulerian: """ Fortran name is Lagrangian_to_Eulerian @@ -445,7 +446,7 @@ def __init__( # NOTE: In GEOS, remap_t is set to True in general # Add in the "remap_option" check later - if(True): + if True: self._remap_t = True self.kmp = grid_indexing.domain[2] - 1 @@ -700,7 +701,7 @@ def __call__( self._map_single_w(w, self._pe1, self._pe2, qs=wsd) self._map_single_delz(delz, self._pe1, self._pe2) - # W_limiter routine will go here + # W_limiter routine will go here self._undo_delz_adjust_and_copy_peln(delp, delz, peln, self._pe0, self._pn2) # if do_omega: # NOTE untested diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 3bd30fcb..106ff7c9 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -1,43 +1,37 @@ -from typing import Dict -from ndsl import ( - Quantity, - QuantityFactory, - StencilFactory, - orchestrate, -) +from ndsl import QuantityFactory, StencilFactory, orchestrate +from ndsl.comm.communicator import Communicator from ndsl.constants import ( + CV_AIR, + GRAV, X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM, Z_INTERFACE_DIM, - GRAV, - CV_AIR, ) from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK -from ndsl.stencils.basic_operations import adjust_divide_stencil from ndsl.grid import GridData -from ndsl.comm.communicator import Communicator +from ndsl.stencils.basic_operations import adjust_divide_stencil from pyFV3._config import RemappingConfig from pyFV3.stencils import moist_cv from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.moist_cv import moist_pt_last_step -from pyFV3.stencils.saturation_adjustment import SatAdjust3d -from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 -from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment +from pyFV3.stencils.mpp_global_sum import mpp_global_sum from pyFV3.stencils.remapping import ( + CONSV_MIN, init_pe, moist_cv_pt_pressure, + pe0_ptop_xmax, + pe_pk_delp_peln, pn2_pk_delp, pressures_mapu, pressures_mapv, - pe0_ptop_xmax, - pe_pk_delp_peln, - CONSV_MIN, ) -from pyFV3.stencils.mpp_global_sum import mpp_global_sum +from pyFV3.stencils.saturation_adjustment import SatAdjust3d +from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 +from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment from pyFV3.tracers import Tracers @@ -214,9 +208,10 @@ def __init__( stencil_factory, quantity_factory, abs(config.kord_tr), + self._nq, fill=config.fill, tracers=tracers, - exclude_tracers=[], + # exclude_tracers=[], ) self._map_single_w = MapSingle( diff --git a/pyFV3/stencils/scale_delz.py b/pyFV3/stencils/scale_delz.py index aae04f30..24aa6b0e 100644 --- a/pyFV3/stencils/scale_delz.py +++ b/pyFV3/stencils/scale_delz.py @@ -1,16 +1,19 @@ +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + from ndsl.dsl.typing import FloatField -from gt4py.cartesian.gtscript import FORWARD, PARALLEL, computation, interval + def rescale_delz_1( - delz: FloatField, - delp: FloatField, + delz: FloatField, + delp: FloatField, ): with computation(PARALLEL), interval(...): delz = -delz / delp + def rescale_delz_2( - delz: FloatField, - dp: FloatField, + delz: FloatField, + dp: FloatField, ): with computation(PARALLEL), interval(...): - delz = -delz * dp \ No newline at end of file + delz = -delz * dp diff --git a/pyFV3/stencils/w_fix_consrv_moment.py b/pyFV3/stencils/w_fix_consrv_moment.py index ed28cb77..26d5c3e1 100644 --- a/pyFV3/stencils/w_fix_consrv_moment.py +++ b/pyFV3/stencils/w_fix_consrv_moment.py @@ -1,5 +1,7 @@ -from ndsl.dsl.typing import FloatField, FloatFieldIJ, Float, BoolFieldIJ -from gt4py.cartesian.gtscript import FORWARD, BACKWARD, PARALLEL, computation, interval +from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval + +from ndsl.dsl.typing import BoolFieldIJ, Float, FloatField, FloatFieldIJ + def W_fix_consrv_moment( w: FloatField, @@ -23,59 +25,59 @@ def W_fix_consrv_moment( with computation(PARALLEL), interval(...): w2 = w - with computation(FORWARD): - with interval(0,1): + with computation(FORWARD): + with interval(0, 1): compute_performed = False - if(w2 > w_max): + if w2 > w_max: gz = (w2 - w_max) * dp2 w2 = w_max compute_performed = True - elif(w2 < w_min): + elif w2 < w_min: gz = (w2 - w_min) * dp2 w2 = w_min compute_performed = True - with interval(1,-1): - if(compute_performed): + with interval(1, -1): + if compute_performed: w2 = w2 + gz / dp2 compute_performed = False - if(w2 > w_max): + if w2 > w_max: gz = (w2 - w_max) * dp2 w2 = w_max compute_performed = True - elif(w2 < w_min): + elif w2 < w_min: gz = (w2 - w_min) * dp2 w2 = w_min compute_performed = True with computation(BACKWARD): - with interval(-1,None): + with interval(-1, None): compute_performed = False - if(w2 > w_max): + if w2 > w_max: gz = (w2 - w_max) * dp2 w2 = w_max compute_performed = True - elif(w2 < w_min): + elif w2 < w_min: gz = (w2 - w_min) * dp2 w2 = w_min compute_performed = True - with interval(1,-1): - if(compute_performed): + with interval(1, -1): + if compute_performed: w2 = w2 + gz / dp2 compute_performed = False - if(w2 > w_max): + if w2 > w_max: gz = (w2 - w_max) * dp2 w2 = w_max compute_performed = True - elif(w2 < w_min): + elif w2 < w_min: gz = (w2 - w_min) * dp2 w2 = w_min compute_performed = True - with computation(FORWARD), interval(0,1): - if(w2 > (w_max * 2.0)): + with computation(FORWARD), interval(0, 1): + if w2 > (w_max * 2.0): w2 = w_max * 2.0 - elif(w2 < (w_min * 2.0)): + elif w2 < (w_min * 2.0): w2 = w_min * 2.0 - + with computation(PARALLEL), interval(...): - w = w2 \ No newline at end of file + w = w2 diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index e3a4fc6c..1e11f4af 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -8,12 +8,14 @@ TranslateDivergenceCorner, TranslateVorticityTransport_Cgrid, ) +from .translate_cond_output import TranslateCond_output from .translate_corners import ( TranslateCopyCorners, TranslateFill4Corners, TranslateFillCorners, TranslateFillCornersVector, ) +from .translate_cs_profile import TranslateCS_Profile from .translate_cubedtolatlon import TranslateCubedToLatLon from .translate_d2a2c_vect import TranslateD2A2C_Vect from .translate_d_sw import ( @@ -32,6 +34,7 @@ from .translate_fvsubgridz import TranslateFVSubgridZ from .translate_fvtp2d import TranslateFvTp2d, TranslateFvTp2d_2 from .translate_fxadv import TranslateFxAdv +from .translate_getMPIprop import TranslateGetMPIProp from .translate_grid import ( TranslateAGrid, TranslateDerivedTrig, @@ -62,45 +65,43 @@ TranslateJablonowskiBaroclinic, TranslatePVarAuxiliaryPressureVars, ) +from .translate_lagrangian_contribution_interp import ( + TranslateLagrangian_Contribution_Interp, +) from .translate_last_step import TranslateLastStep +from .translate_map1_ppm_delz import TranslateMap1_PPM_delz +from .translate_map1_ppm_W import TranslateMap1_PPM_W +from .translate_map_scalar import TranslateMap_Scalar +from .translate_MapN_Tracer_2d import TranslateMapN_Tracer_2d from .translate_moistcvpluspkz_2d import TranslateMoistCVPlusPkz_2d from .translate_moistcvpluspt_2d import TranslateMoistCVPlusPt_2d +from .translate_moistcvpluspt_2d_last_step import TranslateMoistCVPlusPt_2d_last_step +from .translate_moistcvpluste_2d import TranslateMoistCVPlusTe_2d +from .translate_mpp_global_sum import TranslateMpp_global_sum from .translate_neg_adj3 import TranslateNeg_Adj3 from .translate_nh_p_grad import TranslateNH_P_Grad from .translate_pe_halo import TranslatePE_Halo +from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln from .translate_pk3_halo import TranslatePK3_Halo from .translate_pressureadjustedtemperature_nonhydrostatic import ( TranslatePressureAdjustedTemperature_NonHydrostatic, ) +from .translate_Pressures_mapU import TranslatePressures_mapU +from .translate_Pressures_mapV import TranslatePressures_mapV from .translate_qsinit import TranslateQSInit from .translate_ray_fast import TranslateRay_Fast from .translate_remapping import TranslateRemapping +from .translate_remapping_GEOS import TranslateRemapping_GEOS from .translate_riem_solver3 import TranslateRiem_Solver3 from .translate_riem_solver_c import TranslateRiem_Solver_C from .translate_satadjust3d import TranslateSatAdjust3d +from .translate_scalar_profile import TranslateScalar_Profile +from .translate_te_zsum import TranslateTe_Zsum from .translate_tracer2d1l import TranslateTracer2D1L from .translate_updatedzc import TranslateUpdateDzC from .translate_updatedzd import TranslateUpdateDzD +from .translate_w_fix_consrv_moment import TranslateW_fix_consrv_moment from .translate_xppm import TranslateXPPM, TranslateXPPM_2 from .translate_xtp_u import TranslateXTP_U from .translate_yppm import TranslateYPPM, TranslateYPPM_2 from .translate_ytp_v import TranslateYTP_V - -from .translate_w_fix_consrv_moment import TranslateW_fix_consrv_moment -from .translate_lagrangian_contribution_interp import TranslateLagrangian_Contribution_Interp -from .translate_remapping_GEOS import TranslateRemapping_GEOS -from .translate_scalar_profile import TranslateScalar_Profile -from .translate_MapN_Tracer_2d import TranslateMapN_Tracer_2d -from .translate_map_scalar import TranslateMap_Scalar -from .translate_cs_profile import TranslateCS_Profile -from .translate_map1_ppm_W import TranslateMap1_PPM_W -from .translate_map1_ppm_delz import TranslateMap1_PPM_delz -from .translate_Pressures_mapU import TranslatePressures_mapU -from .translate_Pressures_mapV import TranslatePressures_mapV -from .translate_pe_pk_delp_peln import TranslatePE_pk_delp_peln -from .translate_getMPIprop import TranslateGetMPIProp -from .translate_mpp_global_sum import TranslateMpp_global_sum -from .translate_moistcvpluste_2d import TranslateMoistCVPlusTe_2d -from .translate_te_zsum import TranslateTe_Zsum -from .translate_moistcvpluspt_2d_last_step import TranslateMoistCVPlusPt_2d_last_step -from .translate_cond_output import TranslateCond_output \ No newline at end of file diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index 4459a98a..b73570be 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -1,10 +1,9 @@ -from gt4py.cartesian.gtscript import PARALLEL, computation, interval - -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid +from ndsl import Namelist, StencilFactory from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.mapn_tracer import MapNTracer + class TranslateMapN_Tracer_2d(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) @@ -14,91 +13,90 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.in_vars["data_vars"] = { "qvapor": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qliquid": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qice": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qrain": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsnow": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qgraupel": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qcld": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qo3mr": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsgs_tke": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "pe1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, "pe2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, - "dp2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1 + "kend": grid.npz - 1, }, } self.out_vars = { "qvapor": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qliquid": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "qice": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qrain": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsnow": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qgraupel": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qcld": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qo3mr": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsgs_tke": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, } # Value from GEOS - self.kord = 9 + self.kord = 9 # mode / iv set to 1 from GEOS - self.mode = 1 + self.mode = 1 self.nq = 9 @@ -106,15 +104,16 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto def compute_from_storage(self, inputs): - tracers = { "qvapor": inputs["qvapor"], - "qliquid": inputs["qliquid"], - "qice": inputs["qice"], - "qrain": inputs["qrain"], - "qsnow": inputs["qsnow"], - "qgraupel": inputs["qgraupel"], - "qcld": inputs["qcld"], - "qo3mr": inputs["qo3mr"], - "qsgs_tke": inputs["qsgs_tke"], + tracers = { + "qvapor": inputs["qvapor"], + "qliquid": inputs["qliquid"], + "qice": inputs["qice"], + "qrain": inputs["qrain"], + "qsnow": inputs["qsnow"], + "qgraupel": inputs["qgraupel"], + "qcld": inputs["qcld"], + "qo3mr": inputs["qo3mr"], + "qsgs_tke": inputs["qsgs_tke"], } self._compute_func = MapNTracer( @@ -127,10 +126,10 @@ def compute_from_storage(self, inputs): ) self._compute_func( - inputs["pe1_"], - inputs["pe2_"], - inputs["dp2_"], - tracers, - ) + inputs["pe1_"], + inputs["pe2_"], + inputs["dp2_"], + tracers, + ) return inputs diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index 6ea80092..f370a858 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid -from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Y_INTERFACE_DIM +from ndsl import Namelist, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.map_single import MapSingle +from pyFV3.stencils.remapping import pe0_ptop_xmax, pressures_mapu + class TranslatePressures_mapU(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): @@ -14,58 +15,49 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.in_vars["data_vars"] = { "pe_": { - "istart": grid.is_-1, - "iend": grid.ie+1, - "jstart": grid.js-1, - "jend": grid.je+1, - "kend": grid.npz - }, - "ak":{ - - }, - "bk":{ - + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, + "kend": grid.npz, }, + "ak": {}, + "bk": {}, "pe0_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz + "jend": grid.je + 1, + "kend": grid.npz, }, "pe3_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz + "jend": grid.je + 1, + "kend": grid.npz, }, - "u_": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.jsd, - "jend": grid.jed+1, - "kend": grid.npz-1, + "jend": grid.jed + 1, + "kend": grid.npz - 1, }, - "mfy_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, - "cy_": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, - - } self.in_vars["parameters"] = [ "ptop", @@ -75,57 +67,55 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.out_vars = { "pe0_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz + "jend": grid.je + 1, + "kend": grid.npz, }, "pe3_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz + "jend": grid.je + 1, + "kend": grid.npz, }, "u_": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.jsd, - "jend": grid.jed+1, - "kend": grid.npz-1, + "jend": grid.jed + 1, + "kend": grid.npz - 1, }, - "mfy_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, - "cy_": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, } grid_indexing = stencil_factory.grid_indexing - self.dims=[X_DIM, Y_DIM, Z_DIM] + self.dims = [X_DIM, Y_DIM, Z_DIM] self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[0],1,grid_indexing.domain[2] + 1) + domain=(grid_indexing.domain[0], 1, grid_indexing.domain[2] + 1), ) self._pe0_ptop_xmax = stencil_factory.from_origin_domain( pe0_ptop_xmax, - origin=(grid_indexing.domain[0]+3,3,0), - domain=(1,1,grid_indexing.domain[2] + 1) + origin=(grid_indexing.domain[0] + 3, 3, 0), + domain=(1, 1, grid_indexing.domain[2] + 1), ) def compute_from_storage(self, inputs): @@ -138,36 +128,36 @@ def compute_from_storage(self, inputs): ) self._pressures_mapu( - inputs["pe_"], - inputs["ak"], - inputs["bk"], - inputs["pe0_"], - inputs["pe3_"], - inputs["ptop"], - ) + inputs["pe_"], + inputs["ak"], + inputs["bk"], + inputs["pe0_"], + inputs["pe3_"], + inputs["ptop"], + ) self._pe0_ptop_xmax( - inputs["pe0_"], - inputs["ptop"], - ) + inputs["pe0_"], + inputs["ptop"], + ) self._map1_ppm_u( - inputs["u_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + inputs["u_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) self._map1_ppm_u( - inputs["mfy_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) - + inputs["mfy_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) + self._map1_ppm_u( - inputs["cy_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + inputs["cy_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) return inputs diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index d233f669..147339f6 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid -from ndsl.constants import X_DIM, Y_DIM, Z_DIM, X_INTERFACE_DIM +from ndsl import Namelist, StencilFactory +from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils.remapping import pressures_mapv +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.map_single import MapSingle +from pyFV3.stencils.remapping import pressures_mapv + class TranslatePressures_mapV(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): @@ -14,57 +15,49 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.in_vars["data_vars"] = { "pe_": { - "istart": grid.is_-1, - "iend": grid.ie+1, - "jstart": grid.js-1, - "jend": grid.je+1, - "kend": grid.npz+1 - }, + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, + "kend": grid.npz + 1, + }, "pe0_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1 + "jend": grid.je + 1, + "kend": grid.npz + 1, }, "pe3_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1 - }, - "ak":{ - - }, - "bk":{ - + "jend": grid.je + 1, + "kend": grid.npz + 1, }, + "ak": {}, + "bk": {}, "v_": { "istart": grid.isd, - "iend": grid.ied+1, + "iend": grid.ied + 1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "mfx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "cx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - - } self.in_vars["parameters"] = [ "kord_mt", @@ -73,46 +66,44 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.out_vars = { "pe0_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1, + "jend": grid.je + 1, + "kend": grid.npz + 1, }, "pe3_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz+1, + "jend": grid.je + 1, + "kend": grid.npz + 1, }, "v_": { "istart": grid.isd, - "iend": grid.ied+1, + "iend": grid.ied + 1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "mfx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "cx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, } grid_indexing = stencil_factory.grid_indexing - self.dims=[X_DIM, Y_DIM, Z_DIM] + self.dims = [X_DIM, Y_DIM, Z_DIM] self._pressures_mapv = stencil_factory.from_origin_domain( pressures_mapv, @@ -134,31 +125,31 @@ def compute_from_storage(self, inputs): ) self._pressures_mapv( - inputs["pe_"], - inputs["ak"], - inputs["bk"], - inputs["pe0_"], - inputs["pe3_"], - ) + inputs["pe_"], + inputs["ak"], + inputs["bk"], + inputs["pe0_"], + inputs["pe3_"], + ) self._map1_ppm_v( - inputs["v_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) - + inputs["v_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) + self._map1_ppm_v( - inputs["mfx_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) - + inputs["mfx_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) + self._map1_ppm_v( - inputs["cx_"], - inputs["pe0_"], - inputs["pe3_"], - interp=False, - ) + inputs["cx_"], + inputs["pe0_"], + inputs["pe3_"], + interp=False, + ) return inputs diff --git a/tests/savepoint/translate/translate_cond_output.py b/tests/savepoint/translate/translate_cond_output.py index 0dbbe71b..6bd94b2e 100644 --- a/tests/savepoint/translate/translate_cond_output.py +++ b/tests/savepoint/translate/translate_cond_output.py @@ -1,38 +1,35 @@ -from gt4py.cartesian.gtscript import PARALLEL, computation, interval - -from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField, Float -from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j +from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils import moist_cv + class TranslateCond_output(TranslateFortranData2Py): def __init__(self, grid, namelist, stencil_factory): super().__init__(grid, namelist, stencil_factory) self.stencil_factory = stencil_factory self.in_vars["data_vars"] = { "qliquid": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "qice": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qrain": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsnow": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qgraupel": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q_con": { - "kend": grid.npz-1, - } + "kend": grid.npz - 1, + }, } self.out_vars = { "q_con": { - "kend": grid.npz-1, + "kend": grid.npz - 1, } } @@ -44,11 +41,12 @@ def __init__(self, grid, namelist, stencil_factory): def compute_from_storage(self, inputs): - self.compute_func(inputs["q_con"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - ) + self.compute_func( + inputs["q_con"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + ) return inputs diff --git a/tests/savepoint/translate/translate_cs_profile.py b/tests/savepoint/translate/translate_cs_profile.py index 30b73f4a..c63617f2 100644 --- a/tests/savepoint/translate/translate_cs_profile.py +++ b/tests/savepoint/translate/translate_cs_profile.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid +from ndsl import Namelist, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.remap_profile import RemapProfile + class TranslateCS_Profile(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) @@ -12,46 +13,47 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { - "qs_": {"istart": grid.is_, + "qs_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_1": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "q4_2": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_3": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_4": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "dp1_":{ + "dp1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, } self.in_vars["parameters"] = [ @@ -60,29 +62,33 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto ] self.out_vars = { - "q4_1": {"istart": grid.is_, + "q4_1": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "q4_2": {"istart": grid.is_, + "q4_2": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "q4_3": {"istart": grid.is_, + "q4_3": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "q4_4": {"istart": grid.is_, + "q4_4": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, } @@ -92,15 +98,15 @@ def compute_from_storage(self, inputs): self.quantity_factory, inputs["kord_"], inputs["iv_"], - dims=[X_DIM, Y_DIM, Z_DIM] + dims=[X_DIM, Y_DIM, Z_DIM], ) - + self._compute_func( - inputs["qs_"], - inputs["q4_1"], - inputs["q4_2"], - inputs["q4_3"], - inputs["q4_4"], - inputs["dp1_"], - ) + inputs["qs_"], + inputs["q4_1"], + inputs["q4_2"], + inputs["q4_3"], + inputs["q4_4"], + inputs["dp1_"], + ) return inputs diff --git a/tests/savepoint/translate/translate_getMPIprop.py b/tests/savepoint/translate/translate_getMPIprop.py index ba5e1186..236ca062 100644 --- a/tests/savepoint/translate/translate_getMPIprop.py +++ b/tests/savepoint/translate/translate_getMPIprop.py @@ -1,9 +1,11 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid +import numpy as np + +from ndsl import Namelist, StencilFactory +from ndsl.quantity import Quantity from ndsl.stencils.testing import ParallelTranslate +from ndsl.stencils.testing.grid import Grid from ndsl.typing import Communicator -from ndsl.quantity import Quantity -import numpy as np + class TranslateGetMPIProp(ParallelTranslate): def __init__( @@ -14,42 +16,34 @@ def __init__( ): print("Base TranslateGetMPIProp is initialized") super().__init__(grid, namelist, stencil_factory) - self._base.in_vars["data_vars"] = { - "delz" :{ - - } - } - self._base.out_vars = { - "delz" :{ - - } - } + self._base.in_vars["data_vars"] = {"delz": {}} + self._base.out_vars = {"delz": {}} len_k = 10 - a = [1,2,3,4,5] + a = [1, 2, 3, 4, 5] self._testQuantity_1D = Quantity( - data=np.array(a,dtype=np.float32), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - + data=np.array(a, dtype=np.float32), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + self._testQuantity_2D = Quantity( - data=np.ones([5,5],dtype=np.float32), - dims=["I","J"], - units="dunno2", - gt4py_backend=stencil_factory.backend, - ) - + data=np.ones([5, 5], dtype=np.float32), + dims=["I", "J"], + units="dunno2", + gt4py_backend=stencil_factory.backend, + ) + self._testQuantity_3D = Quantity( - data=np.ones([3,3,3],dtype=np.float32), - dims=["I","J","K"], - units="dunno3", - gt4py_backend=stencil_factory.backend, - ) - + data=np.ones([3, 3, 3], dtype=np.float32), + dims=["I", "J", "K"], + units="dunno3", + gt4py_backend=stencil_factory.backend, + ) + def compute_parallel(self, inputs, communicator: Communicator): print("Communicator rank = ", communicator.rank) print("Communicator size = ", communicator.size) @@ -66,4 +60,4 @@ def compute_parallel(self, inputs, communicator: Communicator): print("global_sum_q.data = ", global_sum_q.data) print("global_sum_q.metadata = ", global_sum_q.metadata) - return inputs \ No newline at end of file + return inputs diff --git a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py index 011ac497..33b5f7d6 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py @@ -1,10 +1,11 @@ -from ndsl import StencilFactory, Namelist -from ndsl.dsl.typing import FloatField, BoolFieldIJ, IntField, IntFieldIJ, Int, Bool +from ndsl import Namelist, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import Bool, BoolFieldIJ, FloatField, Int, IntField, IntFieldIJ from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.map_single import lagrangian_contributions_interp + class test_Lagragian_Contribution_Interp: def __init__( self, @@ -18,7 +19,7 @@ def __init__( self._lagrangian_contributions_interp = stencil_factory.from_origin_domain( func=lagrangian_contributions_interp, origin=grid_indexing.origin_compute(), - domain=(grid.nic, 1, grid.npz) + domain=(grid.nic, 1, grid.npz), ) def __call__( @@ -53,74 +54,76 @@ def __call__( lev, ) + class TranslateLagrangian_Contribution_Interp(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid - self.compute_func = test_Lagragian_Contribution_Interp(self.stencil_factory, self.grid) # type: ignore + self.compute_func = test_Lagragian_Contribution_Interp( + self.stencil_factory, self.grid + ) # type: ignore self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { "q1": { - "kend": grid.npz-1, - }, - + "kend": grid.npz - 1, + }, "pe1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, "pe2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, "q4_1": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "q4_2": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_3": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_4": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "dp1_":{ + "dp1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, } self.out_vars = { "q1": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, } def compute_from_storage(self, inputs): @@ -150,19 +153,19 @@ def compute_from_storage(self, inputs): ) self.compute_func( - self.grid.npz, - self._not_exit_loop, - self._INDEX_LM1, - self._INDEX_LP0, - inputs["q1"], - inputs["pe1_"], - inputs["pe2_"], - inputs["q4_1"], - inputs["q4_2"], - inputs["q4_3"], - inputs["q4_4"], - inputs["dp1_"], - self._lev, - ) - - return inputs \ No newline at end of file + self.grid.npz, + self._not_exit_loop, + self._INDEX_LM1, + self._INDEX_LP0, + inputs["q1"], + inputs["pe1_"], + inputs["pe2_"], + inputs["q4_1"], + inputs["q4_2"], + inputs["q4_3"], + inputs["q4_4"], + inputs["dp1_"], + self._lev, + ) + + return inputs diff --git a/tests/savepoint/translate/translate_map1_ppm_W.py b/tests/savepoint/translate/translate_map1_ppm_W.py index 1b6e544b..3948262c 100644 --- a/tests/savepoint/translate/translate_map1_ppm_W.py +++ b/tests/savepoint/translate/translate_map1_ppm_W.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid +from ndsl import Namelist, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.map_single import MapSingle + class TranslateMap1_PPM_W(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) @@ -13,29 +14,28 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.in_vars["data_vars"] = { "w_": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "pe1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, "pe2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, - "ws_":{ + "ws_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, }, - } self.in_vars["parameters"] = [ "kord_wz", @@ -43,15 +43,14 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.out_vars = { "w_": { - "kend": grid.npz-1, - }, - + "kend": grid.npz - 1, + }, } # mode / iv set to -2 from GEOS self.mode = -2 - self.dims=[X_DIM, Y_DIM, Z_DIM] + self.dims = [X_DIM, Y_DIM, Z_DIM] def compute_from_storage(self, inputs): self._compute_func = MapSingle( @@ -63,10 +62,10 @@ def compute_from_storage(self, inputs): ) self._compute_func( - inputs["w_"], - inputs["pe1_"], - inputs["pe2_"], - qs=inputs["ws_"], - interp=False, - ) + inputs["w_"], + inputs["pe1_"], + inputs["pe2_"], + qs=inputs["ws_"], + interp=False, + ) return inputs diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py index ebee7e1b..0da0167e 100644 --- a/tests/savepoint/translate/translate_map1_ppm_delz.py +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -1,24 +1,29 @@ -from ndsl import StencilFactory, Namelist, QuantityFactory -from ndsl.stencils.testing.grid import Grid +from gt4py.cartesian.gtscript import PARALLEL, computation, interval + +from ndsl import Namelist, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.dsl.typing import FloatField from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.map_single import MapSingle -from ndsl.dsl.typing import FloatField -from gt4py.cartesian.gtscript import FORWARD, PARALLEL, computation, interval + def rescale_delz_1( - delz: FloatField, - delp: FloatField, + delz: FloatField, + delp: FloatField, ): with computation(PARALLEL), interval(...): delz = -delz / delp + def rescale_delz_2( - delz: FloatField, - dp: FloatField, + delz: FloatField, + dp: FloatField, ): with computation(PARALLEL), interval(...): delz = -delz * dp + + class TranslateMap1_PPM_delz(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) @@ -28,37 +33,36 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.in_vars["data_vars"] = { "delz_": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "pe1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, "pe2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, "dp2_3d": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, - "gz_":{ + "gz_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, }, "delp": {}, - } self.in_vars["parameters"] = [ "kord_wz", @@ -66,15 +70,14 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.out_vars = { "delz_": { - "kend": grid.npz-1, - }, - + "kend": grid.npz - 1, + }, } # mode / iv set to 1 from GEOS self.mode = 1 - self.dims=[X_DIM, Y_DIM, Z_DIM] + self.dims = [X_DIM, Y_DIM, Z_DIM] self._rescale_delz_1 = stencil_factory.from_origin_domain( rescale_delz_1, @@ -103,13 +106,13 @@ def compute_from_storage(self, inputs): ) self._compute_func( - inputs["delz_"], - inputs["pe1_"], - inputs["pe2_"], - qs=inputs["gz_"], - interp=False, - ) - + inputs["delz_"], + inputs["pe1_"], + inputs["pe2_"], + qs=inputs["gz_"], + interp=False, + ) + self._rescale_delz_2( inputs["delz_"], inputs["dp2_3d"], diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py index 25f82a0f..bba68118 100644 --- a/tests/savepoint/translate/translate_map_scalar.py +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid +from ndsl import Namelist, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.map_single import MapSingle + class TranslateMap_Scalar(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) @@ -13,23 +14,22 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.in_vars["data_vars"] = { "q1": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "pe1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, "pe2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz + "kend": grid.npz, }, - } self.in_vars["parameters"] = [ "q_min", @@ -37,18 +37,17 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.out_vars = { "q1": { - "kend": grid.npz-1, - }, - + "kend": grid.npz - 1, + }, } # Value from GEOS - self._kord_tm = 9 + self._kord_tm = 9 # mode / iv set to 1 from GEOS - self.mode = 1 + self.mode = 1 - self.dims=[X_DIM, Y_DIM, Z_DIM] + self.dims = [X_DIM, Y_DIM, Z_DIM] self._compute_func = MapSingle( self.stencil_factory, @@ -60,10 +59,10 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto def compute_from_storage(self, inputs): self._compute_func( - inputs["q1"], - inputs["pe1_"], - inputs["pe2_"], - qmin=inputs["q_min"], - interp=True, - ) + inputs["q1"], + inputs["pe1_"], + inputs["pe2_"], + qmin=inputs["q_min"], + interp=True, + ) return inputs diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d.py b/tests/savepoint/translate/translate_moistcvpluspt_2d.py index 6e526ea3..3a0748cf 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d.py @@ -1,7 +1,7 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField, Float +from ndsl.dsl.typing import Float, FloatField from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j from pyFV3.stencils import moist_cv diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py index ec1728ee..46cf6402 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py @@ -1,32 +1,30 @@ -from gt4py.cartesian.gtscript import PARALLEL, computation, interval - -from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField, Float -from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j +from ndsl.dsl.typing import Float +from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils import moist_cv + class TranslateMoistCVPlusPt_2d_last_step(TranslateFortranData2Py): def __init__(self, grid, namelist, stencil_factory): super().__init__(grid, namelist, stencil_factory) self.stencil_factory = stencil_factory self.in_vars["data_vars"] = { "qvapor": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qliquid": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "qice": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qrain": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsnow": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qgraupel": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "pt": {}, "pkz": {}, @@ -50,21 +48,23 @@ def __init__(self, grid, namelist, stencil_factory): grid.nid, grid.njd, grid.npz, - ), dtype=Float, + ), + dtype=Float, ) def compute_from_storage(self, inputs): - self.compute_func(inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - self._gz, - inputs["pt"], - inputs["pkz"], - inputs["dtmp"], - inputs["r_vir"], - ) + self.compute_func( + inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + self._gz, + inputs["pt"], + inputs["pkz"], + inputs["dtmp"], + inputs["r_vir"], + ) return inputs diff --git a/tests/savepoint/translate/translate_moistcvpluste_2d.py b/tests/savepoint/translate/translate_moistcvpluste_2d.py index 57c39805..8723bf3a 100644 --- a/tests/savepoint/translate/translate_moistcvpluste_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluste_2d.py @@ -1,6 +1,7 @@ from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j from pyFV3.stencils import moist_cv + class TranslateMoistCVPlusTe_2d(TranslateFortranData2Py): def __init__(self, grid, namelist, stencil_factory): super().__init__(grid, namelist, stencil_factory) @@ -19,7 +20,7 @@ def __init__(self, grid, namelist, stencil_factory): "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, + "kend": grid.npz + 1, }, "te_2d_": { "istart": grid.is_, @@ -29,9 +30,9 @@ def __init__(self, grid, namelist, stencil_factory): }, "u": { "istart": grid.isd, - "iend": grid.ied , + "iend": grid.ied, "jstart": grid.jsd, - "jend": grid.jed+1, + "jend": grid.jed + 1, "kend": grid.npz, }, "v": { @@ -41,7 +42,7 @@ def __init__(self, grid, namelist, stencil_factory): "jend": grid.jed, "kend": grid.npz, }, - "w":{ + "w": { "kend": grid.npz, }, "cosa_s": { @@ -62,7 +63,7 @@ def __init__(self, grid, namelist, stencil_factory): "jstart": grid.jsd, "jend": grid.jed, }, - "delz": {} + "delz": {}, } self.write_vars = ["qvapor", "qliquid", "qice", "qrain", "qsnow", "qgraupel"] for k, v in self.in_vars["data_vars"].items(): @@ -82,7 +83,7 @@ def __init__(self, grid, namelist, stencil_factory): self.compute_func = stencil_factory.from_origin_domain( moist_cv.moist_te, origin=grid.compute_origin(), - domain=(grid.nic, 1, grid.npz+1), + domain=(grid.nic, 1, grid.npz + 1), ) def compute_from_storage(self, inputs): @@ -94,24 +95,25 @@ def compute_from_storage(self, inputs): ) ) - self.compute_func(inputs["qvapor"], - inputs["qliquid"], - inputs["qrain"], - inputs["qsnow"], - inputs["qice"], - inputs["qgraupel"], - inputs["u"], - inputs["v"], - inputs["w"], - inputs["te_2d_"], - inputs["pt"], - inputs["phis_"], - inputs["delp"], - inputs["rsin2"], - inputs["cosa_s"], - inputs["hs"], - inputs["delz"], - inputs["grav"], - ) + self.compute_func( + inputs["qvapor"], + inputs["qliquid"], + inputs["qrain"], + inputs["qsnow"], + inputs["qice"], + inputs["qgraupel"], + inputs["u"], + inputs["v"], + inputs["w"], + inputs["te_2d_"], + inputs["pt"], + inputs["phis_"], + inputs["delp"], + inputs["rsin2"], + inputs["cosa_s"], + inputs["hs"], + inputs["delz"], + inputs["grav"], + ) return inputs diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index 6611daa7..7d495c01 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -1,10 +1,16 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid -from ndsl.stencils.testing import TranslateFortranData2Py, ParallelTranslate import numpy as np -from ndsl.typing import Communicator + +from ndsl import Namelist, StencilFactory +from ndsl.comm.comm_abc import ReductionOperator from ndsl.quantity import Quantity -from pyFV3.stencils.mpp_global_sum import mpp_global_sum +from ndsl.stencils.testing import ParallelTranslate +from ndsl.stencils.testing.grid import Grid +from ndsl.typing import Communicator + + +# from pyFV3.stencils.mpp_global_sum import mpp_global_sum + + class TranslateMpp_global_sum(ParallelTranslate): def __init__( self, @@ -17,25 +23,45 @@ def __init__( self.grid = grid self._base.in_vars["data_vars"] = { - "inputArray" :{ + "inputArray": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, }, - "tesum":{ - - } + "tesum": {}, } - self._base.out_vars = { - "tesum":{ + self._base.out_vars = {"tesum": {}} - } - } - def compute_parallel(self, inputs, communicator: Communicator): - inputs["tesum"] = mpp_global_sum(inputs["inputArray"], communicator, self.stencil_factory) - - return inputs \ No newline at end of file + # inputs["tesum"] = mpp_global_sum(inputs["inputArray"], + # communicator, + # self.stencil_factory) + + ints_sum_reduce = Quantity( + data=np.zeros(inputs["inputArray"].data.shape, dtype=np.float32), + dims=["I", "J"], + units="dunno", + gt4py_backend=self.stencil_factory.backend, + ) + + inputArray = Quantity( + data=inputs["inputArray"].astype(np.float32), + dims=["I", "J"], + units="dunno", + gt4py_backend=self.stencil_factory.backend, + ) + + # print("ints_sum_reduce shape: ", ints_sum_reduce.data.shape) + + # print("inputs[inputArray] type: ", type(inputs["inputArray"])) + # print("ints_sum_reduce type: ", type(ints_sum_reduce)) + + communicator.all_reduce(inputArray, ReductionOperator.SUM, ints_sum_reduce) + + inputs["tesum"] = sum(sum(ints_sum_reduce.data)) + # print("inputs[tesum]: ", inputs["tesum"]) + + return inputs diff --git a/tests/savepoint/translate/translate_pe_pk_delp_peln.py b/tests/savepoint/translate/translate_pe_pk_delp_peln.py index 17fea0da..89dab68c 100644 --- a/tests/savepoint/translate/translate_pe_pk_delp_peln.py +++ b/tests/savepoint/translate/translate_pe_pk_delp_peln.py @@ -1,7 +1,8 @@ -from ndsl import StencilFactory, Namelist +from ndsl import Namelist, StencilFactory from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils.remapping import pe_pk_delp_peln from ndsl.stencils.testing.grid import Grid +from pyFV3.stencils.remapping import pe_pk_delp_peln + class TranslatePE_pk_delp_peln(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): @@ -13,39 +14,39 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.in_vars["data_vars"] = { "pe2_": { "istart": grid.is_, - "iend": grid.ie, + "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "pe_": { - "istart": grid.is_-1, - "iend": grid.ie+1, - "jstart": grid.js-1, - "jend": grid.je+1, - "kend": grid.npz+1, - }, + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, + "kend": grid.npz + 1, + }, "peln_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "pn2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "pk2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "delp": { # "istart": grid.isd, # "iend": grid.ied, @@ -58,12 +59,10 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, - "ak":{ - }, - "bk":{ + "kend": grid.npz + 1, }, + "ak": {}, + "bk": {}, } self.in_vars["parameters"] = [ "akap", @@ -76,36 +75,36 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "pe_": { - "istart": grid.is_-1, - "iend": grid.ie+1, - "jstart": grid.js-1, - "jend": grid.je+1, - "kend": grid.npz+1, - }, + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, + "kend": grid.npz + 1, + }, "peln_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "pn2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "pk2_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "delp": { # "istart": grid.isd, # "iend": grid.ied, @@ -118,8 +117,8 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, } grid_indexing = stencil_factory.grid_indexing @@ -129,7 +128,6 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto grid_indexing.domain[2] + 1, ) - self._pe_pk_delp_peln = stencil_factory.from_origin_domain( pe_pk_delp_peln, origin=grid_indexing.origin_compute(), @@ -138,16 +136,17 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto def compute_from_storage(self, inputs): - self._pe_pk_delp_peln(inputs["pe_"], - inputs["pk"], - inputs["delp"], - inputs["peln_"], - inputs["pe2_"], - inputs["pk2_"], - inputs["pn2_"], - inputs["ak"], - inputs["bk"], - inputs["akap"], - inputs["ptop"], + self._pe_pk_delp_peln( + inputs["pe_"], + inputs["pk"], + inputs["delp"], + inputs["peln_"], + inputs["pe2_"], + inputs["pk2_"], + inputs["pn2_"], + inputs["ak"], + inputs["bk"], + inputs["akap"], + inputs["ptop"], ) return inputs diff --git a/tests/savepoint/translate/translate_pn2_pk_delp.py b/tests/savepoint/translate/translate_pn2_pk_delp.py index 543c83ee..ba29adb4 100644 --- a/tests/savepoint/translate/translate_pn2_pk_delp.py +++ b/tests/savepoint/translate/translate_pn2_pk_delp.py @@ -1,12 +1,9 @@ -from gt4py.cartesian.gtscript import PARALLEL, computation, interval - from ndsl import StencilFactory -from ndsl.dsl.typing import FloatField, Float +from ndsl.dsl.typing import Float, FloatField from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils.remapping import pn2_pk_delp - class testClass: """ Class to test with DaCe orchestration. test class is MoistCVPlusPt_2d @@ -19,8 +16,8 @@ def __init__( ): self._pn2_pk_delp = stencil_factory.from_origin_domain( func=pn2_pk_delp, - origin=(3,3,1), - domain=(24,24,71), + origin=(3, 3, 1), + domain=(24, 24, 71), ) def __call__( @@ -32,14 +29,7 @@ def __call__( pk: FloatField, akap: Float, ): - self._pn2_pk_delp( - dp2, - delp, - pe2, - pn2, - pk, - akap - ) + self._pn2_pk_delp(dp2, delp, pe2, pn2, pk, akap) class TranslatePN2_PK_DelP(TranslateFortranData2Py): @@ -51,51 +41,62 @@ def __init__(self, grid, namelist, stencil_factory): self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { - "pe2_": {"istart": grid.is_, + "pe2_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, - "pn2_": {"istart": grid.is_, + "kend": grid.npz + 1, + }, + "pn2_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1,}, - "pk_": {"istart": grid.is_, + "kend": grid.npz + 1, + }, + "pk_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1,}, - + "kend": grid.npz + 1, + }, } self.in_vars["parameters"] = [ "akap", ] self.out_vars = { - "pe2_": {"istart": grid.is_, + "pe2_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1,}, - "pn2_": {"istart": grid.is_, + "kend": grid.npz + 1, + }, + "pn2_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1,}, - "pk_": {"istart": grid.is_, + "kend": grid.npz + 1, + }, + "pk_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1,}, + "kend": grid.npz + 1, + }, } self._dp2 = self.quantity_factory._numpy.zeros( ( 31, 31, 73, - ), dtype=Float, + ), + dtype=Float, ) self._delp = self.quantity_factory._numpy.zeros( @@ -103,10 +104,10 @@ def __init__(self, grid, namelist, stencil_factory): 31, 31, 73, - ), dtype=Float, + ), + dtype=Float, ) - def compute_from_storage(self, inputs): # print("delp shape = ", self._delp.shape) @@ -119,10 +120,12 @@ def compute_from_storage(self, inputs): # print('self.storage_vars() = ', self.storage_vars()) # self.make_storage_data_input_vars(inputs) # exit(1) - self.compute_func(self._dp2, - self._delp, - inputs["pe2_"], - inputs["pn2_"], - inputs["pk_"], - inputs["akap"]) + self.compute_func( + self._dp2, + self._delp, + inputs["pe2_"], + inputs["pn2_"], + inputs["pk_"], + inputs["akap"], + ) return inputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index d2d5a4d5..ee59ce5e 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,25 +1,26 @@ +from types import SimpleNamespace + from ndsl import Namelist, StencilFactory +from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.dsl.typing import Float +from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig -from pyFV3.stencils.remapping import init_pe, moist_cv_pt_pressure, pn2_pk_delp -from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils import moist_cv -from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 -from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment -from pyFV3.stencils.remapping import pressures_mapu, pe0_ptop_xmax, pressures_mapv, pe_pk_delp_peln +from pyFV3.stencils.map_single import MapSingle +from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.mpp_global_sum import mpp_global_sum -from ndsl.stencils.testing import Grid -from ndsl.stencils.testing import ParallelTranslateBaseSlicing -from ndsl.constants import ( - X_DIM, - X_INTERFACE_DIM, - Y_DIM, - Y_INTERFACE_DIM, - Z_DIM, - Z_INTERFACE_DIM, +from pyFV3.stencils.remapping import ( + init_pe, + moist_cv_pt_pressure, + pe0_ptop_xmax, + pe_pk_delp_peln, + pn2_pk_delp, + pressures_mapu, + pressures_mapv, ) -from ndsl.dsl.typing import Float -from pyFV3.stencils.mapn_tracer import MapNTracer -from types import SimpleNamespace +from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 +from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment + class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): inputs = { @@ -321,6 +322,7 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "units": "No Units", }, } + def __init__( self, grid: Grid, @@ -328,42 +330,41 @@ def __init__( stencil_factory: StencilFactory, ): super().__init__(grid, namelist, stencil_factory) - + self._base.in_vars["data_vars"] = { - "pe_": { - "istart": grid.is_-1, - "iend": grid.ie+1, - "jstart": grid.js-1, - "jend": grid.je+1, + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, "kend": grid.npz + 1, }, "qvapor": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qliquid": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "qice": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qrain": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsnow": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qgraupel": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qcld": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qo3mr": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsgs_tke": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "delp": {}, "delz": {}, @@ -394,9 +395,9 @@ def __init__( "jend": grid.je, }, "w": { - "kend": grid.npz-1, - }, - "ws_":{ + "kend": grid.npz - 1, + }, + "ws_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, @@ -406,12 +407,12 @@ def __init__( "istart": grid.isd, "iend": grid.ied, "jstart": grid.jsd, - "jend": grid.jed+1, + "jend": grid.jed + 1, "kend": grid.npz, }, "v": { "istart": grid.isd, - "iend": grid.ied+1, + "iend": grid.ied + 1, "jstart": grid.jsd, "jend": grid.jed, "kend": grid.npz, @@ -420,31 +421,29 @@ def __init__( "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, - "cy": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, "mfx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "cx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "cosa_s": { "istart": grid.isd, @@ -480,7 +479,6 @@ def __init__( self._base.in_vars["parameters"] = [ "ptop", "r_vir", - # "remap_t", # For some reason, translate test can't accept a logical variable "akap", "t_min", "kord_wz", @@ -502,81 +500,77 @@ def __init__( "delp": {}, "delz": {}, "qvapor": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qliquid": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "qice": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qrain": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsnow": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qgraupel": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qcld": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qo3mr": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "qsgs_tke": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "w": { - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "u": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.jsd, - "jend": grid.jed+1, - "kend": grid.npz-1, + "jend": grid.jed + 1, + "kend": grid.npz - 1, }, "v": { "istart": grid.isd, - "iend": grid.ied+1, + "iend": grid.ied + 1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "mfy": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, - "cy": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.js, - "jend": grid.je+1, - "kend": grid.npz-1, + "jend": grid.je + 1, + "kend": grid.npz - 1, }, "mfx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "cx_": { "istart": grid.is_, - "iend": grid.ie+1, + "iend": grid.ie + 1, "jstart": grid.jsd, "jend": grid.jed, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "peln_3d": { "istart": grid.is_, "iend": grid.ie, @@ -584,22 +578,20 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, - "pe_": { - "istart": grid.is_-1, - "iend": grid.ie+1, - "jstart": grid.js-1, - "jend": grid.je+1, + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, "kend": grid.npz + 1, }, - "pk": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz+1, - }, + "kend": grid.npz + 1, + }, "pkz": { "istart": grid.is_, "iend": grid.ie, @@ -613,7 +605,7 @@ def __init__( # self.namelist found in TranslateDycoreFortranData2Py # self.namelist = DynamicalCoreConfig.from_namelist(namelist) - config=DynamicalCoreConfig.from_namelist(self.namelist).remapping + config = DynamicalCoreConfig.from_namelist(self.namelist).remapping hydrostatic = config.hydrostatic if hydrostatic: @@ -625,15 +617,15 @@ def __init__( grid_indexing.domain[1] + 1, grid_indexing.domain[2] + 1, ) - + # Value from GEOS - self.kord = 9 + self.kord = 9 # Value from GEOS - self._kord_tm = 9 + self._kord_tm = 9 # mode / iv set to 1 from GEOS - self.mode = 1 + self.mode = 1 self.nq = 9 @@ -643,7 +635,8 @@ def __init__( ( grid.nid, grid.njd, - ), dtype=Float, + ), + dtype=Float, ) self._w2 = self.quantity_factory._numpy.zeros( @@ -651,113 +644,125 @@ def __init__( grid.nid, grid.njd, grid.npz, - ), dtype=Float, + ), + dtype=Float, ) self._zsum1 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - ), dtype=Float, + ), + dtype=Float, ) self._compute_performed = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - ), dtype=bool, + ), + dtype=bool, ) self._ps = self._pe1 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - ), dtype=Float, + ), + dtype=Float, ) self._pe0 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._pe1 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._pe2 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._pe3 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._pn1 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._pn2 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._dp2 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, - + grid.npz + 1, + ), + dtype=Float, ) self._pk2 = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._phis = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - grid.npz+1, - ), dtype=Float, + grid.npz + 1, + ), + dtype=Float, ) self._te_2d = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - ), dtype=Float, + ), + dtype=Float, ) - self._init_pe = stencil_factory.from_origin_domain( - init_pe, + init_pe, origin=grid_indexing.origin_compute(), domain=(grid.nic, grid.njc, grid.npz + 1), ) @@ -768,13 +773,17 @@ def __init__( externals={"hydrostatic": hydrostatic}, origin=grid_indexing.origin_compute(), # domain=grid_indexing.domain_compute(add=(0, 0, 1)), - domain=(grid_indexing.domain[0], grid_indexing.domain[1], grid_indexing.domain[2]+1), # Note : Many intervals go from (0,-1) in this stencil + domain=( + grid_indexing.domain[0], + grid_indexing.domain[1], + grid_indexing.domain[2] + 1, + ), # Note : Many intervals go from (0,-1) in this stencil ) self._pn2_pk_delp = stencil_factory.from_origin_domain( pn2_pk_delp, - origin=grid_indexing.origin_compute(add=(0,0,1)), - domain=(grid.nic, grid.njc, grid.npz-1), + origin=grid_indexing.origin_compute(add=(0, 0, 1)), + domain=(grid.nic, grid.njc, grid.npz - 1), ) self._map_scalar = MapSingle( @@ -806,13 +815,21 @@ def __init__( self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[0],grid_indexing.domain[1]+1,grid_indexing.domain[2] + 1) + domain=( + grid_indexing.domain[0], + grid_indexing.domain[1] + 1, + grid_indexing.domain[2] + 1, + ), ) self._pe0_ptop_xmax = stencil_factory.from_origin_domain( pe0_ptop_xmax, - origin=(grid_indexing.n_halo+grid_indexing.domain[0],grid_indexing.n_halo,0), - domain=(1,grid_indexing.domain[1]+1, 1) + origin=( + grid_indexing.n_halo + grid_indexing.domain[0], + grid_indexing.n_halo, + 0, + ), + domain=(1, grid_indexing.domain[1] + 1, 1), ) self._pressures_mapv = stencil_factory.from_origin_domain( @@ -828,7 +845,10 @@ def __init__( self._pe_pk_delp_peln = stencil_factory.from_origin_domain( pe_pk_delp_peln, origin=grid_indexing.origin_compute(), - domain=(grid_indexing.domain[0], grid_indexing.domain[1], grid_indexing.domain[2] + 1, + domain=( + grid_indexing.domain[0], + grid_indexing.domain[1], + grid_indexing.domain[2] + 1, ), ) @@ -841,7 +861,7 @@ def __init__( self._moist_cv_te = stencil_factory.from_origin_domain( moist_cv.moist_te, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz+1), + domain=(grid.nic, grid.njc, grid.npz + 1), ) self._te_zsum = stencil_factory.from_origin_domain( @@ -909,22 +929,23 @@ def compute_parallel(self, inputs, communicator): ) self._map_scalar( - state_namespace.pt, - self._pn1, - self._pn2, - qmin=state_namespace.t_min, - interp=True, + state_namespace.pt, + self._pn1, + self._pn2, + qmin=state_namespace.t_min, + interp=True, ) - tracers = { "qvapor": state_namespace.qvapor, - "qliquid": state_namespace.qliquid, - "qice": state_namespace.qice, - "qrain": state_namespace.qrain, - "qsnow": state_namespace.qsnow, - "qgraupel": state_namespace.qgraupel, - "qcld": state_namespace.qcld, - "qo3mr": state_namespace.qo3mr, - "qsgs_tke": state_namespace.qsgs_tke, + tracers = { + "qvapor": state_namespace.qvapor, + "qliquid": state_namespace.qliquid, + "qice": state_namespace.qice, + "qrain": state_namespace.qrain, + "qsnow": state_namespace.qsnow, + "qgraupel": state_namespace.qgraupel, + "qcld": state_namespace.qcld, + "qo3mr": state_namespace.qo3mr, + "qsgs_tke": state_namespace.qsgs_tke, } self._mapn_tracer = MapNTracer( @@ -936,10 +957,7 @@ def compute_parallel(self, inputs, communicator): tracers=tracers, ) - self._mapn_tracer(self._pe1, - self._pe2, - self._dp2, - tracers) + self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) self._map_single_w = MapSingle( self.stencil_factory, @@ -956,35 +974,35 @@ def compute_parallel(self, inputs, communicator): 1, dims=[X_DIM, Y_DIM, Z_DIM], ) - self._map_single_w(state_namespace.w, - self._pe1, - self._pe2, - qs=state_namespace.ws_, - interp=False) - + self._map_single_w( + state_namespace.w, + self._pe1, + self._pe2, + qs=state_namespace.ws_, + interp=False, + ) + self._rescale_delz_1( state_namespace.delz, state_namespace.delp, ) - - self._map_single_delz(state_namespace.delz, - self._pe1, - self._pe2) + + self._map_single_delz(state_namespace.delz, self._pe1, self._pe2) self._rescale_delz_2( state_namespace.delz, self._dp2, ) - + self._w_fix_consrv_moment( - state_namespace.w, - self._w2, - self._dp2, - self._gz, - state_namespace.w_max, - state_namespace.w_min, - self._compute_performed - ) + state_namespace.w, + self._w2, + self._dp2, + self._gz, + state_namespace.w_max, + state_namespace.w_min, + self._compute_performed, + ) self._map1_ppm_u = MapSingle( self.stencil_factory, @@ -995,40 +1013,40 @@ def compute_parallel(self, inputs, communicator): ) self._pressures_mapu( - state_namespace.pe_, - state_namespace.ak, - state_namespace.bk, - self._pe0, - self._pe3, - state_namespace.ptop, - ) - + state_namespace.pe_, + state_namespace.ak, + state_namespace.bk, + self._pe0, + self._pe3, + state_namespace.ptop, + ) + self._pe0_ptop_xmax( - self._pe0, - state_namespace.ptop, - ) + self._pe0, + state_namespace.ptop, + ) self._map1_ppm_u( - state_namespace.u, - self._pe0, - self._pe3, - interp=False, - ) - + state_namespace.u, + self._pe0, + self._pe3, + interp=False, + ) + self._map1_ppm_u( - state_namespace.mfy, - self._pe0, - self._pe3, - interp=False, - ) - + state_namespace.mfy, + self._pe0, + self._pe3, + interp=False, + ) + self._map1_ppm_u( - state_namespace.cy, - self._pe0, - self._pe3, - interp=False, - ) - + state_namespace.cy, + self._pe0, + self._pe3, + interp=False, + ) + self._map1_ppm_v = MapSingle( self.stencil_factory, self.quantity_factory, @@ -1038,45 +1056,46 @@ def compute_parallel(self, inputs, communicator): ) self._pressures_mapv( - state_namespace.pe_, - state_namespace.ak, - state_namespace.bk, - self._pe0, - self._pe3, - ) + state_namespace.pe_, + state_namespace.ak, + state_namespace.bk, + self._pe0, + self._pe3, + ) self._map1_ppm_v( - state_namespace.v, - self._pe0, - self._pe3, - interp=False, - ) - + state_namespace.v, + self._pe0, + self._pe3, + interp=False, + ) + self._map1_ppm_v( - state_namespace.mfx_, - self._pe0, - self._pe3, - interp=False, - ) - + state_namespace.mfx_, + self._pe0, + self._pe3, + interp=False, + ) + self._map1_ppm_v( - state_namespace.cx_, - self._pe0, - self._pe3, - interp=False, - ) + state_namespace.cx_, + self._pe0, + self._pe3, + interp=False, + ) - self._pe_pk_delp_peln(state_namespace.pe_, - state_namespace.pk, - state_namespace.delp, - state_namespace.peln_3d, - self._pe2, - self._pk2, - self._pn2, - state_namespace.ak, - state_namespace.bk, - state_namespace.akap, - state_namespace.ptop, + self._pe_pk_delp_peln( + state_namespace.pe_, + state_namespace.pk, + state_namespace.delp, + state_namespace.peln_3d, + self._pe2, + self._pk2, + self._pn2, + state_namespace.ak, + state_namespace.bk, + state_namespace.akap, + state_namespace.ptop, ) self._moist_cv_pkz( @@ -1093,70 +1112,80 @@ def compute_parallel(self, inputs, communicator): state_namespace.delz, Float(state_namespace.r_vir), ) - + if state_namespace.last_step and not state_namespace.do_adiabatic_init: - + if state_namespace.consv > state_namespace.consv_min: - self._moist_cv_te(state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, - state_namespace.u, - state_namespace.v, - state_namespace.w, - self._te_2d, - state_namespace.pt, - self._phis, - state_namespace.delp, - state_namespace.rsin2, - state_namespace.cosa_s, - state_namespace.hs, - state_namespace.delz, - state_namespace.grav, - ) - - self._te_zsum(self._te_2d, - state_namespace.te0_2d_, - state_namespace.delp, - state_namespace.pkz, - self._zsum1, - ) - + self._moist_cv_te( + state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.u, + state_namespace.v, + state_namespace.w, + self._te_2d, + state_namespace.pt, + self._phis, + state_namespace.delp, + state_namespace.rsin2, + state_namespace.cosa_s, + state_namespace.hs, + state_namespace.delz, + state_namespace.grav, + ) + + self._te_zsum( + self._te_2d, + state_namespace.te0_2d_, + state_namespace.delp, + state_namespace.pkz, + self._zsum1, + ) + # Note, mpp_global_sum is currently set up for the C24 TBC setup - inputArray = self._te_2d.data*state_namespace.area_64_.data[0:-1,0:-1] - tesum = mpp_global_sum(inputArray[3:27,3:27], communicator, self.stencil_factory) + inputArray = ( + self._te_2d.data * state_namespace.area_64_.data[0:-1, 0:-1] + ) + tesum = mpp_global_sum( + inputArray[3:27, 3:27], communicator, self.stencil_factory + ) # print("tesum: ", tesum) - inputArray = self._zsum1*state_namespace.area_64_.data[0:-1,0:-1] - zsum = mpp_global_sum(inputArray[3:27,3:27], communicator, self.stencil_factory) + inputArray = self._zsum1 * state_namespace.area_64_.data[0:-1, 0:-1] + zsum = mpp_global_sum( + inputArray[3:27, 3:27], communicator, self.stencil_factory + ) # print("zsum: ", zsum) dtmp = tesum / (state_namespace.cv_air.data * zsum) # print("dtmp: ", dtmp) - # I ignore the E_flux computation since it's not used elsewhere in our current setup once it's computed - + # I ignore the E_flux computation since it's not used elsewhere + # in our current setup once it's computed if state_namespace.last_step and not state_namespace.adiabatic: - self._most_cv_pt_last_step(state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, - state_namespace.pt, - state_namespace.pkz, - Float(dtmp), - state_namespace.r_vir, - ) - - self._fill_cond(state_namespace.q_con, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, + self._most_cv_pt_last_step( + state_namespace.qvapor, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, + state_namespace.pt, + state_namespace.pkz, + Float(dtmp), + state_namespace.r_vir, + ) + + self._fill_cond( + state_namespace.q_con, + state_namespace.qliquid, + state_namespace.qrain, + state_namespace.qsnow, + state_namespace.qice, + state_namespace.qgraupel, ) - return self.outputs_from_state(state) \ No newline at end of file + return self.outputs_from_state(state) diff --git a/tests/savepoint/translate/translate_scalar_profile.py b/tests/savepoint/translate/translate_scalar_profile.py index b0779d34..12c08355 100644 --- a/tests/savepoint/translate/translate_scalar_profile.py +++ b/tests/savepoint/translate/translate_scalar_profile.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory, Namelist -from ndsl.stencils.testing.grid import Grid +from ndsl import Namelist, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.remap_profile import RemapProfile + class TranslateScalar_Profile(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): super().__init__(grid, stencil_factory) @@ -12,46 +13,47 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { - "qs_": {"istart": grid.is_, + "qs_": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_1": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "q4_2": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_3": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, "q4_4": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "dp1_":{ + "dp1_": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, } self.in_vars["parameters"] = [ @@ -59,54 +61,58 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto ] self.out_vars = { - "q4_1": {"istart": grid.is_, + "q4_1": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "q4_2": {"istart": grid.is_, + "q4_2": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "q4_3": {"istart": grid.is_, + "q4_3": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, - "q4_4": {"istart": grid.is_, + "q4_4": { + "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz-1, + "kend": grid.npz - 1, }, } # Value from GEOS - self.kord = 9 + self.kord = 9 # mode / iv set to 1 from GEOS - self.mode = 1 + self.mode = 1 self._compute_func = RemapProfile( self.stencil_factory, self.quantity_factory, self.kord, self.mode, - dims=[X_DIM, Y_DIM, Z_DIM] + dims=[X_DIM, Y_DIM, Z_DIM], ) def compute_from_storage(self, inputs): self._compute_func( - inputs["qs_"], - inputs["q4_1"], - inputs["q4_2"], - inputs["q4_3"], - inputs["q4_4"], - inputs["dp1_"], - inputs["q_min"], - ) + inputs["qs_"], + inputs["q4_1"], + inputs["q4_2"], + inputs["q4_3"], + inputs["q4_4"], + inputs["dp1_"], + inputs["q_min"], + ) return inputs diff --git a/tests/savepoint/translate/translate_te_zsum.py b/tests/savepoint/translate/translate_te_zsum.py index 6ff0e028..2b39e1a5 100644 --- a/tests/savepoint/translate/translate_te_zsum.py +++ b/tests/savepoint/translate/translate_te_zsum.py @@ -1,6 +1,7 @@ from ndsl.stencils.testing import TranslateFortranData2Py from pyFV3.stencils import moist_cv + class TranslateTe_Zsum(TranslateFortranData2Py): def __init__(self, grid, namelist, stencil_factory): super().__init__(grid, namelist, stencil_factory) @@ -57,12 +58,13 @@ def __init__(self, grid, namelist, stencil_factory): ) def compute_from_storage(self, inputs): - - self.compute_func(inputs["te_2d_"], - inputs["te0_2d_"], - inputs["delp"], - inputs["pkz"], - inputs["zsum1"], - ) + + self.compute_func( + inputs["te_2d_"], + inputs["te0_2d_"], + inputs["delp"], + inputs["pkz"], + inputs["zsum1"], + ) return inputs diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py index 6a298bc2..6a642bc2 100644 --- a/tests/savepoint/translate/translate_w_fix_consrv_moment.py +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -1,8 +1,9 @@ from ndsl.dsl.typing import Float -from ndsl.stencils.testing.grid import Grid from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment + class TranslateW_fix_consrv_moment(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist, stencil_factory): super().__init__(grid, stencil_factory) @@ -18,8 +19,8 @@ def __init__(self, grid: Grid, namelist, stencil_factory): self.in_vars["data_vars"] = { "w": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, "dp2_W": grid.compute_dict(), } @@ -27,14 +28,15 @@ def __init__(self, grid: Grid, namelist, stencil_factory): self.out_vars = { "w": { - "kend": grid.npz-1, - }, + "kend": grid.npz - 1, + }, } self._gz = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - ), dtype=Float, + ), + dtype=Float, ) self._w2 = self.quantity_factory._numpy.zeros( @@ -42,26 +44,28 @@ def __init__(self, grid: Grid, namelist, stencil_factory): grid.nid, grid.njd, grid.npz, - ), dtype=Float, + ), + dtype=Float, ) self._compute_performed = self.quantity_factory._numpy.zeros( ( grid.nid, grid.njd, - ), dtype=bool, + ), + dtype=bool, ) def compute_from_storage(self, inputs): self.compute_func( - inputs["w"], - self._w2, - inputs["dp2_W"], - self._gz, - inputs["w_max"], - inputs["w_min"], - self._compute_performed - ) - + inputs["w"], + self._w2, + inputs["dp2_W"], + self._gz, + inputs["w_max"], + inputs["w_min"], + self._compute_performed, + ) + return inputs From 18a5218386da17c0b28ea66735be4594b773feeb Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 29 Jan 2025 07:59:36 -0800 Subject: [PATCH 153/252] Enabled mpp_global_sum with flag to choose 'precise' or 'simplified' reduction calculation --- pyFV3/stencils/mpp_global_sum.py | 107 +++++++++++------- .../translate/translate_mpp_global_sum.py | 38 +------ 2 files changed, 73 insertions(+), 72 deletions(-) diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/stencils/mpp_global_sum.py index 39f93518..e5ad9c5a 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/stencils/mpp_global_sum.py @@ -1,48 +1,77 @@ import numpy as np from ndsl.comm.comm_abc import ReductionOperator +from ndsl.dsl.typing import Float from ndsl.quantity import Quantity -def mpp_global_sum(inputArray, communicator, stencil_factory=None): - NUMINT = 6 - NUMBIT = 46 - r_prec = 2.0 ** NUMBIT - prec = 2 ** NUMBIT - I_prec = 1.0 / (2.0 ** NUMBIT) - pr = [r_prec ** 2, r_prec, 1.0, 1.0 / r_prec, 1.0 / r_prec ** 2, 1.0 / r_prec ** 3] - I_pr = [1.0 / r_prec ** 2, 1.0 / r_prec, 1.0, r_prec, r_prec ** 2, r_prec ** 3] - prec_error = (2 ** 62 + (2 ** 62 - 1)) / 6 - mag_max_term = 0.0 - - ints_sum = Quantity( - data=np.zeros((NUMINT), dtype=np.float32), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - - ints_sum_reduce = Quantity( - data=np.zeros((NUMINT), dtype=np.float32), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - - # Note: This loop range in i and j are for the TBC test case. - for j in range(inputArray.shape[1]): - for i in range(inputArray.shape[0]): - increment_ints_faster( - ints_sum.data, pr, I_pr, inputArray[i, j], mag_max_term - ) - - carry_overflow(ints_sum.data, prec, I_prec, prec_error) - - communicator.all_reduce(ints_sum, ReductionOperator.SUM, ints_sum_reduce) - - regularize_ints(ints_sum_reduce.data, prec, I_prec) - - sum_ = ints_to_real(ints_sum_reduce.data, pr) +def mpp_global_sum( + inputArray, communicator, stencil_factory=None, simplified_reduce=False +): + if not simplified_reduce: + NUMINT = 6 + NUMBIT = 46 + r_prec = 2.0 ** NUMBIT + prec = 2 ** NUMBIT + I_prec = 1.0 / (2.0 ** NUMBIT) + pr = [ + r_prec ** 2, + r_prec, + 1.0, + 1.0 / r_prec, + 1.0 / r_prec ** 2, + 1.0 / r_prec ** 3, + ] + I_pr = [1.0 / r_prec ** 2, 1.0 / r_prec, 1.0, r_prec, r_prec ** 2, r_prec ** 3] + prec_error = (2 ** 62 + (2 ** 62 - 1)) / 6 + mag_max_term = 0.0 + + ints_sum = Quantity( + data=np.zeros((NUMINT), dtype=Float), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + ints_sum_reduce = Quantity( + data=np.zeros((NUMINT), dtype=Float), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + # Note: This loop range in i and j are for the TBC test case. + for j in range(inputArray.shape[1]): + for i in range(inputArray.shape[0]): + increment_ints_faster( + ints_sum.data, pr, I_pr, inputArray[i, j], mag_max_term + ) + + carry_overflow(ints_sum.data, prec, I_prec, prec_error) + + communicator.all_reduce(ints_sum, ReductionOperator.SUM, ints_sum_reduce) + + regularize_ints(ints_sum_reduce.data, prec, I_prec) + + sum_ = ints_to_real(ints_sum_reduce.data, pr) + else: + inputArray_ = Quantity( + data=inputArray.astype(Float), + dims=["I", "J"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + ints_sum_reduce = Quantity( + data=np.zeros(inputArray.data.shape, dtype=Float), + dims=["I", "J"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + communicator.all_reduce(inputArray_, ReductionOperator.SUM, ints_sum_reduce) + + sum_ = sum(sum(ints_sum_reduce.data)) return sum_ diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index 7d495c01..0980b7f1 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -1,14 +1,8 @@ -import numpy as np - from ndsl import Namelist, StencilFactory -from ndsl.comm.comm_abc import ReductionOperator -from ndsl.quantity import Quantity from ndsl.stencils.testing import ParallelTranslate from ndsl.stencils.testing.grid import Grid from ndsl.typing import Communicator - - -# from pyFV3.stencils.mpp_global_sum import mpp_global_sum +from pyFV3.stencils.mpp_global_sum import mpp_global_sum class TranslateMpp_global_sum(ParallelTranslate): @@ -36,32 +30,10 @@ def __init__( def compute_parallel(self, inputs, communicator: Communicator): - # inputs["tesum"] = mpp_global_sum(inputs["inputArray"], - # communicator, - # self.stencil_factory) - - ints_sum_reduce = Quantity( - data=np.zeros(inputs["inputArray"].data.shape, dtype=np.float32), - dims=["I", "J"], - units="dunno", - gt4py_backend=self.stencil_factory.backend, - ) - - inputArray = Quantity( - data=inputs["inputArray"].astype(np.float32), - dims=["I", "J"], - units="dunno", - gt4py_backend=self.stencil_factory.backend, + inputs["tesum"] = mpp_global_sum( + inputs["inputArray"], + communicator, + self.stencil_factory, ) - # print("ints_sum_reduce shape: ", ints_sum_reduce.data.shape) - - # print("inputs[inputArray] type: ", type(inputs["inputArray"])) - # print("ints_sum_reduce type: ", type(ints_sum_reduce)) - - communicator.all_reduce(inputArray, ReductionOperator.SUM, ints_sum_reduce) - - inputs["tesum"] = sum(sum(ints_sum_reduce.data)) - # print("inputs[tesum]: ", inputs["tesum"]) - return inputs From dd4e8dec60350353f32158b06c90e691571d908a Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 30 Jan 2025 09:21:26 -0800 Subject: [PATCH 154/252] Small fixes to input and output dictionaries in translate test. Not sure if these actually make a difference in the computation but it does providing more details about each serialized variable --- .../translate/translate_remapping_GEOS.py | 33 +++++++++++-------- 1 file changed, 20 insertions(+), 13 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index ee59ce5e..2e3274aa 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,7 +1,14 @@ from types import SimpleNamespace from ndsl import Namelist, StencilFactory -from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.constants import ( + X_DIM, + X_INTERFACE_DIM, + Y_DIM, + Y_INTERFACE_DIM, + Z_DIM, + Z_INTERFACE_DIM, +) from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig @@ -26,7 +33,7 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): inputs = { "pe_": { "name": "pe_", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "qvapor": { @@ -106,22 +113,22 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): }, "peln_3d": { "name": "peln_3d", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "ak": { "name": "ak", - "dims": [X_DIM], + "dims": [Z_INTERFACE_DIM], "units": "No Units", }, "bk": { "name": "bk", - "dims": [X_DIM], + "dims": [Z_INTERFACE_DIM], "units": "No Units", }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "pkz": { @@ -141,12 +148,12 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): }, "u": { "name": "u", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], "units": "No Units", }, "v": { "name": "v", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], "units": "No Units", }, "mfy": { @@ -268,12 +275,12 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): }, "u": { "name": "u", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], "units": "No Units", }, "v": { "name": "v", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], "units": "No Units", }, "mfy": { @@ -298,17 +305,17 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): }, "peln_3d": { "name": "peln_3d", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "pe_": { "name": "pe_", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "pkz": { From 94844538b9602ac207549b479e589a3f5d6693cd Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 30 Jan 2025 09:59:31 -0800 Subject: [PATCH 155/252] Using domain_compute in domain within stencil_factory calls --- .../translate/translate_remapping_GEOS.py | 50 ++++++------------- 1 file changed, 16 insertions(+), 34 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 2e3274aa..4e7301f9 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -771,26 +771,20 @@ def __init__( self._init_pe = stencil_factory.from_origin_domain( init_pe, origin=grid_indexing.origin_compute(), - domain=(grid.nic, grid.njc, grid.npz + 1), + domain=grid_indexing.domain_compute(add=(0, 1, 1)), ) self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( moist_cv_pt_pressure, - # externals={"kord_tm": config.kord_tm, "hydrostatic": hydrostatic}, externals={"hydrostatic": hydrostatic}, origin=grid_indexing.origin_compute(), - # domain=grid_indexing.domain_compute(add=(0, 0, 1)), - domain=( - grid_indexing.domain[0], - grid_indexing.domain[1], - grid_indexing.domain[2] + 1, - ), # Note : Many intervals go from (0,-1) in this stencil + domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._pn2_pk_delp = stencil_factory.from_origin_domain( pn2_pk_delp, origin=grid_indexing.origin_compute(add=(0, 0, 1)), - domain=(grid.nic, grid.njc, grid.npz - 1), + domain=grid_indexing.domain_compute(add=(0, 0, -1)), ) self._map_scalar = MapSingle( @@ -804,29 +798,25 @@ def __init__( self._rescale_delz_1 = stencil_factory.from_origin_domain( rescale_delz_1, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz), + domain=grid_indexing.domain_compute(), ) self._rescale_delz_2 = stencil_factory.from_origin_domain( rescale_delz_2, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz), + domain=grid_indexing.domain_compute(), ) self._w_fix_consrv_moment = stencil_factory.from_origin_domain( func=W_fix_consrv_moment, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz), + domain=grid_indexing.domain_compute(), ) self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, origin=grid_indexing.origin_compute(), - domain=( - grid_indexing.domain[0], - grid_indexing.domain[1] + 1, - grid_indexing.domain[2] + 1, - ), + domain=grid_indexing.domain_compute(add=(0, 1, 1)), ) self._pe0_ptop_xmax = stencil_factory.from_origin_domain( @@ -842,51 +832,43 @@ def __init__( self._pressures_mapv = stencil_factory.from_origin_domain( pressures_mapv, origin=grid_indexing.origin_compute(), - domain=( - grid_indexing.domain[0] + 1, - grid_indexing.domain[1], - grid_indexing.domain[2] + 1, - ), + domain=grid_indexing.domain_compute(add=(1, 0, 1)), ) self._pe_pk_delp_peln = stencil_factory.from_origin_domain( pe_pk_delp_peln, origin=grid_indexing.origin_compute(), - domain=( - grid_indexing.domain[0], - grid_indexing.domain[1], - grid_indexing.domain[2] + 1, - ), + domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._moist_cv_pkz = stencil_factory.from_origin_domain( moist_cv.moist_pkz, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz), + domain=grid_indexing.domain_compute(), ) self._moist_cv_te = stencil_factory.from_origin_domain( moist_cv.moist_te, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz + 1), + domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._te_zsum = stencil_factory.from_origin_domain( moist_cv.te_zsum, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz), + domain=grid_indexing.domain_compute(), ) - self._most_cv_pt_last_step = stencil_factory.from_origin_domain( + self._moist_cv_pt_last_step = stencil_factory.from_origin_domain( moist_cv.moist_pt_last_step, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz), + domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._fill_cond = stencil_factory.from_origin_domain( moist_cv.cond_output, origin=grid.compute_origin(), - domain=(grid.nic, grid.njc, grid.npz), + domain=grid_indexing.domain_compute(), ) def compute_sequential(self, inputs_list, communicator_list): @@ -1173,7 +1155,7 @@ def compute_parallel(self, inputs, communicator): if state_namespace.last_step and not state_namespace.adiabatic: - self._most_cv_pt_last_step( + self._moist_cv_pt_last_step( state_namespace.qvapor, state_namespace.qliquid, state_namespace.qrain, From 830ae0b9fa3ebf13ffd3efa8a51835f83ba6ddb8 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 30 Jan 2025 12:22:05 -0800 Subject: [PATCH 156/252] Updated translate test with quantity factor to create temporary data --- .../translate/translate_remapping_GEOS.py | 138 +++++++----------- 1 file changed, 51 insertions(+), 87 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 4e7301f9..6bd724be 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -638,133 +638,93 @@ def __init__( self.fill = True - self._gz = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - ), + self._gz = self.quantity_factory.zeros( + [X_DIM, Y_DIM], + units="m^2 s^-2", dtype=Float, ) - self._w2 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz, - ), + self._w2 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="temp W", dtype=Float, ) - self._zsum1 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - ), + self._zsum1 = self.quantity_factory.zeros( + [X_DIM, Y_DIM], + units="Pa", dtype=Float, ) - self._compute_performed = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - ), + self._compute_performed = self.quantity_factory.zeros( + [X_DIM, Y_DIM], + units="mask", dtype=bool, ) - self._ps = self._pe1 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - ), + self._ps = self.quantity_factory.zeros( + [X_DIM, Y_DIM], + units="Pa", dtype=Float, ) - self._pe0 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._pe0 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", dtype=Float, ) - self._pe1 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._pe1 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", dtype=Float, ) - self._pe2 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._pe2 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", dtype=Float, ) - self._pe3 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._pe3 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="Pa", dtype=Float, ) - self._pn1 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._pn1 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", dtype=Float, ) - self._pn2 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._pn2 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", dtype=Float, ) - self._dp2 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._dp2 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", dtype=Float, ) - self._pk2 = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._pk2 = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="Pa", dtype=Float, ) - self._phis = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - grid.npz + 1, - ), + self._phis = self.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_INTERFACE_DIM], + units="n/a", dtype=Float, ) - self._te_2d = self.quantity_factory._numpy.zeros( - ( - grid.nid, - grid.njd, - ), + self._te_2d = self.quantity_factory.zeros( + [X_DIM, Y_DIM], + units="Pa", dtype=Float, ) @@ -1137,13 +1097,17 @@ def compute_parallel(self, inputs, communicator): # Note, mpp_global_sum is currently set up for the C24 TBC setup inputArray = ( - self._te_2d.data * state_namespace.area_64_.data[0:-1, 0:-1] + self._te_2d.data[0:-1, 0:-1] + * state_namespace.area_64_.data[0:-1, 0:-1] ) tesum = mpp_global_sum( inputArray[3:27, 3:27], communicator, self.stencil_factory ) # print("tesum: ", tesum) - inputArray = self._zsum1 * state_namespace.area_64_.data[0:-1, 0:-1] + inputArray = ( + self._zsum1.data[0:-1, 0:-1] + * state_namespace.area_64_.data[0:-1, 0:-1] + ) zsum = mpp_global_sum( inputArray[3:27, 3:27], communicator, self.stencil_factory ) From 76bfbecad6633de7e55c96afbeefe5a8fb4e060e Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 30 Jan 2025 12:45:57 -0800 Subject: [PATCH 157/252] Added updated fv_dynamics.F90.SER serialization file --- fv_dynamics.F90.SER | 1489 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 1489 insertions(+) create mode 100644 fv_dynamics.F90.SER diff --git a/fv_dynamics.F90.SER b/fv_dynamics.F90.SER new file mode 100644 index 00000000..fbe9da9c --- /dev/null +++ b/fv_dynamics.F90.SER @@ -0,0 +1,1489 @@ +!*********************************************************************** +!* GNU Lesser General Public License +!* +!* This file is part of the FV3 dynamical core. +!* +!* The FV3 dynamical core is free software: you can redistribute it +!* and/or modify it under the terms of the +!* GNU Lesser General Public License as published by the +!* Free Software Foundation, either version 3 of the License, or +!* (at your option) any later version. +!* +!* The FV3 dynamical core is distributed in the hope that it will be +!* useful, but WITHOUT ANYWARRANTY; without even the implied warranty +!* of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. +!* See the GNU General Public License for more details. +!* +!* You should have received a copy of the GNU Lesser General Public +!* License along with the FV3 dynamical core. +!* If not, see . +!*********************************************************************** + +!>@brief The module 'fv_dynamics' is the top-level routine for the dynamical core. +!>@details It executes on the dt_atmos (or p_split) loop. + +module fv_dynamics_mod + +! Modules Included: +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +! +!
Module NameFunctions Included
boundary_modnested_grid_BC_apply_intT
constants_modgrav, pi=>pi_8, radius, hlv, rdgas, omega, rvgas, cp_vapor
diag_manager_modsend_data
dyn_core_moddyn_core, del2_cubed, init_ijk_mem
field_manager_modMODEL_ATMOS
fv_arrays_modfv_grid_type, fv_flags_type, fv_atmos_type, fv_nest_type, fv_diag_type, fv_grid_bounds_type
fv_diagnostics_modfv_time, prt_mxm, range_check, prt_minmax
fv_fill_modfill2D
fv_grid_utils_modcubed_to_latlon, c2l_ord2, g_sum
fv_mapz_modcompute_total_energy, Lagrangian_to_Eulerian, moist_cv, moist_cp
fv_mp_modis_master,group_halo_update_type, start_group_halo_update, complete_group_halo_update
fv_nesting_modsetup_nested_grid_BCs
fv_nwp_moddo_adiabatic_init
fv_restart_modd2a_setup, d2c_setup
fv_sg_modneg_adj3
fv_timing_modtiming_on, timing_off
fv_tracer2d_modtracer_2d, tracer_2d_1L, tracer_2d_nested
mpp_mod/td> +! mpp_pe
mpp_domains_mod/td> +! DGRID_NE, CGRID_NE, mpp_update_domains, domain2D
fv_sg_modneg_adj3
tracer_manager_modget_tracer_index
+ + use constants_mod, only: grav, pi=>pi_8, radius, hlv, rdgas, omega, rvgas, cp_vapor + use dyn_core_mod, only: dyn_core, del2_cubed, init_ijk_mem + use fv_mapz_mod, only: compute_total_energy, Lagrangian_to_Eulerian, moist_cv, moist_cp + use fv_tracer2d_mod, only: tracer_2d, tracer_2d_1L, tracer_2d_nested + use fv_grid_utils_mod, only: cubed_to_latlon, c2l_ord2, g_sum + use fv_fill_mod, only: fill2D + use fv_mp_mod, only: is_master + use fv_mp_mod, only: group_halo_update_type + use fv_mp_mod, only: start_group_halo_update, complete_group_halo_update + use fv_timing_mod, only: timing_on, timing_off + use diag_manager_mod, only: send_data + use fv_diagnostics_mod, only: fv_time, prt_mxm, range_check, prt_minmax + use mpp_domains_mod, only: DGRID_NE, CGRID_NE, mpp_update_domains, domain2D + use mpp_mod, only: mpp_pe + use field_manager_mod, only: MODEL_ATMOS + use tracer_manager_mod, only: get_tracer_index + use fv_sg_mod, only: neg_adj3 + use fv_nesting_mod, only: setup_nested_grid_BCs + use boundary_mod, only: nested_grid_BC_apply_intT + use fv_arrays_mod, only: fv_grid_type, fv_flags_type, fv_atmos_type, fv_nest_type, fv_diag_type, fv_grid_bounds_type + use fv_nwp_nudge_mod, only: do_adiabatic_init + +implicit none + + logical :: RF_initialized = .false. + logical :: bad_range = .false. + real, allocatable :: rf(:) + integer :: kmax=1 + integer :: k_rf = 0 + + real :: agrav +#ifdef HIWPP + real, allocatable:: u00(:,:,:), v00(:,:,:) +#endif +private +public :: fv_dynamics + +contains + +!----------------------------------------------------------------------- +! fv_dynamics :: FV dynamical core driver +!----------------------------------------------------------------------- + + subroutine fv_dynamics(npx, npy, npz, nq_tot, ng, bdt, consv_te, fill, & + reproduce_sum, kappa, cp_air, zvir, ptop, ks, ncnst, & + k_split, n_split, & + q_split, u, v, w, delz, hydrostatic, pt, delp, q, & + ps, pe, pk, peln, pkz, phis, varflt, q_con, omga, ua, va, uc, vc, & + ak, bk, mfx, mfy, cx, cy, ze0, hybrid_z, & + gridstruct, flagstruct, neststruct, idiag, bd, & + parent_grid, domain, diss_est, & + dudt_rf, dvdt_rf, dwdt_rf, dtdt_rf, & + time_total) + + real, intent(IN) :: bdt !< Large time-step + real, intent(IN) :: consv_te + real, intent(IN) :: kappa, cp_air + real, intent(IN) :: zvir, ptop + real, intent(IN), optional :: time_total + + integer, intent(IN) :: npx + integer, intent(IN) :: npy + integer, intent(IN) :: npz + integer, intent(IN) :: nq_tot !< transported tracers + integer, intent(IN) :: ng + integer, intent(IN) :: ks + integer, intent(IN) :: ncnst + integer, intent(IN) :: k_split !< remap-step vertical dynamics + integer, intent(IN) :: n_split !< small-step horizontal dynamics + integer, intent(IN) :: q_split !< tracer + logical, intent(IN) :: fill + logical, intent(IN) :: reproduce_sum + logical, intent(IN) :: hydrostatic + logical, intent(IN) :: hybrid_z !< Using hybrid_z for remapping + + type(fv_grid_bounds_type), intent(IN) :: bd + real, intent(inout), dimension(bd%isd:bd%ied ,bd%jsd:bd%jed+1,npz) :: u !< D grid zonal wind (m/s) + real, intent(inout), dimension(bd%isd:bd%ied+1,bd%jsd:bd%jed ,npz) :: v !< D grid meridional wind (m/s) + real, intent(inout) :: w( bd%isd: ,bd%jsd: ,1:) !< W (m/s) + real, intent(inout) :: pt( bd%isd:bd%ied ,bd%jsd:bd%jed ,npz) !< temperature (K) + real, intent(inout) :: delp(bd%isd:bd%ied ,bd%jsd:bd%jed ,npz) !< pressure thickness (pascal) + real, intent(inout) :: q( bd%isd:bd%ied ,bd%jsd:bd%jed ,npz, ncnst) !< specific humidity and constituents + real, intent(inout) :: delz(bd%isd:,bd%jsd:,1:) !< delta-height (m); non-hydrostatic only + real, intent(inout) :: ze0(bd%is:, bd%js: ,1:) !< height at edges (m); non-hydrostatic + real, intent(inout), dimension(bd%isd:bd%ied ,bd%jsd:bd%jed, npz) :: diss_est! diffusion estimate for SKEB +! ze0 no longer used + real, intent(inout), dimension(bd%is :bd%ie ,bd%js :bd%je ,npz) :: dudt_rf ! U-wind tendency from Rayleigh friction + real, intent(inout), dimension(bd%is :bd%ie ,bd%js :bd%je ,npz) :: dvdt_rf ! V-wind tendency from Rayleigh friction + real, intent(inout), dimension(bd%is :bd%ie ,bd%js :bd%je ,npz) :: dwdt_rf ! W tendency from Rayleigh friction + real, intent(inout), dimension(bd%is :bd%ie ,bd%js :bd%je ,npz) :: dtdt_rf ! Temp tendency from Rayleigh friction + +!----------------------------------------------------------------------- +! Auxilliary pressure arrays: +! The 5 vars below can be re-computed from delp and ptop. +!----------------------------------------------------------------------- +! dyn_aux: + real, intent(inout) :: ps (bd%isd:bd%ied ,bd%jsd:bd%jed) !< Surface pressure (pascal) + real, intent(inout) :: pe (bd%is-1:bd%ie+1, npz+1,bd%js-1:bd%je+1) !< edge pressure (pascal) + real, intent(inout) :: pk (bd%is:bd%ie,bd%js:bd%je, npz+1) !< pe**kappa + real, intent(inout) :: peln(bd%is:bd%ie,npz+1,bd%js:bd%je) !< ln(pe) + real, intent(inout) :: pkz (bd%is:bd%ie,bd%js:bd%je,npz) !< finite-volume mean pk + real, intent(inout):: q_con(bd%isd:, bd%jsd:, 1:) + +!----------------------------------------------------------------------- +! Others: +!----------------------------------------------------------------------- + real, intent(inout) :: phis(bd%isd:bd%ied,bd%jsd:bd%jed) !< Surface geopotential (g*Z_surf) + real, intent(inout) :: varflt(bd%is:bd%ie,bd%js:bd%je) !< Variance of the filtered topography (m^2) + real, intent(inout) :: omga(bd%isd:bd%ied,bd%jsd:bd%jed,npz) !< Vertical pressure velocity (pa/s) + real, intent(inout) :: uc(bd%isd:bd%ied+1,bd%jsd:bd%jed ,npz) !< (uc,vc) mostly used as the C grid winds + real, intent(inout) :: vc(bd%isd:bd%ied ,bd%jsd:bd%jed+1,npz) + + real, intent(inout), dimension(bd%isd:bd%ied ,bd%jsd:bd%jed ,npz):: ua, va + real, intent(in), dimension(npz+1):: ak, bk + +! Accumulated Mass flux arrays: the "Flux Capacitor" + real, intent(inout) :: mfx(bd%is:bd%ie+1, bd%js:bd%je, npz) + real, intent(inout) :: mfy(bd%is:bd%ie , bd%js:bd%je+1, npz) +! Accumulated Courant number arrays + real, intent(inout) :: cx(bd%is:bd%ie+1, bd%jsd:bd%jed, npz) + real, intent(inout) :: cy(bd%isd:bd%ied ,bd%js:bd%je+1, npz) + + type(fv_grid_type), intent(inout), target :: gridstruct + type(fv_flags_type), intent(INOUT) :: flagstruct + type(fv_nest_type), intent(INOUT) :: neststruct + type(domain2d), intent(INOUT) :: domain + type(fv_atmos_type), pointer, intent(INOUT) :: parent_grid + type(fv_diag_type), intent(IN) :: idiag + +! Local Arrays +#ifdef SINGLE_FV +! R8 Mass flux arrays: the "Flux Capacitor" + real(kind=8) :: mfxR8(bd%is:bd%ie+1, bd%js:bd%je, npz) + real(kind=8) :: mfyR8(bd%is:bd%ie , bd%js:bd%je+1, npz) +! R8 Courant number arrays + real(kind=8) :: cxR8(bd%is:bd%ie+1, bd%jsd:bd%jed, npz) + real(kind=8) :: cyR8(bd%isd:bd%ied ,bd%js:bd%je+1, npz) + !$ser verbatim real :: mfxR8_ser(bd%is:bd%ie+1, bd%js:bd%je, npz) + !$ser verbatim real :: mfyR8_ser(bd%is:bd%ie , bd%js:bd%je+1, npz) + !$ser verbatim real :: cxR8_ser(bd%is:bd%ie+1, bd%jsd:bd%jed, npz) + !$ser verbatim real :: cyR8_ser(bd%isd:bd%ied ,bd%js:bd%je+1, npz) +#endif +! Local Mass flux arrays: the "Flux Capacitor" + real :: mfxL(bd%is:bd%ie+1, bd%js:bd%je, npz) + real :: mfyL(bd%is:bd%ie , bd%js:bd%je+1, npz) +! Local Courant number arrays + real :: cxL(bd%is:bd%ie+1, bd%jsd:bd%jed, npz) + real :: cyL(bd%isd:bd%ied ,bd%js:bd%je+1, npz) +! More Local arrays + real:: ws(bd%is:bd%ie,bd%js:bd%je) + real(kind=8):: te_2d(bd%is:bd%ie,bd%js:bd%je) + real:: aam(bd%is:bd%ie,bd%js:bd%je) + real:: teq(bd%is:bd%ie,bd%js:bd%je) + real:: ps2(bd%isd:bd%ied,bd%jsd:bd%jed) + real:: m_fac(bd%is:bd%ie,bd%js:bd%je) + real:: pfull(npz) + real, dimension(bd%is:bd%ie):: cvm + real, allocatable :: dp1(:,:,:), dtdt_m(:,:,:), cappa(:,:,:) + real(kind=8), allocatable :: psx(:,:) + !real(kind=8), allocatable :: dpx(:,:) + real(kind=8), allocatable :: dpx(:,:,:) !needed for OpenMP + real:: akap, rdg, ph1, ph2, mdt, gam, amdt, u0 + integer:: kord_tracer(ncnst) + integer :: i,j,k, n, iq, n_map, nq, nwat + integer :: sphum, liq_wat = -999, ice_wat = -999 ! GFDL physics + integer :: rainwat = -999, snowwat = -999, graupel = -999, cld_amt = -999 + integer :: theta_d = -999 + logical used, last_step, do_omega + type(group_halo_update_type), save :: i_pack(12) + integer :: is, ie, js, je + integer :: isd, ied, jsd, jed + real :: dt2 + + !$ser verbatim integer:: mode, n_map_step + !$ser verbatim real :: ph1v(npz), ph2v(npz) + + is = bd%is + ie = bd%ie + js = bd%js + je = bd%je + isd = bd%isd + ied = bd%ied + jsd = bd%jsd + jed = bd%jed + +! Empty the accumulated mass flux and courant numbers + mfx = 0.0 + mfy = 0.0 + cx = 0.0 + cy = 0.0 + +! cv_air = cp_air - rdgas + agrav = 1. / grav + dt2 = 0.5*bdt + nwat = flagstruct%nwat + nq = nq_tot - flagstruct%dnats + rdg = -rdgas * agrav + allocate ( dp1(isd:ied, jsd:jed, 1:npz) ) + +#ifdef MOIST_CAPPA + allocate ( cappa(isd:ied,jsd:jed,npz) ) + call init_ijk_mem(isd,ied, jsd,jed, npz, cappa, 0.) +#else + allocate ( cappa(isd:isd,jsd:jsd,1) ) + cappa = 0. +#endif + !We call this BEFORE converting pt to virtual potential temperature, + !since we interpolate on (regular) temperature rather than theta. + if (gridstruct%nested .or. ANY(neststruct%child_grids)) then + call timing_on('NEST_BCs') + call setup_nested_grid_BCs(npx, npy, npz, zvir, ncnst, & + u, v, w, pt, delp, delz, q, uc, vc, pkz, & + neststruct%nested, flagstruct%inline_q, flagstruct%make_nh, ng, & + gridstruct, flagstruct, neststruct, & + neststruct%nest_timestep, neststruct%tracer_nest_timestep, & + domain, bd, nwat) + +#ifndef SW_DYNAMICS + if (gridstruct%nested) then + !Correct halo values have now been set up for BCs; we can go ahead and apply them too... + call nested_grid_BC_apply_intT(pt, & + 0, 0, npx, npy, npz, bd, 1., 1., & + neststruct%pt_BC, bctype=neststruct%nestbctype ) +#ifdef USE_COND + call nested_grid_BC_apply_intT(q_con, & + 0, 0, npx, npy, npz, bd, 1., 1., & + neststruct%q_con_BC, bctype=neststruct%nestbctype ) +#ifdef MOIST_CAPPA + call nested_grid_BC_apply_intT(cappa, & + 0, 0, npx, npy, npz, bd, 1., 1., & + neststruct%cappa_BC, bctype=neststruct%nestbctype ) +#endif +#endif + endif +#endif + call timing_off('NEST_BCs') + endif + + + if ( flagstruct%no_dycore ) then + if ( nwat.eq.2 .and. (.not.hydrostatic) ) then + sphum = get_tracer_index (MODEL_ATMOS, 'sphum') + endif + goto 911 + endif + +#ifdef MAPL_MODE + select case(nwat) + case(0) + sphum = 1 + liq_wat = -1 + ice_wat = -1 + rainwat = -1 + snowwat = -1 + graupel = -1 + cld_amt = -1 + case(1) + sphum = 1 + liq_wat = -1 + ice_wat = -1 + rainwat = -1 + snowwat = -1 + graupel = -1 + cld_amt = -1 + case(3) + sphum = 1 + liq_wat = 2 + ice_wat = 3 + rainwat = -1 + snowwat = -1 + graupel = -1 + cld_amt = -1 + case(6:7) + sphum = 1 + liq_wat = 2 + ice_wat = 3 + rainwat = 4 + snowwat = 5 + graupel = 6 + cld_amt = 7 + end select +#else + if ( nwat==0 ) then + sphum = 1 + cld_amt = -1 ! to cause trouble if (mis)used + else + sphum = get_tracer_index (MODEL_ATMOS, 'sphum') + liq_wat = get_tracer_index (MODEL_ATMOS, 'liq_wat') + ice_wat = get_tracer_index (MODEL_ATMOS, 'ice_wat') + rainwat = get_tracer_index (MODEL_ATMOS, 'rainwat') + snowwat = get_tracer_index (MODEL_ATMOS, 'snowwat') + graupel = get_tracer_index (MODEL_ATMOS, 'graupel') + cld_amt = get_tracer_index (MODEL_ATMOS, 'cld_amt') + endif + + theta_d = get_tracer_index (MODEL_ATMOS, 'theta_d') +#endif + +#ifdef SW_DYNAMICS + akap = 1. + pfull(1) = 0.5*flagstruct%p_ref +#else + akap = kappa + +!$OMP parallel do default(none) shared(npz,ak,bk,flagstruct,pfull) & +!$OMP private(ph1, ph2) + do k=1,npz + ph1 = ak(k ) + bk(k )*flagstruct%p_ref + ph2 = ak(k+1) + bk(k+1)*flagstruct%p_ref + pfull(k) = (ph2 - ph1) / log(ph2/ph1) + enddo + + if ( hydrostatic ) then +!$OMP parallel do default(none) shared(is,ie,js,je,isd,ied,jsd,jed,npz,dp1,zvir,nwat,q,q_con,sphum,liq_wat, & +!$OMP rainwat,ice_wat,snowwat,graupel) private(cvm) + do k=1,npz + do j=js,je + do i=is,ie + dp1(i,j,k) = zvir*q(i,j,k,sphum) + enddo + enddo + enddo + else +!$OMP parallel do default(none) shared(is,ie,js,je,isd,ied,jsd,jed,npz,dp1,zvir,q,q_con,sphum,liq_wat, & +!$OMP rainwat,ice_wat,snowwat,graupel,pkz,flagstruct, & +!$OMP cappa,kappa,rdg,delp,pt,delz,nwat) & +!$OMP private(cvm) + do k=1,npz + if ( flagstruct%moist_phys ) then + do j=js,je +#ifdef MOIST_CAPPA + call moist_cv(is,ie,isd,ied,jsd,jed, npz, j, k, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, q, q_con(is:ie,j,k), cvm) +#endif + do i=is,ie + dp1(i,j,k) = zvir*q(i,j,k,sphum) +#ifdef MOIST_CAPPA + cappa(i,j,k) = rdgas/(rdgas + cvm(i)/(1.+dp1(i,j,k))) + pkz(i,j,k) = exp(cappa(i,j,k)*log(rdg*delp(i,j,k)*pt(i,j,k)* & + (1.+dp1(i,j,k))*(1.-q_con(i,j,k))/delz(i,j,k)) ) +#else + pkz(i,j,k) = exp( kappa*log(rdg*delp(i,j,k)*pt(i,j,k)* & + (1.+dp1(i,j,k))/delz(i,j,k)) ) +! Using dry pressure for the definition of the virtual potential temperature +! pkz(i,j,k) = exp( kappa*log(rdg*delp(i,j,k)*pt(i,j,k)* & +! (1.-q(i,j,k,sphum))/delz(i,j,k)) ) +#endif + enddo + enddo + else + do j=js,je + do i=is,ie + dp1(i,j,k) = 0. + pkz(i,j,k) = exp(kappa*log(rdg*delp(i,j,k)*pt(i,j,k)/delz(i,j,k))) + enddo + enddo + endif + enddo + endif + + if ( flagstruct%fv_debug ) then + call prt_mxm('PS_b', ps, is, ie, js, je, ng, 1, 0.01, gridstruct%area_64, domain) + call prt_mxm('T_dyn_b', pt, is, ie, js, je, ng, npz, 1., gridstruct%area_64, domain) + if ( .not. hydrostatic) call prt_mxm('delz_b', delz, is, ie, js, je, ng, npz, 1., gridstruct%area_64, domain) + if ( .not. hydrostatic) call prt_mxm('W_b', w, is, ie, js, je, ng, npz, 1., gridstruct%area_64, domain) + call prt_mxm('delp_b ', delp, is, ie, js, je, ng, npz, 0.01, gridstruct%area_64, domain) + call prt_mxm('pk_b', pk, is, ie, js, je, 0, npz+1, 1.,gridstruct%area_64, domain) + call prt_mxm('pkz_b', pkz,is, ie, js, je, 0, npz, 1.,gridstruct%area_64, domain) + endif + +!--------------------- +! Compute Total Energy +!--------------------- + if ( consv_te > 0. .and. (.not.do_adiabatic_init) ) then + call compute_total_energy(is, ie, js, je, isd, ied, jsd, jed, npz, & + u, v, w, delz, pt, delp, q, dp1, pe, peln, phis, & + gridstruct%rsin2, gridstruct%cosa_s, & + zvir, cp_air, rdgas, hlv, te_2d, ua, va, teq, & + flagstruct%moist_phys, nwat, sphum, liq_wat, rainwat, & + ice_wat, snowwat, graupel, hydrostatic, idiag%id_te) + if( idiag%id_te>0 ) then + used = send_data(idiag%id_te, teq, fv_time) +! te_den=1.E-9*g_sum(teq, is, ie, js, je, ng, area, 0)/(grav*4.*pi*radius**2) +! if(is_master()) write(*,*) 'Total Energy Density (Giga J/m**2)=',te_den + endif + endif + + if( (flagstruct%consv_am.or.idiag%id_amdt>0) .and. (.not.do_adiabatic_init) ) then + call compute_aam(npz, is, ie, js, je, isd, ied, jsd, jed, gridstruct, bd, & + ptop, ua, va, u, v, delp, teq, ps2, m_fac) + endif + + dudt_rf = u(is:ie,js:je,:) + dvdt_rf = v(is:ie,js:je,:) + dtdt_rf = pt(is:ie,js:je,:) ! at this time PT is dry T (K) + if (.not. hydrostatic) dwdt_rf = w(is:ie,js:je,:) + if (.not.flagstruct%RF_fast .and. flagstruct%tau > 0. ) then + if ( gridstruct%grid_type<4 ) then + call Rayleigh_Super(abs(bdt), npx, npy, npz, ks, pfull, phis, flagstruct%tau, u, v, w, pt, & + ua, va, delz, gridstruct%agrid, cp_air, rdgas, ptop, hydrostatic, & + (.not. neststruct%nested), flagstruct%rf_cutoff, gridstruct, domain, bd) + else + call Rayleigh_Friction(abs(bdt), npx, npy, npz, ks, pfull, flagstruct%tau, u, v, w, pt, & + ua, va, delz, cp_air, rdgas, ptop, hydrostatic, .true., flagstruct%rf_cutoff, gridstruct, domain, bd) + endif + endif + dudt_rf = ( u(is:ie,js:je,:) - dudt_rf)/bdt + dvdt_rf = ( v(is:ie,js:je,:) - dvdt_rf)/bdt + dtdt_rf = (pt(is:ie,js:je,:) - dtdt_rf)/bdt + if (.not. hydrostatic) dwdt_rf = ( w(is:ie,js:je,:) - dwdt_rf)/bdt + +#endif + +#ifndef SW_DYNAMICS +! Convert pt to virtual potential temperature on the first timestep + if ( flagstruct%adiabatic ) then +!$OMP parallel do default(none) shared(theta_d,is,ie,js,je,npz,pt,pkz,q) + do k=1,npz + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k)/pkz(i,j,k) + enddo + enddo + if ( theta_d>0 ) then + do j=js,je + do i=is,ie + q(i,j,k,theta_d) = pt(i,j,k) + enddo + enddo + endif + enddo + else + if (hydrostatic) then +!$OMP parallel do default(none) shared(is,ie,js,je,npz,pt,dp1,pkz) + do k=1,npz + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k)*(1.+dp1(i,j,k))/pkz(i,j,k) + enddo + enddo + enddo + else +!$OMP parallel do default(none) shared(is,ie,js,je,npz,pt,dp1,pkz,q_con) + do k=1,npz + do j=js,je + do i=is,ie +#ifdef USE_COND + pt(i,j,k) = pt(i,j,k)*(1.+dp1(i,j,k))*(1.-q_con(i,j,k))/pkz(i,j,k) +#else + pt(i,j,k) = pt(i,j,k)*(1.+dp1(i,j,k))/pkz(i,j,k) +#endif + enddo + enddo + enddo + endif + endif +#endif + + last_step = .false. + mdt = bdt / real(k_split) + flagstruct%m_split = CEILING(1. + abs(bdt)/real(k_split*n_split)) + + allocate ( dtdt_m(is:ie,js:je,npz) ) + if ( idiag%id_mdt > 0 .and. (.not. do_adiabatic_init) ) then +!$OMP parallel do default(none) shared(is,ie,js,je,npz,dtdt_m) + do k=1,npz + do j=js,je + do i=is,ie + dtdt_m(i,j,k) = 0. + enddo + enddo + enddo + endif + +!DryMassRoundoffControl + allocate(psx(isd:ied,jsd:jed),dpx(is:ie,js:je,npz)) + psx(:,:) = 0.0 + do j=js,je + do i=is,ie + psx(i,j) = pe(i,npz+1,j) + dpx(i,j,:) = 0.0 + enddo + enddo + call timing_on('FV_DYN_LOOP') + do n_map=1, k_split ! first level of time-split + call timing_on('COMM_TOTAL') + !$ser verbatim n_map_step=n_map + !$ser savepoint DynCore-In + !$ser verbatim ! No use before this point, we init to zero + !$ser verbatim mfxR8_ser = 0 + !$ser verbatim mfyR8_ser = 0 + !$ser verbatim cxR8_ser = 0 + !$ser verbatim cyR8_ser = 0 + !$ser data nq=nq mdt=mdt n_split=n_split zvir=zvir akap=akap cappa=cappa u=u v=v w=w delz=delz pt=pt delp=delp pe=pe pk=pk phis=phis wsd=ws omga=omga ptop=ptop pfull=pfull ua=ua va=va uc=uc vc=vc + !$ser data mfxd_R8=mfxR8 mfyd_R8=mfyR8 cxd_R8=cxR8 cyd_R8=cyR8 + !$ser data pkz=pkz peln=peln q_con=q_con ak=ak bk=bk ks=ks diss_estd=diss_est n_map=n_map_step + if (.not. hydrostatic) then +#ifdef USE_COND + call start_group_halo_update(i_pack(11), q_con, domain) +#ifdef MOIST_CAPPA + call start_group_halo_update(i_pack(12), cappa, domain) +#endif +#endif + endif + + call start_group_halo_update(i_pack(1), delp, domain, complete=.false.) + call start_group_halo_update(i_pack(1), pt, domain, complete=.true.) +#ifndef ROT3 + call start_group_halo_update(i_pack(8), u, v, domain, gridtype=DGRID_NE) +#endif + call timing_off('COMM_TOTAL') +!$OMP parallel do default(none) shared(isd,ied,jsd,jed,npz,dp1,delp) + do k=1,npz + do j=jsd,jed + do i=isd,ied + dp1(i,j,k) = delp(i,j,k) + enddo + enddo + enddo + + if ( n_map==k_split ) last_step = .true. + + if (.not. hydrostatic) then +#ifdef USE_COND + call timing_on('COMM_TOTAL') + call complete_group_halo_update(i_pack(11), domain) +#ifdef MOIST_CAPPA + call complete_group_halo_update(i_pack(12), domain) +#endif + call timing_off('COMM_TOTAL') +#endif + endif + + call timing_on('DYN_CORE') + call dyn_core(npx, npy, npz, ng, sphum, nq, mdt, k_split, n_split, zvir, cp_air, akap, cappa, grav, hydrostatic, & + u, v, w, delz, pt, q, delp, pe, pk, phis, varflt, ws, omga, ptop, pfull, ua, va, & + uc, vc, & +#ifdef SINGLE_FV + mfxR8, mfyR8, cxR8, cyR8, & +#else + mfxL, mfyL, cxL, cyL, & +#endif + pkz, peln, q_con, ak, bk, dpx, ks, & + gridstruct, flagstruct, neststruct, idiag, bd, & + domain, n_map==1, i_pack, last_step, diss_est,time_total) + call timing_off('DYN_CORE') + !$ser savepoint DynCore-Out + !$ser verbatim ! Copy back to ensure 32-bit save for Serialbox +#ifdef SINGLE_FV + !$ser verbatim mfxR8_ser = mfxR8 + !$ser verbatim mfyR8_ser = mfyR8 + !$ser verbatim cxR8_ser = cxR8 + !$ser verbatim cyR8_ser = cyR8 +#else + !$ser verbatim mfxR8_ser = mfxL + !$ser verbatim mfyR8_ser = mfyL + !$ser verbatim cxR8_ser = cxL + !$ser verbatim cyR8_ser = cyL +#endif + +!MassFluxRoundoffControl +#ifdef SINGLE_FV + mfxL=mfxR8 + mfyL=mfyR8 + cxL= cxR8 + cyL= cyR8 +#endif + !$ser data cappa=cappa u=u v=v w=w delz=delz pt=pt delp=delp pe=pe pk=pk phis=phis wsd=ws omga=omga ptop=ptop pfull=pfull ua=ua va=va uc=uc vc=vc + !$ser data mfxd_R8=mfxR8 mfyd_R8=mfyR8 cxd_R8=cxR8 cyd_R8=cyR8 + !$ser data mfxL=mfxL mfyL=mfyL cxL=cxL cyL=cyL + !$ser data pkz=pkz peln=peln q_con=q_con diss_estd=diss_est + +!DryMassRoundoffControl + if(last_step) then + if (hydrostatic) then + do k = 2, npz + dpx(:,:,1) = dpx(:,:,1) + dpx(:,:,k) + enddo + do j=js,je + do i=is,ie + psx(i,j) = psx(i,j) + dpx(i,j,1) + enddo + enddo + call timing_on('COMM_TOTAL') + call mpp_update_domains(psx, domain) + call timing_off('COMM_TOTAL') + do j=js-1,je+1 + do i=is-1,ie+1 + pe(i,npz+1,j) = psx(i,j) + enddo + enddo + endif + deallocate(psx,dpx) + end if + +#ifdef SW_DYNAMICS +!!$OMP parallel do default(none) shared(is,ie,js,je,ps,delp,agrav) + do j=js,je + do i=is,ie + ps(i,j) = delp(i,j,1) * agrav + enddo + enddo +#else + if( .not. flagstruct%inline_q .and. nq /= 0 ) then +!-------------------------------------------------------- +! Perform large-time-step scalar transport using the accumulated CFL and +! mass fluxes + call timing_on('tracer_2d') + !!! CLEANUP: merge these two calls? + if (gridstruct%nested) then + call tracer_2d_nested(q, dp1, mfxL, mfyL, cxL, cyL, gridstruct, bd, domain, npx, npy, npz, nq, & + flagstruct%hord_tr, q_split, mdt, idiag%id_divg, i_pack(10), & + flagstruct%nord_tr, flagstruct%trdm2, & + k_split, neststruct, parent_grid, flagstruct%lim_fac) + else + if ( flagstruct%z_tracer ) then + !$ser savepoint Tracer2D1L-In + !$ser data tracers=q dp1=dp1 + !$ser data mfxd_R4=mfxL mfyd_R4=mfyL cxd_R4=cxL cyd_R4=cyL + !$ser data nq=nq q_split=q_split mdt=mdt + call tracer_2d_1L(q, dp1, mfxL, mfyL, cxL, cyL, gridstruct, bd, domain, npx, npy, npz, nq, & + flagstruct%hord_tr, q_split, mdt, idiag%id_divg, i_pack(10), & + flagstruct%nord_tr, flagstruct%trdm2, flagstruct%lim_fac) + !$ser savepoint Tracer2D1L-Out + !$ser data tracers=q dp1=dp1 + !$ser data mfxd_R4=mfxL mfyd_R4=mfyL cxd_R4=cxL cyd_R4=cyL + else + call tracer_2d(q, dp1, mfxL, mfyL, cxL, cyL, gridstruct, bd, domain, npx, npy, npz, nq, & + flagstruct%hord_tr, q_split, mdt, idiag%id_divg, i_pack(10), & + flagstruct%nord_tr, flagstruct%trdm2, flagstruct%lim_fac) + endif + endif + call timing_off('tracer_2d') + +#ifdef FILL2D + if ( flagstruct%hord_tr<8 .and. flagstruct%moist_phys ) then + call timing_on('Fill2D') + if ( liq_wat > 0 ) & + call fill2D(is, ie, js, je, ng, npz, q(isd,jsd,1,liq_wat), delp, gridstruct%area, domain, neststruct%nested, npx, npy) + if ( rainwat > 0 ) & + call fill2D(is, ie, js, je, ng, npz, q(isd,jsd,1,rainwat), delp, gridstruct%area, domain, neststruct%nested, npx, npy) + if ( ice_wat > 0 ) & + call fill2D(is, ie, js, je, ng, npz, q(isd,jsd,1,ice_wat), delp, gridstruct%area, domain, neststruct%nested, npx, npy) + if ( snowwat > 0 ) & + call fill2D(is, ie, js, je, ng, npz, q(isd,jsd,1,snowwat), delp, gridstruct%area, domain, neststruct%nested, npx, npy) + if ( graupel > 0 ) & + call fill2D(is, ie, js, je, ng, npz, q(isd,jsd,1,graupel), delp, gridstruct%area, domain, neststruct%nested, npx, npy) + call timing_off('Fill2D') + endif +#endif + + if( last_step .and. idiag%id_divg>0 ) then + used = send_data(idiag%id_divg, dp1, fv_time) + if(flagstruct%fv_debug) call prt_mxm('divg', dp1, is, ie, js, je, 0, npz, 1.,gridstruct%area_64, domain) + endif + endif + + if ( npz > 4 ) then +!------------------------------------------------------------------------ +! Perform vertical remapping from Lagrangian control-volume to +! the Eulerian coordinate as specified by the routine set_eta. +! Note that this finite-volume dycore is otherwise independent of the vertical +! Eulerian coordinate. +!------------------------------------------------------------------------ + + do iq=1,nq + kord_tracer(iq) = flagstruct%kord_tr + if ( iq==cld_amt ) kord_tracer(iq) = 9 ! monotonic + enddo + + do_omega = hydrostatic .and. last_step + call timing_on('Remapping') +#ifdef AVEC_TIMERS + call avec_timer_start(6) +#endif + + !$ser savepoint Remapping-In + !$ser data iq=iq last_step=last_step consv_te=consv_te ps=ps pe=pe delp=delp pkz=pkz pk=pk mdt=mdt bdt=bdt sphum=sphum q_con=q_con u=u v=v w=w delz=delz pt=pt tracers=q phis=phis zvir=zvir akap=akap cappa=cappa kord_tracer=kord_tracer peln=peln te_2d=te_2d ua=ua va=va omga=omga dp1=dp1 wsd=ws reproduce_sum=reproduce_sum ptop=ptop ak=ak bk=bk pfull=pfull hybrid_z=hybrid_z do_adiabatic_init=do_adiabatic_init nq=nq + !$ser savepoint Lagrangian_to_Eulerian-In + !$ser data last_step=last_step consv=consv_te ps=ps pe=pe delp=delp pkz=pkz mdt=mdt pdt=bdt + !$ser data km=npz nq=nq nwat=nwat sphum=sphum q_con=q_con u=u v=v w=w delz=delz pt=pt q_=q + !$ser data hs=phis r_vir=zvir cp=cp_air akap=akap cappa=cappa kord_tr=kord_tracer + !$ser data peln=peln te0_2d=te_2d ng=ng ua=ua va=va omga=omga te=dp1 ws_=ws fill=fill + !$ser data reproduce_sum=reproduce_sum dtdt=dtdt_m ptop=ptop ak=ak bk=bk pfull=pfull + !$ser data hydrostatic=hydrostatic hybrid_z=hybrid_z do_omega=do_omega + !$ser data do_adiabatic_init=do_adiabatic_init + !$ser data mfx=mfxL + !$ser data mfy=mfyL + !$ser data cx=cxL + !$ser data cy=cyL + + call Lagrangian_to_Eulerian(last_step, consv_te, ps, pe, delp, & + pkz, pk, mdt, bdt, npz, is,ie,js,je, isd,ied,jsd,jed, & + nq, nwat, sphum, q_con, u, v, w, delz, pt, q, phis, & + zvir, cp_air, akap, cappa, flagstruct%kord_mt, flagstruct%kord_wz, & + kord_tracer, flagstruct%kord_tm, peln, te_2d, & + ng, ua, va, omga, dp1, ws, fill, reproduce_sum, & + idiag%id_mdt>0, dtdt_m, ptop, ak, bk, pfull, flagstruct, gridstruct, domain, & + flagstruct%do_sat_adj, hydrostatic, hybrid_z, do_omega, & + flagstruct%adiabatic, do_adiabatic_init, & + mfxL, mfyL, cxL, cyL, flagstruct%remap_option, flagstruct%gmao_remap) + !$ser savepoint Remapping-Out + !$ser data te_2d=te_2d pk=pk tracers=q delp=delp pe=pe ps=ps u=u v=v w=w pt=pt delz=delz q_con=q_con cappa=cappa ua=ua va=va omga=omga peln=peln pkz=pkz dp1=dp1 + !$ser savepoint Lagrangian_to_Eulerian-Out + !$ser data ps=ps pe=pe pkz=pkz q_con=q_con u=u v=v w=w delz=delz pt=pt q_=q cappa=cappa + !$ser data peln=peln te0_2d=te_2d ua=ua va=va omga=omga te=dp1 dtdt=dtdt_m + !$ser data mfx=mfxL + !$ser data mfy=mfyL + !$ser data cx=cxL + !$ser data cy=cyL + +#ifdef AVEC_TIMERS + call avec_timer_stop(6) +#endif + call timing_off('Remapping') +#ifdef MOIST_CAPPA + if ( neststruct%nested .and. .not. last_step) then + call nested_grid_BC_apply_intT(cappa, & + 0, 0, npx, npy, npz, bd, real(n_map+1), real(k_split), & + neststruct%cappa_BC, bctype=neststruct%nestbctype ) + endif +#endif + +! Accumulate the total Mass flux and Courant numbers for export + mfx = mfx + mfxL + mfy = mfy + mfyL + cx = cx + cxL + cy = cy + cyL + + if( last_step ) then + if( .not. hydrostatic ) then +!$OMP parallel do default(none) shared(is,ie,js,je,npz,omga,delp,delz,w) + do k=1,npz + do j=js,je + do i=is,ie + omga(i,j,k) = delp(i,j,k)/delz(i,j,k)*w(i,j,k) + enddo + enddo + enddo + endif +!-------------------------- +! Filter omega for physics: +!-------------------------- + if(flagstruct%nf_omega>0) & + call del2_cubed(omga, 0.18*gridstruct%da_min, gridstruct, domain, npx, npy, npz, flagstruct%nf_omega, bd) + endif + end if +#endif + enddo ! n_map loop + call timing_off('FV_DYN_LOOP') + if ( idiag%id_mdt > 0 .and. (.not.do_adiabatic_init) ) then +! Output temperature tendency due to inline moist physics: +!$OMP parallel do default(none) shared(is,ie,js,je,npz,dtdt_m,bdt) + do k=1,npz + do j=js,je + do i=is,ie + dtdt_m(i,j,k) = dtdt_m(i,j,k) / bdt * 86400. + enddo + enddo + enddo +! call prt_mxm('Fast DTDT (deg/Day)', dtdt_m, is, ie, js, je, 0, npz, 1., gridstruct%area_64, domain) + used = send_data(idiag%id_mdt, dtdt_m, fv_time) + endif + deallocate ( dtdt_m ) + + if( nwat>=6 ) then + !$ser savepoint Neg_Adj3-In + !$ser data peln=peln delz=delz pt=pt delp=delp qvapor=q(:,:,:,sphum) qliquid=q(:,:,:,liq_wat) qice=q(:,:,:,ice_wat) qrain=q(:,:,:,rainwat) qsnow=q(:,:,:,snowwat) qgraupel=q(:,:,:,graupel) qcld=q(:,:,:,cld_amt) + if (cld_amt > 0) then + call neg_adj3(is, ie, js, je, ng, npz, & + flagstruct%hydrostatic, & + peln, delz, & + pt, delp, q(isd,jsd,1,sphum), & + q(isd,jsd,1,liq_wat), & + q(isd,jsd,1,rainwat), & + q(isd,jsd,1,ice_wat), & + q(isd,jsd,1,snowwat), & + q(isd,jsd,1,graupel), & + q(isd,jsd,1,cld_amt), flagstruct%check_negative) + else + call neg_adj3(is, ie, js, je, ng, npz, & + flagstruct%hydrostatic, & + peln, delz, & + pt, delp, q(isd,jsd,1,sphum), & + q(isd,jsd,1,liq_wat), & + q(isd,jsd,1,rainwat), & + q(isd,jsd,1,ice_wat), & + q(isd,jsd,1,snowwat), & + q(isd,jsd,1,graupel), check_negative=flagstruct%check_negative) + endif + !$ser savepoint Neg_Adj3-Out + !$ser data qvapor=q(:,:,:,sphum) qliquid=q(:,:,:,liq_wat) qice=q(:,:,:,ice_wat) qrain=q(:,:,:,rainwat) qsnow=q(:,:,:,snowwat) qgraupel=q(:,:,:,graupel) qcld=q(:,:,:,cld_amt) pt=pt + if ( flagstruct%fv_debug ) then + call prt_mxm('SPHUM_dyn', q(isd,jsd,1,sphum ), is, ie, js, je, ng, npz, 1.,gridstruct%area_64, domain) + call prt_mxm('liq_wat_dyn', q(isd,jsd,1,liq_wat), is, ie, js, je, ng, npz, 1.,gridstruct%area_64, domain) + call prt_mxm('rainwat_dyn', q(isd,jsd,1,rainwat), is, ie, js, je, ng, npz, 1.,gridstruct%area_64, domain) + call prt_mxm('ice_wat_dyn', q(isd,jsd,1,ice_wat), is, ie, js, je, ng, npz, 1.,gridstruct%area_64, domain) + call prt_mxm('snowwat_dyn', q(isd,jsd,1,snowwat), is, ie, js, je, ng, npz, 1.,gridstruct%area_64, domain) + call prt_mxm('graupel_dyn', q(isd,jsd,1,graupel), is, ie, js, je, ng, npz, 1.,gridstruct%area_64, domain) + endif + endif + + if( (flagstruct%consv_am.or.idiag%id_amdt>0.or.idiag%id_aam>0) .and. (.not.do_adiabatic_init) ) then + call compute_aam(npz, is, ie, js, je, isd, ied, jsd, jed, gridstruct, bd, & + ptop, ua, va, u, v, delp, aam, ps, m_fac) + if( idiag%id_aam>0 ) then + used = send_data(idiag%id_aam, aam, fv_time) + if ( prt_minmax ) then + gam = g_sum( domain, aam, is, ie, js, je, ng, gridstruct%area_64, 0) + if( is_master() ) write(6,*) 'Total AAM =', gam + endif + endif + endif + + if( (flagstruct%consv_am.or.idiag%id_amdt>0) .and. (.not.do_adiabatic_init) ) then +!$OMP parallel do default(none) shared(is,ie,js,je,aam,teq,dt2,ps2,ps,idiag) + do j=js,je + do i=is,ie +! Note: the mountain torque computation contains also numerical error +! The numerical error is mostly from the zonal gradient of the terrain (zxg) + aam(i,j) = aam(i,j)-teq(i,j) + dt2*(ps2(i,j)+ps(i,j))*idiag%zxg(i,j) + enddo + enddo + if( idiag%id_amdt>0 ) used = send_data(idiag%id_amdt, aam/bdt, fv_time) + + if ( flagstruct%consv_am .or. prt_minmax ) then + amdt = g_sum( domain, aam, is, ie, js, je, ng, gridstruct%area_64, 0, reproduce=.true.) + u0 = -radius*amdt/g_sum( domain, m_fac, is, ie, js, je, ng, gridstruct%area_64, 0,reproduce=.true.) + if(is_master() .and. prt_minmax) & + write(6,*) 'Dynamic AM tendency (Hadleys)=', amdt/(bdt*1.e18), 'del-u (per day)=', u0*86400./bdt + endif + + if( flagstruct%consv_am ) then +!$OMP parallel do default(none) shared(is,ie,js,je,m_fac,u0,gridstruct) + do j=js,je + do i=is,ie + m_fac(i,j) = u0*cos(gridstruct%agrid(i,j,2)) + enddo + enddo +!$OMP parallel do default(none) shared(is,ie,js,je,npz,hydrostatic,pt,m_fac,ua,cp_air, & +!$OMP u,u0,gridstruct,v ) + do k=1,npz + do j=js,je+1 + do i=is,ie + u(i,j,k) = u(i,j,k) + u0*gridstruct%l2c_u(i,j) + enddo + enddo + do j=js,je + do i=is,ie+1 + v(i,j,k) = v(i,j,k) + u0*gridstruct%l2c_v(i,j) + enddo + enddo + enddo + endif ! consv_am + endif + +!$ser savepoint CubedToLatLon-In +!$ser verbatim mode=1 +!$ser data u=u v=v ua=ua va=va mode=mode +911 call cubed_to_latlon(u, v, ua, va, gridstruct, & + npx, npy, npz, 1, gridstruct%grid_type, domain, gridstruct%nested, flagstruct%c2l_ord, bd) +!$ser savepoint CubedToLatLon-Out +!$ser data u=u v=v ua=ua va=va + + deallocate(dp1) + deallocate(cappa) + + if ( flagstruct%fv_debug ) then + call prt_mxm('UA', ua, is, ie, js, je, ng, npz, 1., gridstruct%area_64, domain) + call prt_mxm('VA', va, is, ie, js, je, ng, npz, 1., gridstruct%area_64, domain) + call prt_mxm('TA', pt, is, ie, js, je, ng, npz, 1., gridstruct%area_64, domain) + if (.not. hydrostatic) call prt_mxm('W ', w, is, ie, js, je, ng, npz, 1., gridstruct%area_64, domain) + endif + + if ( flagstruct%range_warn ) then + call range_check('UA_dyn', ua, is, ie, js, je, ng, npz, gridstruct%agrid, & + -280., 280., bad_range) + call range_check('VA_dyn', ua, is, ie, js, je, ng, npz, gridstruct%agrid, & + -280., 280., bad_range) + call range_check('TA_dyn', pt, is, ie, js, je, ng, npz, gridstruct%agrid, & + 100., 335., bad_range) + !if ( .not. hydrostatic ) then + ! call range_check('W_dyn', w, is, ie, js, je, ng, npz, gridstruct%agrid, & + ! -100., 100., bad_range) + !endif + endif + + end subroutine fv_dynamics + +#ifdef USE_RF_FAST + subroutine Rayleigh_fast(dt, npx, npy, npz, pfull, tau, u, v, w, & + ks, dp, ptop, hydrostatic, rf_cutoff, bd) +! Simple "inline" version of the Rayleigh friction + real, intent(in):: dt + real, intent(in):: tau !< time scale (days) + real, intent(in):: ptop, rf_cutoff + real, intent(in), dimension(npz):: pfull + integer, intent(in):: npx, npy, npz, ks + logical, intent(in):: hydrostatic + type(fv_grid_bounds_type), intent(IN) :: bd + real, intent(inout):: u(bd%isd:bd%ied ,bd%jsd:bd%jed+1,npz) !< D grid zonal wind (m/s) + real, intent(inout):: v(bd%isd:bd%ied+1,bd%jsd:bd%jed,npz) !< D grid meridional wind (m/s) + real, intent(inout):: w(bd%isd: ,bd%jsd: ,1: ) !< cell center vertical wind (m/s) + real, intent(in):: dp(npz) +! + real(kind=R_GRID):: rff(npz) + real, parameter:: sday = 86400. !< seconds per day + real, dimension(bd%is:bd%ie+1):: dmv + real, dimension(bd%is:bd%ie):: dmu + real:: tau0, dm + integer i, j, k + + integer :: is, ie, js, je + integer :: isd, ied, jsd, jed + + is = bd%is + ie = bd%ie + js = bd%js + je = bd%je + isd = bd%isd + ied = bd%ied + jsd = bd%jsd + jed = bd%jed + + if ( .not. RF_initialized ) then + tau0 = tau * sday + allocate( rf(npz) ) + rf(:) = 1. + + if( is_master() ) write(6,*) 'Fast Rayleigh friction E-folding time (days):' + do k=1, npz + if ( pfull(k) < rf_cutoff ) then + rff(k) = dt/tau0*sin(0.5*pi*log(rf_cutoff/pfull(k))/log(rf_cutoff/ptop))**2 +! Re-FACTOR rf + if( is_master() ) write(6,*) k, 0.01*pfull(k), dt/(rff(k)*sday) + kmax = k + rff(k) = 1.d0 / (1.0d0+rff(k)) + rf(k) = rff(k) + else + exit + endif + enddo + dm = 0. + do k=1,ks + if ( pfull(k) < 100.E2 ) then + dm = dm + dp(k) + k_rf = k + else + exit + endif + enddo + if( is_master() ) write(6,*) 'k_rf=', k_rf, 0.01*pfull(k_rf), 'dm=', dm + RF_initialized = .true. + endif + +!$OMP parallel do default(none) shared(k_rf,is,ie,js,je,kmax,pfull,rf_cutoff,w,rf,dp,u,v,hydrostatic) & +!$OMP private(dm, dmu, dmv) + do j=js,je+1 + + dm = 0. + do k=1, k_rf + dm = dm + dp(k) + enddo + + dmu(:) = 0. + dmv(:) = 0. + do k=1,kmax + do i=is,ie + dmu(i) = dmu(i) + (1.-rf(k))*dp(k)*u(i,j,k) + u(i,j,k) = rf(k)*u(i,j,k) + enddo + if ( j/=je+1 ) then + do i=is,ie+1 + dmv(i) = dmv(i) + (1.-rf(k))*dp(k)*v(i,j,k) + v(i,j,k) = rf(k)*v(i,j,k) + enddo + if ( .not. hydrostatic ) then + do i=is,ie + w(i,j,k) = rf(k)*w(i,j,k) + enddo + endif + endif + enddo + + do i=is,ie + dmu(i) = dmu(i) / dm + enddo + if ( j/=je+1 ) then + do i=is,ie+1 + dmv(i) = dmv(i) / dm + enddo + endif + + do k=1, k_rf + do i=is,ie + u(i,j,k) = u(i,j,k) + dmu(i) + enddo + if ( j/=je+1 ) then + do i=is,ie+1 + v(i,j,k) = v(i,j,k) + dmv(i) + enddo + endif + enddo + + enddo + + end subroutine Rayleigh_fast +#endif + + + + subroutine Rayleigh_Super(dt, npx, npy, npz, ks, pm, phis, tau, u, v, w, pt, & + ua, va, delz, agrid, cp, rg, ptop, hydrostatic, & + conserve, rf_cutoff, gridstruct, domain, bd) + real, intent(in):: dt + real, intent(in):: tau !< time scale (days) + real, intent(in):: cp, rg, ptop, rf_cutoff + real, intent(in), dimension(npz):: pm + integer, intent(in):: npx, npy, npz, ks + logical, intent(in):: hydrostatic + logical, intent(in):: conserve + type(fv_grid_bounds_type), intent(IN) :: bd + real, intent(inout):: u(bd%isd:bd%ied ,bd%jsd:bd%jed+1,npz) !< D grid zonal wind (m/s) + real, intent(inout):: v(bd%isd:bd%ied+1,bd%jsd:bd%jed,npz) !< D grid meridional wind (m/s) + real, intent(inout):: w(bd%isd: ,bd%jsd: ,1: ) !< cell center vertical wind (m/s) + real, intent(inout):: pt(bd%isd:bd%ied,bd%jsd:bd%jed,npz) !< temp + real, intent(inout):: ua(bd%isd:bd%ied,bd%jsd:bd%jed,npz) ! + real, intent(inout):: va(bd%isd:bd%ied,bd%jsd:bd%jed,npz) ! + real, intent(inout):: delz(bd%isd: ,bd%jsd: ,1: ) !< delta-height (m); non-hydrostatic only + real, intent(in) :: agrid(bd%isd:bd%ied, bd%jsd:bd%jed,2) + real, intent(in) :: phis(bd%isd:bd%ied,bd%jsd:bd%jed) !< Surface geopotential (g*Z_surf) + type(fv_grid_type), intent(IN) :: gridstruct + type(domain2d), intent(INOUT) :: domain +! + real, allocatable :: u2f(:,:,:) + real, parameter:: u0 = 60. !< scaling velocity + real, parameter:: sday = 86400. + real rcv, tau0 + integer i, j, k + + integer :: is, ie, js, je + integer :: isd, ied, jsd, jed + + is = bd%is + ie = bd%ie + js = bd%js + je = bd%je + isd = bd%isd + ied = bd%ied + jsd = bd%jsd + jed = bd%jed + + rcv = 1. / (cp - rg) + + if ( .not. RF_initialized ) then +#ifdef HIWPP + allocate ( u00(is:ie, js:je+1,npz) ) + allocate ( v00(is:ie+1,js:je ,npz) ) +!$OMP parallel do default(none) shared(is,ie,js,je,npz,u00,u,v00,v) + do k=1,npz + do j=js,je+1 + do i=is,ie + u00(i,j,k) = u(i,j,k) + enddo + enddo + do j=js,je + do i=is,ie+1 + v00(i,j,k) = v(i,j,k) + enddo + enddo + enddo +#endif +#ifdef SMALL_EARTH + tau0 = tau +#else + tau0 = tau * sday +#endif + allocate( rf(npz) ) + rf(:) = 0. + + do k=1, ks+1 + if( is_master() ) write(6,*) k, 0.01*pm(k) + enddo + if( is_master() ) write(6,*) 'Rayleigh friction E-folding time (days):' + do k=1, npz + if ( pm(k) < rf_cutoff ) then + rf(k) = dt/tau0*sin(0.5*pi*log(rf_cutoff/pm(k))/log(rf_cutoff/ptop))**2 + if( is_master() ) write(6,*) k, 0.01*pm(k), dt/(rf(k)*sday) + kmax = k + else + exit + endif + enddo + RF_initialized = .true. + endif + + !$ser savepoint C2L_Ord2-In + !$ser data u=u v=v ua=ua va=va + call c2l_ord2(u, v, ua, va, gridstruct, npz, gridstruct%grid_type, bd, gridstruct%nested) + !$ser savepoint C2L_Ord2-Out + !$ser data ua=ua va=va + + allocate( u2f(isd:ied,jsd:jed,kmax) ) + +!$OMP parallel do default(none) shared(is,ie,js,je,kmax,pm,rf_cutoff,hydrostatic,ua,va,agrid, & +!$OMP u2f,rf,w) + do k=1,kmax + if ( pm(k) < rf_cutoff ) then + u2f(:,:,k) = 1. / (1.+rf(k)) + else + u2f(:,:,k) = 1. + endif + enddo + call timing_on('COMM_TOTAL') + call mpp_update_domains(u2f, domain) + call timing_off('COMM_TOTAL') + +!$OMP parallel do default(none) shared(is,ie,js,je,kmax,pm,rf_cutoff,w,rf,u,v, & +#ifdef HIWPP +!$OMP u00,v00, & +#endif +!$OMP conserve,hydrostatic,pt,ua,va,u2f,cp,rg,ptop,rcv) + do k=1,kmax + if ( pm(k) < rf_cutoff ) then +#ifdef HIWPP + if (.not. hydrostatic) then + do j=js,je + do i=is,ie + w(i,j,k) = w(i,j,k)/(1.+rf(k)) + enddo + enddo + endif + do j=js,je+1 + do i=is,ie + u(i,j,k) = (u(i,j,k)+rf(k)*u00(i,j,k))/(1.+rf(k)) + enddo + enddo + do j=js,je + do i=is,ie+1 + v(i,j,k) = (v(i,j,k)+rf(k)*v00(i,j,k))/(1.+rf(k)) + enddo + enddo +#else +! Add heat so as to conserve TE + if ( conserve ) then + if ( hydrostatic ) then + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k) + 0.5*(ua(i,j,k)**2+va(i,j,k)**2)*(1.-u2f(i,j,k)**2)/(cp-rg*ptop/pm(k)) + enddo + enddo + else + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k) + 0.5*(ua(i,j,k)**2+va(i,j,k)**2+w(i,j,k)**2)*(1.-u2f(i,j,k)**2)*rcv + enddo + enddo + endif + endif + + do j=js,je+1 + do i=is,ie + u(i,j,k) = 0.5*(u2f(i,j-1,k)+u2f(i,j,k))*u(i,j,k) + enddo + enddo + do j=js,je + do i=is,ie+1 + v(i,j,k) = 0.5*(u2f(i-1,j,k)+u2f(i,j,k))*v(i,j,k) + enddo + enddo + if ( .not. hydrostatic ) then + do j=js,je + do i=is,ie + w(i,j,k) = u2f(i,j,k)*w(i,j,k) + enddo + enddo + endif +#endif + endif + enddo + + deallocate ( u2f ) + + end subroutine Rayleigh_Super + + + subroutine Rayleigh_Friction(dt, npx, npy, npz, ks, pm, tau, u, v, w, pt, & + ua, va, delz, cp, rg, ptop, hydrostatic, conserve, & + rf_cutoff, gridstruct, domain, bd) + real, intent(in):: dt + real, intent(in):: tau !< time scale (days) + real, intent(in):: cp, rg, ptop, rf_cutoff + real, intent(in), dimension(npz):: pm + integer, intent(in):: npx, npy, npz, ks + logical, intent(in):: hydrostatic + logical, intent(in):: conserve + type(fv_grid_bounds_type), intent(IN) :: bd + real, intent(inout):: u(bd%isd:bd%ied ,bd%jsd:bd%jed+1,npz) !< D grid zonal wind (m/s) + real, intent(inout):: v(bd%isd:bd%ied+1,bd%jsd:bd%jed,npz) !< D grid meridional wind (m/s) + real, intent(inout):: w(bd%isd: ,bd%jsd: ,1: ) !< cell center vertical wind (m/s) + real, intent(inout):: pt(bd%isd:bd%ied,bd%jsd:bd%jed,npz) !< temp + real, intent(inout):: ua(bd%isd:bd%ied,bd%jsd:bd%jed,npz) ! + real, intent(inout):: va(bd%isd:bd%ied,bd%jsd:bd%jed,npz) ! + real, intent(inout):: delz(bd%isd: ,bd%jsd: ,1: ) !< delta-height (m); non-hydrostatic only + type(fv_grid_type), intent(IN) :: gridstruct + type(domain2d), intent(INOUT) :: domain +! local: + real, allocatable :: u2f(:,:,:) + real, parameter:: sday = 86400. !< seconds per day + real, parameter:: u000 = 4900. !< scaling velocity **2 + real rcv + integer i, j, k + + integer :: is, ie, js, je + integer :: isd, ied, jsd, jed + + + is = bd%is + ie = bd%ie + js = bd%js + je = bd%je + isd = bd%isd + ied = bd%ied + jsd = bd%jsd + jed = bd%jed + + + rcv = 1. / (cp - rg) + + if ( .not. RF_initialized ) then + allocate( rf(npz) ) + if( is_master() ) write(6,*) 'Rayleigh friction E-folding time (days):' + do k=1, npz + if ( pm(k) < rf_cutoff ) then + rf(k) = dt/(tau*sday)*sin(0.5*pi*log(rf_cutoff/pm(k))/log(rf_cutoff/ptop))**2 + if( is_master() ) write(6,*) k, 0.01*pm(k), dt/(rf(k)*sday) + kmax = k + else + exit + endif + enddo + RF_initialized = .true. + endif + + allocate( u2f(isd:ied,jsd:jed,kmax) ) + + call c2l_ord2(u, v, ua, va, gridstruct, npz, gridstruct%grid_type, bd, gridstruct%nested) + +!$OMP parallel do default(none) shared(is,ie,js,je,kmax,u2f,hydrostatic,ua,va,w) + do k=1,kmax + if ( hydrostatic ) then + do j=js,je + do i=is,ie + u2f(i,j,k) = ua(i,j,k)**2 + va(i,j,k)**2 + enddo + enddo + else + do j=js,je + do i=is,ie + u2f(i,j,k) = ua(i,j,k)**2 + va(i,j,k)**2 + w(i,j,k)**2 + enddo + enddo + endif + enddo + + call timing_on('COMM_TOTAL') + call mpp_update_domains(u2f, domain) + call timing_off('COMM_TOTAL') + +!$OMP parallel do default(none) shared(is,ie,js,je,kmax,conserve,hydrostatic,pt,u2f,cp,rg, & +!$OMP ptop,pm,rf,delz,rcv,u,v,w) + do k=1,kmax + + if ( conserve ) then + if ( hydrostatic ) then + do j=js,je + do i=is,ie + pt(i,j,k) = pt(i,j,k) + 0.5*u2f(i,j,k)/(cp-rg*ptop/pm(k)) & + * ( 1. - 1./(1.+rf(k)*sqrt(u2f(i,j,k)/u000))**2 ) + enddo + enddo + else + do j=js,je + do i=is,ie + delz(i,j,k) = delz(i,j,k) / pt(i,j,k) + pt(i,j,k) = pt(i,j,k) + 0.5*u2f(i,j,k) * rcv & + * ( 1. - 1./(1.+rf(k)*sqrt(u2f(i,j,k)/u000))**2 ) + delz(i,j,k) = delz(i,j,k) * pt(i,j,k) + enddo + enddo + endif + endif + + do j=js-1,je+1 + do i=is-1,ie+1 + u2f(i,j,k) = rf(k)*sqrt(u2f(i,j,k)/u000) + enddo + enddo + + do j=js,je+1 + do i=is,ie + u(i,j,k) = u(i,j,k) / (1.+0.5*(u2f(i,j-1,k)+u2f(i,j,k))) + enddo + enddo + do j=js,je + do i=is,ie+1 + v(i,j,k) = v(i,j,k) / (1.+0.5*(u2f(i-1,j,k)+u2f(i,j,k))) + enddo + enddo + + if ( .not. hydrostatic ) then + do j=js,je + do i=is,ie + w(i,j,k) = w(i,j,k) / (1.+u2f(i,j,k)) + enddo + enddo + endif + + enddo + + deallocate ( u2f ) + + end subroutine Rayleigh_Friction + +!>@brief The subroutine 'compute_aam' computes vertically (mass) integrated Atmospheric Angular Momentum. + subroutine compute_aam(npz, is, ie, js, je, isd, ied, jsd, jed, gridstruct, bd, & + ptop, ua, va, u, v, delp, aam, ps, m_fac) +! Compute vertically (mass) integrated Atmospheric Angular Momentum + integer, intent(in):: npz + integer, intent(in):: is, ie, js, je + integer, intent(in):: isd, ied, jsd, jed + real, intent(in):: ptop + real, intent(inout):: u(isd:ied ,jsd:jed+1,npz) !< D grid zonal wind (m/s) + real, intent(inout):: v(isd:ied+1,jsd:jed,npz) !< D grid meridional wind (m/s) + real, intent(inout):: delp(isd:ied,jsd:jed,npz) + real, intent(inout), dimension(isd:ied,jsd:jed, npz):: ua, va + real, intent(out):: aam(is:ie,js:je) + real, intent(out):: m_fac(is:ie,js:je) + real, intent(out):: ps(isd:ied,jsd:jed) + type(fv_grid_bounds_type), intent(IN) :: bd + type(fv_grid_type), intent(IN) :: gridstruct +! local: + real, dimension(is:ie):: r1, r2, dm + integer i, j, k + + call c2l_ord2(u, v, ua, va, gridstruct, npz, gridstruct%grid_type, bd, gridstruct%nested) + +!$OMP parallel do default(none) shared(is,ie,js,je,npz,gridstruct,aam,m_fac,ps,ptop,delp,agrav,ua) & +!$OMP private(r1, r2, dm) + do j=js,je + do i=is,ie + r1(i) = radius*cos(gridstruct%agrid(i,j,2)) + r2(i) = r1(i)*r1(i) + aam(i,j) = 0. + m_fac(i,j) = 0. + ps(i,j) = ptop + enddo + do k=1,npz + do i=is,ie + dm(i) = delp(i,j,k) + ps(i,j) = ps(i,j) + dm(i) + dm(i) = dm(i)*agrav + aam(i,j) = aam(i,j) + (r2(i)*omega + r1(i)*ua(i,j,k)) * dm(i) + m_fac(i,j) = m_fac(i,j) + dm(i)*r2(i) + enddo + enddo + enddo + + end subroutine compute_aam + +end module fv_dynamics_mod From 87c7f28c1f5c1613c95200abf143ee8d8915d2ec Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 31 Jan 2025 08:53:54 -0800 Subject: [PATCH 158/252] Made specific tracer parameter calls within remapping as calls from the tracer dictionary. Note that the tracer key names right now need to match up with 'tracer_varaibles' from gt4py_utils.py --- .../translate/translate_remapping_GEOS.py | 98 +++++++++---------- 1 file changed, 49 insertions(+), 49 deletions(-) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 6bd724be..6a5d49a6 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -757,19 +757,19 @@ def __init__( self._rescale_delz_1 = stencil_factory.from_origin_domain( rescale_delz_1, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._rescale_delz_2 = stencil_factory.from_origin_domain( rescale_delz_2, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._w_fix_consrv_moment = stencil_factory.from_origin_domain( func=W_fix_consrv_moment, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) @@ -803,31 +803,31 @@ def __init__( self._moist_cv_pkz = stencil_factory.from_origin_domain( moist_cv.moist_pkz, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._moist_cv_te = stencil_factory.from_origin_domain( moist_cv.moist_te, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._te_zsum = stencil_factory.from_origin_domain( moist_cv.te_zsum, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) self._moist_cv_pt_last_step = stencil_factory.from_origin_domain( moist_cv.moist_pt_last_step, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(0, 0, 1)), ) self._fill_cond = stencil_factory.from_origin_domain( moist_cv.cond_output, - origin=grid.compute_origin(), + origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(), ) @@ -838,6 +838,18 @@ def compute_parallel(self, inputs, communicator): state = self.state_from_inputs(inputs) state_namespace = SimpleNamespace(**state) + tracers = { + "qvapor": state_namespace.qvapor, + "qliquid": state_namespace.qliquid, + "qice": state_namespace.qice, + "qrain": state_namespace.qrain, + "qsnow": state_namespace.qsnow, + "qgraupel": state_namespace.qgraupel, + "qcld": state_namespace.qcld, + "qo3mr": state_namespace.qo3mr, + "qsgs_tke": state_namespace.qsgs_tke, + } + self._init_pe( state_namespace.pe_, self._pe1, @@ -846,12 +858,12 @@ def compute_parallel(self, inputs, communicator): ) self._moist_cv_pt_pressure( - state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, + tracers["qvapor"], + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], state_namespace.q_con, state_namespace.pt, state_namespace.cappa, @@ -885,18 +897,6 @@ def compute_parallel(self, inputs, communicator): interp=True, ) - tracers = { - "qvapor": state_namespace.qvapor, - "qliquid": state_namespace.qliquid, - "qice": state_namespace.qice, - "qrain": state_namespace.qrain, - "qsnow": state_namespace.qsnow, - "qgraupel": state_namespace.qgraupel, - "qcld": state_namespace.qcld, - "qo3mr": state_namespace.qo3mr, - "qsgs_tke": state_namespace.qsgs_tke, - } - self._mapn_tracer = MapNTracer( self.stencil_factory, self.quantity_factory, @@ -1048,12 +1048,12 @@ def compute_parallel(self, inputs, communicator): ) self._moist_cv_pkz( - state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, + tracers["qvapor"], + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], state_namespace.pkz, state_namespace.pt, state_namespace.cappa, @@ -1067,12 +1067,12 @@ def compute_parallel(self, inputs, communicator): if state_namespace.consv > state_namespace.consv_min: self._moist_cv_te( - state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, + tracers["qvapor"], + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], state_namespace.u, state_namespace.v, state_namespace.w, @@ -1120,12 +1120,12 @@ def compute_parallel(self, inputs, communicator): if state_namespace.last_step and not state_namespace.adiabatic: self._moist_cv_pt_last_step( - state_namespace.qvapor, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, + tracers["qvapor"], + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], state_namespace.pt, state_namespace.pkz, Float(dtmp), @@ -1134,11 +1134,11 @@ def compute_parallel(self, inputs, communicator): self._fill_cond( state_namespace.q_con, - state_namespace.qliquid, - state_namespace.qrain, - state_namespace.qsnow, - state_namespace.qice, - state_namespace.qgraupel, + tracers["qliquid"], + tracers["qrain"], + tracers["qsnow"], + tracers["qice"], + tracers["qgraupel"], ) return self.outputs_from_state(state) From e7596f9bb32beb28dda854a3245e2eab8eb5eea7 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 3 Feb 2025 16:32:17 -0800 Subject: [PATCH 159/252] Reorganizing code and accessing constants from PyFV3 and NDSL instead of serializing --- fv_mapz.F90.SER | 35 ++- .../translate/translate_remapping_GEOS.py | 201 +++++++++--------- 2 files changed, 111 insertions(+), 125 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 9aad009b..77ce438c 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -396,13 +396,12 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser verbatim enddo -!$ser verbatim qmin=t_min -!$ser verbatim w_max_ = w_max -!$ser verbatim w_min_ = w_min -!$ser verbatim rsin2=gridstruct%rsin2 -!$ser verbatim cosa_s=gridstruct%cosa_s -!$ser verbatim grav_=grav -!$ser verbatim area_64_ = real(gridstruct%area_64) +!!$ser verbatim qmin=t_min +!!$ser verbatim w_max_ = w_max +!!$ser verbatim w_min_ = w_min +!!$ser verbatim rsin2=gridstruct%rsin2 +!!$ser verbatim grav_=grav +!!$ser verbatim area_64_ = real(gridstruct%area_64) !$ser verbatim te0_2d_f32 = real(te0_2d) !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d @@ -430,8 +429,8 @@ contains !$ser data pk=pk !$ser data pkz=pkz !$ser data akap=akap -!$ser data t_min=qmin -!$ser data kord_wz=kord_wz +!!$ser data t_min=qmin +!!$ser data kord_wz=kord_wz !$ser data ws_=ws !$ser data w=w !$ser data u=u @@ -448,20 +447,19 @@ contains !$ser verbatim if (present(cx)) then !$ser data cx_=cx !$ser verbatim endif -!$ser data w_max=w_max_ -!$ser data w_min=w_min_ -!$ser data kord_mt=kord_mt -!$ser data cosa_s=cosa_s -!$ser data rsin2=rsin2 +!!$ser data w_max=w_max_ +!!$ser data w_min=w_min_ +!!$ser data kord_mt=kord_mt +!!$ser data rsin2=rsin2 !$ser data hs=hs -!$ser data grav=grav_ +!!$ser data grav=grav_ !$ser data te0_2d_=te0_2d_f32 -!$ser data area_64_=area_64_ +!!$ser data area_64_=area_64_ !$ser data last_step=last_step !$ser data do_adiabatic_init=do_adiabatic_init !$ser data consv=consv -!$ser data consv_min=consv_min_ -!$ser data cv_air=cv_air_ +!!$ser data consv_min=consv_min_ +!!$ser data cv_air=cv_air_ !$ser data adiabatic=adiabatic do 1000 j=js,je+1 @@ -1784,7 +1782,6 @@ contains !$ser data cx_=cx !$ser verbatim endif !$ser data kord_mt=kord_mt -!$ser data cosa_s=cosa_s !$ser data rsin2=rsin2 !$ser data hs=hs !$ser data grav=grav_ diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 6a5d49a6..0ed9cea3 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -2,6 +2,8 @@ from ndsl import Namelist, StencilFactory from ndsl.constants import ( + CV_AIR, + GRAV, X_DIM, X_INTERFACE_DIM, Y_DIM, @@ -12,11 +14,14 @@ from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig + +# from pyFV3._config import RemappingConfig from pyFV3.stencils import moist_cv from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.mpp_global_sum import mpp_global_sum from pyFV3.stencils.remapping import ( + CONSV_MIN, init_pe, moist_cv_pt_pressure, pe0_ptop_xmax, @@ -176,16 +181,11 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "cosa_s": { - "name": "cosa_s", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, - "rsin2": { - "name": "rsin2", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, + # "rsin2": { + # "name": "rsin2", + # "dims": [X_DIM, Y_DIM], + # "units": "No Units", + # }, "hs": { "name": "hs", "dims": [X_DIM, Y_DIM], @@ -196,11 +196,11 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "area_64_": { - "name": "area_64_", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, + # "area_64_": { + # "name": "area_64_", + # "dims": [X_DIM, Y_DIM], + # "units": "No Units", + # }, } outputs = { "pt": { @@ -452,18 +452,18 @@ def __init__( "jend": grid.jed, "kend": grid.npz - 1, }, - "cosa_s": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.jsd, - "jend": grid.jed, - }, - "rsin2": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.jsd, - "jend": grid.jed, - }, + # "cosa_s": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed, + # }, + # "rsin2": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed, + # }, "hs": { "istart": grid.isd, "iend": grid.ied, @@ -476,28 +476,23 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, - "area_64_": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.jsd, - "jend": grid.jed, - }, + # "area_64_": { + # "istart": grid.isd, + # "iend": grid.ied, + # "jstart": grid.jsd, + # "jend": grid.jed, + # }, } self._base.in_vars["parameters"] = [ "ptop", "r_vir", "akap", - "t_min", - "kord_wz", - "w_max", - "w_min", - "kord_mt", - "grav", + # "grav", "last_step", "do_adiabatic_init", "consv", - "consv_min", - "cv_air", + # "consv_min", + # "cv_air", "adiabatic", ] self._base.out_vars = { @@ -619,17 +614,18 @@ def __init__( raise NotImplementedError("Hydrostatic is not implemented") grid_indexing = stencil_factory.grid_indexing - self._domain_jextra = ( - grid_indexing.domain[0], - grid_indexing.domain[1] + 1, - grid_indexing.domain[2] + 1, - ) - - # Value from GEOS - self.kord = 9 - # Value from GEOS - self._kord_tm = 9 + self._t_min = Float(184.0) + self._w_max = Float(90.0) + self._w_min = Float(-60.0) + self._area_64 = self.grid.grid_data.area_64 + self._cosa_s = self.grid.grid_data.cosa_s + self._rsin2 = self.grid.grid_data.rsin2 + self._kord_tr = config.kord_tr + self._kord_tm = abs(config.kord_tm) + self._kord_wz = config.kord_wz + self._kord_mt = config.kord_mt + self._do_sat_adjust = config.do_sat_adj # mode / iv set to 1 from GEOS self.mode = 1 @@ -831,6 +827,38 @@ def __init__( domain=grid_indexing.domain_compute(), ) + self._map1_ppm_u = MapSingle( + self.stencil_factory, + self.quantity_factory, + self._kord_mt, + -1, + dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + ) + + self._map1_ppm_v = MapSingle( + self.stencil_factory, + self.quantity_factory, + self._kord_mt, + -1, + dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + ) + + self._map_single_w = MapSingle( + self.stencil_factory, + self.quantity_factory, + self._kord_wz, + -2, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + + self._map_single_delz = MapSingle( + self.stencil_factory, + self.quantity_factory, + self._kord_wz, + 1, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") @@ -893,14 +921,14 @@ def compute_parallel(self, inputs, communicator): state_namespace.pt, self._pn1, self._pn2, - qmin=state_namespace.t_min, + self._t_min, interp=True, ) self._mapn_tracer = MapNTracer( self.stencil_factory, self.quantity_factory, - abs(self.kord), + abs(self._kord_tr), self.nq, fill=self.fill, tracers=tracers, @@ -908,21 +936,6 @@ def compute_parallel(self, inputs, communicator): self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) - self._map_single_w = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_wz, - -2, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - - self._map_single_delz = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_wz, - 1, - dims=[X_DIM, Y_DIM, Z_DIM], - ) self._map_single_w( state_namespace.w, self._pe1, @@ -948,19 +961,11 @@ def compute_parallel(self, inputs, communicator): self._w2, self._dp2, self._gz, - state_namespace.w_max, - state_namespace.w_min, + self._w_max, + self._w_min, self._compute_performed, ) - self._map1_ppm_u = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_mt, - -1, - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], - ) - self._pressures_mapu( state_namespace.pe_, state_namespace.ak, @@ -996,14 +1001,6 @@ def compute_parallel(self, inputs, communicator): interp=False, ) - self._map1_ppm_v = MapSingle( - self.stencil_factory, - self.quantity_factory, - state_namespace.kord_mt, - -1, - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], - ) - self._pressures_mapv( state_namespace.pe_, state_namespace.ak, @@ -1064,7 +1061,7 @@ def compute_parallel(self, inputs, communicator): if state_namespace.last_step and not state_namespace.do_adiabatic_init: - if state_namespace.consv > state_namespace.consv_min: + if state_namespace.consv > CONSV_MIN: self._moist_cv_te( tracers["qvapor"], @@ -1080,11 +1077,11 @@ def compute_parallel(self, inputs, communicator): state_namespace.pt, self._phis, state_namespace.delp, - state_namespace.rsin2, - state_namespace.cosa_s, + self._rsin2, + self._cosa_s, state_namespace.hs, state_namespace.delz, - state_namespace.grav, + GRAV, ) self._te_zsum( @@ -1096,24 +1093,16 @@ def compute_parallel(self, inputs, communicator): ) # Note, mpp_global_sum is currently set up for the C24 TBC setup - inputArray = ( - self._te_2d.data[0:-1, 0:-1] - * state_namespace.area_64_.data[0:-1, 0:-1] - ) - tesum = mpp_global_sum( - inputArray[3:27, 3:27], communicator, self.stencil_factory - ) + inputArray = self._te_2d.view[:] * self._area_64.view[:] + tesum = mpp_global_sum(inputArray, communicator, self.stencil_factory) # print("tesum: ", tesum) - inputArray = ( - self._zsum1.data[0:-1, 0:-1] - * state_namespace.area_64_.data[0:-1, 0:-1] - ) - zsum = mpp_global_sum( - inputArray[3:27, 3:27], communicator, self.stencil_factory - ) + # print("type(self._zsum1) = ",type(self._zsum1)) + # print("type(self._area_64) = ",type(self._area_64)) + inputArray = self._zsum1.view[:] * self._area_64.view[:] + zsum = mpp_global_sum(inputArray, communicator, self.stencil_factory) # print("zsum: ", zsum) - dtmp = tesum / (state_namespace.cv_air.data * zsum) - # print("dtmp: ", dtmp) + dtmp = tesum / (CV_AIR * zsum) + print("dtmp: ", dtmp) # I ignore the E_flux computation since it's not used elsewhere # in our current setup once it's computed From 3ad0b51f65a95f168dd3998bdb95c0bf58a9645a Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Mon, 3 Feb 2025 17:12:15 -0800 Subject: [PATCH 160/252] Code cleanup --- fv_mapz.F90.SER | 27 ++---------- .../translate/translate_remapping_GEOS.py | 41 ++----------------- 2 files changed, 7 insertions(+), 61 deletions(-) diff --git a/fv_mapz.F90.SER b/fv_mapz.F90.SER index 77ce438c..7b479895 100644 --- a/fv_mapz.F90.SER +++ b/fv_mapz.F90.SER @@ -230,11 +230,11 @@ contains real, dimension(is:ie, js:je, km) :: w2_3d, dp2_3d real, dimension(is:ie, js:je) :: gz_2d, te_2d_f32, te0_2d_f32, zsum1_f32 - real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s, area_64_ + real, dimension(isd:ied,jsd:jed) :: rsin2, cosa_s logical :: serial_flag - real :: grav_, consv_min_, cv_air_ + real :: grav_ !$ser verbatim real :: w_max_, w_min_ @@ -255,8 +255,6 @@ contains !$ser verbatim pk2_3d = pk !$ser verbatim gz_2d = 0.0 !$ser verbatim phis_3d = 0.0 - !$ser verbatim consv_min_=consv_min - !$ser verbatim cv_air_=cv_air remap_t = .false. remap_pt = .false. @@ -380,7 +378,7 @@ contains !$OMP pe_3d, pe1_3d, pe2_3d, pn1_3d, pn2_3d, pk2_3d, w2_3d, dp2_3d, & !$OMP peln_3d, gz_2d, pe0_3d, pe3_3d, phis_3d, te_2d, te_2d_f32, te0_2d, te0_2d_f32) & !$OMP private(gz,cvm,bkh,dp2, & -!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s, grav_, area_64_) +!$OMP pe0,pe1,pe2,pe3,pk1,pk2,pn1,pn2,phis,q2,w2,dpln,dlnp, w_max_, w_min_, rsin2, cosa_s) !$ser verbatim do k = 1, km+1 !$ser verbatim do j = js, je @@ -396,12 +394,6 @@ contains !$ser verbatim enddo !$ser verbatim enddo !$ser verbatim enddo -!!$ser verbatim qmin=t_min -!!$ser verbatim w_max_ = w_max -!!$ser verbatim w_min_ = w_min -!!$ser verbatim rsin2=gridstruct%rsin2 -!!$ser verbatim grav_=grav -!!$ser verbatim area_64_ = real(gridstruct%area_64) !$ser verbatim te0_2d_f32 = real(te0_2d) !$ser savepoint Remapping_GEOS-In !$ser data pe_=pe_3d @@ -429,8 +421,6 @@ contains !$ser data pk=pk !$ser data pkz=pkz !$ser data akap=akap -!!$ser data t_min=qmin -!!$ser data kord_wz=kord_wz !$ser data ws_=ws !$ser data w=w !$ser data u=u @@ -447,20 +437,13 @@ contains !$ser verbatim if (present(cx)) then !$ser data cx_=cx !$ser verbatim endif -!!$ser data w_max=w_max_ -!!$ser data w_min=w_min_ -!!$ser data kord_mt=kord_mt -!!$ser data rsin2=rsin2 !$ser data hs=hs -!!$ser data grav=grav_ !$ser data te0_2d_=te0_2d_f32 -!!$ser data area_64_=area_64_ !$ser data last_step=last_step !$ser data do_adiabatic_init=do_adiabatic_init !$ser data consv=consv -!!$ser data consv_min=consv_min_ -!!$ser data cv_air=cv_air_ !$ser data adiabatic=adiabatic +!$ser data nq=nq do 1000 j=js,je+1 @@ -1782,9 +1765,7 @@ contains !$ser data cx_=cx !$ser verbatim endif !$ser data kord_mt=kord_mt -!$ser data rsin2=rsin2 !$ser data hs=hs -!$ser data grav=grav_ !$ser data te0_2d_=te0_2d_f32 !$ser savepoint GetMPIProp-In diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 0ed9cea3..70d3892f 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -181,11 +181,6 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - # "rsin2": { - # "name": "rsin2", - # "dims": [X_DIM, Y_DIM], - # "units": "No Units", - # }, "hs": { "name": "hs", "dims": [X_DIM, Y_DIM], @@ -196,11 +191,6 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - # "area_64_": { - # "name": "area_64_", - # "dims": [X_DIM, Y_DIM], - # "units": "No Units", - # }, } outputs = { "pt": { @@ -452,18 +442,6 @@ def __init__( "jend": grid.jed, "kend": grid.npz - 1, }, - # "cosa_s": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.jsd, - # "jend": grid.jed, - # }, - # "rsin2": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.jsd, - # "jend": grid.jed, - # }, "hs": { "istart": grid.isd, "iend": grid.ied, @@ -476,24 +454,16 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, - # "area_64_": { - # "istart": grid.isd, - # "iend": grid.ied, - # "jstart": grid.jsd, - # "jend": grid.jed, - # }, } self._base.in_vars["parameters"] = [ "ptop", "r_vir", "akap", - # "grav", "last_step", "do_adiabatic_init", "consv", - # "consv_min", - # "cv_air", "adiabatic", + "nq", ] self._base.out_vars = { "pt": {}, @@ -627,11 +597,6 @@ def __init__( self._kord_mt = config.kord_mt self._do_sat_adjust = config.do_sat_adj - # mode / iv set to 1 from GEOS - self.mode = 1 - - self.nq = 9 - self.fill = True self._gz = self.quantity_factory.zeros( @@ -747,7 +712,7 @@ def __init__( self.stencil_factory, self.quantity_factory, self._kord_tm, - self.mode, + 1, dims=[X_DIM, Y_DIM, Z_DIM], ) @@ -929,7 +894,7 @@ def compute_parallel(self, inputs, communicator): self.stencil_factory, self.quantity_factory, abs(self._kord_tr), - self.nq, + state_namespace.nq, fill=self.fill, tracers=tracers, ) From a1498b21f389e352f760db42099d7f87734a95cf Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 4 Feb 2025 11:45:55 -0800 Subject: [PATCH 161/252] Laying groundwork to use translate test to test LangrangianToEulerian --- pyFV3/tracers.py | 132 +++++ .../translate/translate_remapping_GEOS.py | 3 +- .../translate/translate_remapping_GEOS_v2.py | 464 ++++++++++++++++++ 3 files changed, 597 insertions(+), 2 deletions(-) create mode 100644 pyFV3/tracers.py create mode 100644 tests/savepoint/translate/translate_remapping_GEOS_v2.py diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py new file mode 100644 index 00000000..4b938c5f --- /dev/null +++ b/pyFV3/tracers.py @@ -0,0 +1,132 @@ +from __future__ import annotations + +from typing import Dict, List + +import numpy as np + +from ndsl import Quantity, QuantityFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM + + +# FOR REFERENCE - previous descriptive of the tracers, lining up with Pace work +# tracer_variables = [ +# "qvapor", +# "qliquid", +# "qrain", +# "qice", +# "qsnow", +# "qgraupel", +# "qo3mr", +# "qsgs_tke", +# "qcld", +# ] + + +class Tracers: + unit = "g/kg" + dims = [X_DIM, Y_DIM, Z_DIM] + + def __init__(self, factory: QuantityFactory) -> None: + self._quantities: Dict[str, Quantity] = {} + self._quantity_factory = factory + + def copy_tracer_data( + self, + name: str, + data: np.ndarray, + unit="unknown", + ): + qty = self._quantity_factory.empty(dims=self.dims, units=unit) + if data.shape > qty.data.shape: + raise ValueError( + f"[pyFV3] Tracer {name} size ({data.shape}" + f" is bigger than grid {qty.data.shape})" + ) + qty.data[: data.shape[0], : data.shape[1], : data.shape[2]] = data + self._quantities[name] = qty + + @property + def count(self) -> int: + return len(self._quantities) + + def values(self): + return self._quantities.values() + + def names(self): + return self._quantities.keys() + + def items(self): + return self._quantities.items() + + def as_4D_array(self) -> np.ndarray: + shape = self._quantity_factory.sizer.get_shape(self.dims) + var4d = np.empty( + ( + shape[0] - 1, + shape[1] - 1, + shape[2] - 1, + self.count, + ) + ) + # Skip the extra data point that is meant to align interface + # and non interface fields + for idx, q in enumerate(self.values()): + var4d[:, :, :, idx] = q.data[:-1, :-1, :-1] + return var4d + + def __getitem__(self, key): + return self._quantities[key] + + def __setitem__(self, key, value): + self._quantities[key] = value + + def __str__(self) -> str: + return self.__repr__() + + def __repr__(self) -> str: + msg = f"[pyFV3] {self.count} Tracers:\n" + for q in self.names(): + msg += f" {q}\n" + return msg + + @classmethod + def make( + cls, + quantity_factory: QuantityFactory, + tracer_mapping: List[str], + ): + tracers = cls(quantity_factory) + for name in tracer_mapping: + qty = quantity_factory.empty(dims=cls.dims, units=cls.unit) + tracers._quantities[name] = qty + return tracers + + @staticmethod + def blind_mapping_from_data(tracer_data: np.ndarray): + if len(tracer_data.shape) != 4: + raise ValueError("Expected 4D field as input") + return [f"Tracer_{idx}" for idx in range(tracer_data.shape[3])] + + @classmethod + def make_from_4D_array( + cls, + quantity_factory: QuantityFactory, + tracer_mapping: List[str], + tracer_data: np.ndarray, + ) -> Tracers: + if len(tracer_data.shape) != 4: + raise ValueError("Expected 4D field as input") + count = len(tracer_mapping) + if count > tracer_data.shape[3]: + raise ValueError( + f"Mapping size {len(tracer_mapping)} is bigger than" + f" data dimensionality {tracer_data.shape[3]}" + ) + tracers = cls(quantity_factory) + for idx in range(0, count): + tracers.copy_tracer_data( + name=tracer_mapping[idx] or f"Tracer_{idx}", + data=tracer_data[:, :, :, idx], + unit=cls.unit, + ) + return tracers diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 70d3892f..c92f519f 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -575,8 +575,7 @@ def __init__( self.stencil_factory = stencil_factory self.quantity_factory = grid.quantity_factory - # self.namelist found in TranslateDycoreFortranData2Py - # self.namelist = DynamicalCoreConfig.from_namelist(namelist) + # self.namelist found in ParallelTranslateBaseSlicing config = DynamicalCoreConfig.from_namelist(self.namelist).remapping hydrostatic = config.hydrostatic diff --git a/tests/savepoint/translate/translate_remapping_GEOS_v2.py b/tests/savepoint/translate/translate_remapping_GEOS_v2.py new file mode 100644 index 00000000..2148bb46 --- /dev/null +++ b/tests/savepoint/translate/translate_remapping_GEOS_v2.py @@ -0,0 +1,464 @@ +from types import SimpleNamespace + +from ndsl import Namelist, StencilFactory +from ndsl.constants import ( + X_DIM, + X_INTERFACE_DIM, + Y_DIM, + Y_INTERFACE_DIM, + Z_DIM, + Z_INTERFACE_DIM, +) +from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing +from pyFV3 import DynamicalCoreConfig +from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS + + +# from pyFV3._config import RemappingConfig + + +class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): + inputs = { + "pe_": { + "name": "pe_", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "No Units", + }, + "qvapor": { + "name": "qvapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qliquid": { + "name": "qliquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qice": { + "name": "qice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qrain": { + "name": "qrain", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsnow": { + "name": "qsnow", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qgraupel": { + "name": "qgraupel", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qcld": { + "name": "qcld", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qo3mr": { + "name": "qo3mr", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsgs_tke": { + "name": "qsgs_tke", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delp": { + "name": "delp", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delz": { + "name": "delz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "q_con": { + "name": "q_con", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "pt": { + "name": "pt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cappa": { + "name": "cappa", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "ps": { + "name": "ps", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "peln_3d": { + "name": "peln_3d", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "No Units", + }, + "ak": { + "name": "ak", + "dims": [Z_INTERFACE_DIM], + "units": "No Units", + }, + "bk": { + "name": "bk", + "dims": [Z_INTERFACE_DIM], + "units": "No Units", + }, + "pk": { + "name": "pk", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "No Units", + }, + "pkz": { + "name": "pkz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "w": { + "name": "w", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "ws_": { + "name": "ws_", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "u": { + "name": "u", + "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], + "units": "No Units", + }, + "v": { + "name": "v", + "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfy": { + "name": "mfy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cy": { + "name": "cy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfx_": { + "name": "mfx_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cx_": { + "name": "cx_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "hs": { + "name": "hs", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "te0_2d_": { + "name": "te0_2d_", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + } + outputs = { + "pt": { + "name": "pt", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cappa": { + "name": "cappa", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delp": { + "name": "delp", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "delz": { + "name": "delz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qvapor": { + "name": "qvapor", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qliquid": { + "name": "qliquid", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qice": { + "name": "qice", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qrain": { + "name": "qrain", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsnow": { + "name": "qsnow", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qgraupel": { + "name": "qgraupel", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qcld": { + "name": "qcld", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qo3mr": { + "name": "qo3mr", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "qsgs_tke": { + "name": "qsgs_tke", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "w": { + "name": "w", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "u": { + "name": "u", + "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], + "units": "No Units", + }, + "v": { + "name": "v", + "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfy": { + "name": "mfy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cy": { + "name": "cy", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "mfx_": { + "name": "mfx_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "cx_": { + "name": "cx_", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "peln_3d": { + "name": "peln_3d", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "No Units", + }, + "pe_": { + "name": "pe_", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "No Units", + }, + "pk": { + "name": "pk", + "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "units": "No Units", + }, + "pkz": { + "name": "pkz", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "q_con": { + "name": "q_con", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + } + + def __init__( + self, + grid: Grid, + namelist: Namelist, + stencil_factory: StencilFactory, + ): + super().__init__(grid, namelist, stencil_factory) + + self._base.in_vars["data_vars"] = { + "tracers": {}, + "w": { + "kend": grid.npz - 1, + }, + "u": grid.y3d_domain_dict(), + "v": grid.x3d_domain_dict(), + "delz": {}, + "pt": {}, + "dp1": {}, + "delp": {}, + "cappa": {}, + "q_con": {}, + "pkz": grid.compute_dict(), + "pk": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz, + }, + "peln": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + "kaxis": 1, + "kend": grid.npz, + }, + "pe": { + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, + "kend": grid.npz + 1, + "kaxis": 1, + }, + "hs": {"serialname": "phis"}, + "ps": {}, + "wsd": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, + # column variables... + "ak": {}, + "bk": {}, + "pfull": {}, + } + self._base.in_vars["parameters"] = [ + "ptop", + "akap", + "zvir", + "last_step", + "consv_te", + "mdt", + "nq", + ] + self._base.out_vars = {} + for k in [ + "pe", + "pkz", + "pk", + "peln", + "pt", + "tracers", + "cappa", + "delp", + "delz", + "q_con", + "u", + "v", + "w", + "ps", + "dp1", + ]: + self._base.out_vars[k] = self._base.in_vars["data_vars"][k] + + self.stencil_factory = stencil_factory + self.quantity_factory = grid.quantity_factory + + self.stencil_factory = stencil_factory + self.namelist = DynamicalCoreConfig.from_namelist(namelist) + self.grid = grid + + def compute_sequential(self, inputs_list, communicator_list): + print("No serial test available") + + def compute_parallel(self, inputs, communicator): + state = self.state_from_inputs(inputs) + state_namespace = SimpleNamespace(**state) + + tracers = { + "qvapor": state_namespace.qvapor, + "qliquid": state_namespace.qliquid, + "qice": state_namespace.qice, + "qrain": state_namespace.qrain, + "qsnow": state_namespace.qsnow, + "qgraupel": state_namespace.qgraupel, + "qcld": state_namespace.qcld, + "qo3mr": state_namespace.qo3mr, + "qsgs_tke": state_namespace.qsgs_tke, + } + + l_to_e = LagrangianToEulerian_GEOS( + self.stencil_factory, + self.quantity_factory, + DynamicalCoreConfig.from_namelist(self.namelist).remapping, + communicator, + self.grid.grid_data, + state_namespace.nq, + state_namespace.pfull, + tracers, + state_namespace.adiabatic, + ) + + l_to_e( + tracers, + state_namespace.pt, + state_namespace.delp, + state_namespace.delz, + state_namespace.peln, + state_namespace.u, + state_namespace.v, + state_namespace.w, + state_namespace.mfx, + state_namespace.mfy, + state_namespace.cx, + state_namespace.cy, + state_namespace.cappa, + state_namespace.q_con, + state_namespace.pkz, + state_namespace.pk, + state_namespace.pe, + state_namespace.hs, + state_namespace.te0_2d, + state_namespace.ps, + state_namespace.wsd, + state_namespace.ak, + state_namespace.bk, + state_namespace.dp1, + state_namespace.ptop, + state_namespace.akap, + state_namespace.zvir, + state_namespace.last_step, + state_namespace.consv_te, + state_namespace.mdt, + ) + + return self.outputs_from_state(state) From ad5cba213305418ea9282e4a01a7b432920177bd Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 4 Feb 2025 18:05:13 -0800 Subject: [PATCH 162/252] Continue development of LtoE translate test --- fv_dynamics.F90.SER | 37 ++- tests/savepoint/translate/__init__.py | 1 + .../translate/translate_remapping_GEOS_v2.py | 275 ++++++++++++------ 3 files changed, 221 insertions(+), 92 deletions(-) diff --git a/fv_dynamics.F90.SER b/fv_dynamics.F90.SER index fbe9da9c..1f8cae76 100644 --- a/fv_dynamics.F90.SER +++ b/fv_dynamics.F90.SER @@ -794,11 +794,11 @@ contains !$ser savepoint Remapping-In !$ser data iq=iq last_step=last_step consv_te=consv_te ps=ps pe=pe delp=delp pkz=pkz pk=pk mdt=mdt bdt=bdt sphum=sphum q_con=q_con u=u v=v w=w delz=delz pt=pt tracers=q phis=phis zvir=zvir akap=akap cappa=cappa kord_tracer=kord_tracer peln=peln te_2d=te_2d ua=ua va=va omga=omga dp1=dp1 wsd=ws reproduce_sum=reproduce_sum ptop=ptop ak=ak bk=bk pfull=pfull hybrid_z=hybrid_z do_adiabatic_init=do_adiabatic_init nq=nq - !$ser savepoint Lagrangian_to_Eulerian-In - !$ser data last_step=last_step consv=consv_te ps=ps pe=pe delp=delp pkz=pkz mdt=mdt pdt=bdt - !$ser data km=npz nq=nq nwat=nwat sphum=sphum q_con=q_con u=u v=v w=w delz=delz pt=pt q_=q - !$ser data hs=phis r_vir=zvir cp=cp_air akap=akap cappa=cappa kord_tr=kord_tracer - !$ser data peln=peln te0_2d=te_2d ng=ng ua=ua va=va omga=omga te=dp1 ws_=ws fill=fill + !$ser savepoint Remapping_GEOS_v2-In + !$ser data last_step=last_step consv_te=consv_te ps=ps pe=pe delp=delp pkz=pkz mdt=mdt pdt=bdt + !$ser data km=npz nq=nq nwat=nwat sphum=sphum q_con=q_con u=u v=v w=w delz=delz pt=pt tracers=q + !$ser data hs=phis zvir=zvir cp=cp_air akap=akap cappa=cappa kord_tr=kord_tracer + !$ser data peln=peln te0_2d=te_2d ng=ng ua=ua va=va omga=omga dp1=dp1 wsd=ws fill=fill !$ser data reproduce_sum=reproduce_sum dtdt=dtdt_m ptop=ptop ak=ak bk=bk pfull=pfull !$ser data hydrostatic=hydrostatic hybrid_z=hybrid_z do_omega=do_omega !$ser data do_adiabatic_init=do_adiabatic_init @@ -806,6 +806,16 @@ contains !$ser data mfy=mfyL !$ser data cx=cxL !$ser data cy=cyL + !$ser data qvapor=q(:,:,:,sphum) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) + !$ser data pk=pk call Lagrangian_to_Eulerian(last_step, consv_te, ps, pe, delp, & pkz, pk, mdt, bdt, npz, is,ie,js,je, isd,ied,jsd,jed, & @@ -819,13 +829,24 @@ contains mfxL, mfyL, cxL, cyL, flagstruct%remap_option, flagstruct%gmao_remap) !$ser savepoint Remapping-Out !$ser data te_2d=te_2d pk=pk tracers=q delp=delp pe=pe ps=ps u=u v=v w=w pt=pt delz=delz q_con=q_con cappa=cappa ua=ua va=va omga=omga peln=peln pkz=pkz dp1=dp1 - !$ser savepoint Lagrangian_to_Eulerian-Out - !$ser data ps=ps pe=pe pkz=pkz q_con=q_con u=u v=v w=w delz=delz pt=pt q_=q cappa=cappa - !$ser data peln=peln te0_2d=te_2d ua=ua va=va omga=omga te=dp1 dtdt=dtdt_m + !$ser savepoint Remapping_GEOS_v2-Out + !$ser data ps=ps pe=pe pkz=pkz q_con=q_con u=u v=v w=w delz=delz pt=pt tracers=q cappa=cappa + !$ser data peln=peln te0_2d=te_2d ua=ua va=va omga=omga dp1=dp1 dtdt=dtdt_m !$ser data mfx=mfxL !$ser data mfy=mfyL !$ser data cx=cxL !$ser data cy=cyL + !$ser data delp=delp + !$ser data pk=pk + !$ser data qvapor=q(:,:,:,sphum) + !$ser data qliquid=q(:,:,:,2) + !$ser data qice=q(:,:,:,3) + !$ser data qrain=q(:,:,:,4) + !$ser data qsnow=q(:,:,:,5) + !$ser data qgraupel=q(:,:,:,6) + !$ser data qcld=q(:,:,:,7) + !$ser data qo3mr=q(:,:,:,8) + !$ser data qsgs_tke=q(:,:,:,9) #ifdef AVEC_TIMERS call avec_timer_stop(6) diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 1e11f4af..55412ce6 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -92,6 +92,7 @@ from .translate_ray_fast import TranslateRay_Fast from .translate_remapping import TranslateRemapping from .translate_remapping_GEOS import TranslateRemapping_GEOS +from .translate_remapping_GEOS_v2 import TranslateRemapping_GEOS_v2 from .translate_riem_solver3 import TranslateRiem_Solver3 from .translate_riem_solver_c import TranslateRiem_Solver_C from .translate_satadjust3d import TranslateSatAdjust3d diff --git a/tests/savepoint/translate/translate_remapping_GEOS_v2.py b/tests/savepoint/translate/translate_remapping_GEOS_v2.py index 2148bb46..024ac16b 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS_v2.py +++ b/tests/savepoint/translate/translate_remapping_GEOS_v2.py @@ -9,9 +9,12 @@ Z_DIM, Z_INTERFACE_DIM, ) +from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig -from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS + + +# from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS # from pyFV3._config import RemappingConfig @@ -19,8 +22,8 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): inputs = { - "pe_": { - "name": "pe_", + "pe": { + "name": "pe", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, @@ -99,8 +102,8 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "peln_3d": { - "name": "peln_3d", + "peln": { + "name": "peln", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, @@ -116,7 +119,7 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, "pkz": { @@ -129,11 +132,6 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "ws_": { - "name": "ws_", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, "u": { "name": "u", "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], @@ -154,13 +152,13 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "mfx_": { - "name": "mfx_", + "mfx": { + "name": "mfx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "cx_": { - "name": "cx_", + "cx": { + "name": "cx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, @@ -169,11 +167,21 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "te0_2d_": { - "name": "te0_2d_", + "te0_2d": { + "name": "te0_2d", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "wsd": { + "name": "wsd", "dims": [X_DIM, Y_DIM], "units": "No Units", }, + "dp1": { + "name": "dp1", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, } outputs = { "pt": { @@ -266,29 +274,29 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "mfx_": { - "name": "mfx_", + "mfx": { + "name": "mfx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "cx_": { - "name": "cx_", + "cx": { + "name": "cx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "peln_3d": { - "name": "peln_3d", + "peln": { + "name": "peln", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, - "pe_": { - "name": "pe_", + "pe": { + "name": "pe", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, "pkz": { @@ -301,6 +309,16 @@ class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, + "dp1": { + "name": "dp1", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "ps": { + "name": "ps", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, } def __init__( @@ -312,7 +330,34 @@ def __init__( super().__init__(grid, namelist, stencil_factory) self._base.in_vars["data_vars"] = { - "tracers": {}, + # "tracers": {}, + "qvapor": { + "kend": grid.npz - 1, + }, + "qliquid": { + "kend": grid.npz - 1, + }, + "qice": { + "kend": grid.npz - 1, + }, + "qrain": { + "kend": grid.npz - 1, + }, + "qsnow": { + "kend": grid.npz - 1, + }, + "qgraupel": { + "kend": grid.npz - 1, + }, + "qcld": { + "kend": grid.npz - 1, + }, + "qo3mr": { + "kend": grid.npz - 1, + }, + "qsgs_tke": { + "kend": grid.npz - 1, + }, "w": { "kend": grid.npz - 1, }, @@ -330,7 +375,7 @@ def __init__( "iend": grid.ie, "jstart": grid.js, "jend": grid.je, - "kend": grid.npz, + "kend": grid.npz + 1, }, "peln": { "istart": grid.is_, @@ -348,7 +393,7 @@ def __init__( "kend": grid.npz + 1, "kaxis": 1, }, - "hs": {"serialname": "phis"}, + # "hs": {"serialname": "phis"}, "ps": {}, "wsd": { "istart": grid.is_, @@ -356,6 +401,46 @@ def __init__( "jstart": grid.js, "jend": grid.je, }, + "mfy": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je + 1, + "kend": grid.npz - 1, + }, + "cy": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.js, + "jend": grid.je + 1, + "kend": grid.npz - 1, + }, + "mfx": { + "istart": grid.is_, + "iend": grid.ie + 1, + "jstart": grid.js, + "jend": grid.je, + "kend": grid.npz - 1, + }, + "cx": { + "istart": grid.is_, + "iend": grid.ie + 1, + "jstart": grid.jsd, + "jend": grid.jed, + "kend": grid.npz - 1, + }, + "hs": { + "istart": grid.isd, + "iend": grid.ied, + "jstart": grid.jsd, + "jend": grid.jed, + }, + "te0_2d": { + "istart": grid.is_, + "iend": grid.ie, + "jstart": grid.js, + "jend": grid.je, + }, # column variables... "ak": {}, "bk": {}, @@ -377,7 +462,16 @@ def __init__( "pk", "peln", "pt", - "tracers", + # "tracers", + "qvapor", + "qliquid", + "qice", + "qrain", + "qsnow", + "qgraupel", + "qcld", + "qo3mr", + "qsgs_tke", "cappa", "delp", "delz", @@ -387,6 +481,10 @@ def __init__( "w", "ps", "dp1", + "mfy", + "cy", + "mfx", + "cx", ]: self._base.out_vars[k] = self._base.in_vars["data_vars"][k] @@ -401,64 +499,73 @@ def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") def compute_parallel(self, inputs, communicator): + inputs["te0_2d"] = inputs["te0_2d"].astype(Float) state = self.state_from_inputs(inputs) state_namespace = SimpleNamespace(**state) - tracers = { - "qvapor": state_namespace.qvapor, - "qliquid": state_namespace.qliquid, - "qice": state_namespace.qice, - "qrain": state_namespace.qrain, - "qsnow": state_namespace.qsnow, - "qgraupel": state_namespace.qgraupel, - "qcld": state_namespace.qcld, - "qo3mr": state_namespace.qo3mr, - "qsgs_tke": state_namespace.qsgs_tke, - } + # tracers = { + # "qvapor": state_namespace.qvapor, + # "qliquid": state_namespace.qliquid, + # "qice": state_namespace.qice, + # "qrain": state_namespace.qrain, + # "qsnow": state_namespace.qsnow, + # "qgraupel": state_namespace.qgraupel, + # "qcld": state_namespace.qcld, + # "qo3mr": state_namespace.qo3mr, + # "qsgs_tke": state_namespace.qsgs_tke, + # } + + # l_to_e = LagrangianToEulerian_GEOS( + # self.stencil_factory, + # self.quantity_factory, + # DynamicalCoreConfig.from_namelist(self.namelist).remapping, + # communicator, + # self.grid.grid_data, + # state_namespace.nq, + # state_namespace.pfull, + # tracers, + # state_namespace.adiabatic, + # ) - l_to_e = LagrangianToEulerian_GEOS( - self.stencil_factory, - self.quantity_factory, - DynamicalCoreConfig.from_namelist(self.namelist).remapping, - communicator, - self.grid.grid_data, - state_namespace.nq, - state_namespace.pfull, - tracers, - state_namespace.adiabatic, - ) + # print("type(state_namespace.ptop): ", type(state_namespace.ptop)) + # print("type(state_namespace.akap): ", type(state_namespace.akap)) + # print("type(state_namespace.zvir): ", type(state_namespace.zvir)) + # print("type(state_namespace.last_step): ", type(state_namespace.last_step)) + # print("type(state_namespace.consv_te): ", type(state_namespace.consv_te)) + # print("type(state_namespace.mdt): ", type(state_namespace.mdt)) + # print("type(nq): ", type(state_namespace.nq)) - l_to_e( - tracers, - state_namespace.pt, - state_namespace.delp, - state_namespace.delz, - state_namespace.peln, - state_namespace.u, - state_namespace.v, - state_namespace.w, - state_namespace.mfx, - state_namespace.mfy, - state_namespace.cx, - state_namespace.cy, - state_namespace.cappa, - state_namespace.q_con, - state_namespace.pkz, - state_namespace.pk, - state_namespace.pe, - state_namespace.hs, - state_namespace.te0_2d, - state_namespace.ps, - state_namespace.wsd, - state_namespace.ak, - state_namespace.bk, - state_namespace.dp1, - state_namespace.ptop, - state_namespace.akap, - state_namespace.zvir, - state_namespace.last_step, - state_namespace.consv_te, - state_namespace.mdt, - ) + # l_to_e( + # tracers, + # state_namespace.pt, + # state_namespace.delp, + # state_namespace.delz, + # state_namespace.peln, + # state_namespace.u, + # state_namespace.v, + # state_namespace.w, + # state_namespace.mfx, + # state_namespace.mfy, + # state_namespace.cx, + # state_namespace.cy, + # state_namespace.cappa, + # state_namespace.q_con, + # state_namespace.pkz, + # state_namespace.pk, + # state_namespace.pe, + # state_namespace.hs, + # state_namespace.te0_2d, + # state_namespace.ps, + # state_namespace.wsd, + # state_namespace.ak, + # state_namespace.bk, + # state_namespace.dp1, + # Float(state_namespace.ptop), + # Float(state_namespace.akap), + # state_namespace.zvir, + # state_namespace.last_step, + # state_namespace.consv_te, + # state_namespace.mdt, + # ) return self.outputs_from_state(state) From 27c67b2564b6511b4a8e04b4bf8841c98c7009c1 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 5 Feb 2025 11:01:39 -0800 Subject: [PATCH 163/252] Continue L_to_E integration in translate remapping test --- pyFV3/stencils/mapn_tracer.py | 12 +- pyFV3/stencils/remapping_GEOS.py | 3 +- .../translate/translate_remapping_GEOS_v2.py | 120 ++++++++---------- 3 files changed, 65 insertions(+), 70 deletions(-) diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index 0696d145..050e5610 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -1,7 +1,9 @@ -from typing import Dict +# from typing import Dict import ndsl.dsl.gt4py_utils as utils -from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate + +# from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.fillz import FillNegativeTracerValues @@ -20,7 +22,8 @@ def __init__( kord: int, nq: int, fill: bool, - tracers: Dict[str, Quantity], + # tracers: Dict[str, Quantity], + tracers, ): orchestrate( obj=self, @@ -64,7 +67,8 @@ def __call__( pe1: FloatField, pe2: FloatField, dp2: FloatField, - tracers: Dict[str, Quantity], + # tracers: Dict[str, Quantity], + tracers, ): """ Remaps the tracer species onto the Eulerian grid diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 106ff7c9..6fca74c6 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -85,7 +85,8 @@ def __init__( self._adiabatic = adiabatic self.kmp = grid_indexing.domain[2] - 1 for k in range(pfull.shape[0]): - if pfull.view[k] > 10.0e2: + # if pfull.view[k] > 10.0e2: + if pfull[k] > 10.0e2: self.kmp = k break # do_omega = hydrostatic and last_step # TODO pull into inputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS_v2.py b/tests/savepoint/translate/translate_remapping_GEOS_v2.py index 024ac16b..56df0259 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS_v2.py +++ b/tests/savepoint/translate/translate_remapping_GEOS_v2.py @@ -12,9 +12,7 @@ from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig - - -# from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS # from pyFV3._config import RemappingConfig @@ -503,69 +501,61 @@ def compute_parallel(self, inputs, communicator): state = self.state_from_inputs(inputs) state_namespace = SimpleNamespace(**state) - # tracers = { - # "qvapor": state_namespace.qvapor, - # "qliquid": state_namespace.qliquid, - # "qice": state_namespace.qice, - # "qrain": state_namespace.qrain, - # "qsnow": state_namespace.qsnow, - # "qgraupel": state_namespace.qgraupel, - # "qcld": state_namespace.qcld, - # "qo3mr": state_namespace.qo3mr, - # "qsgs_tke": state_namespace.qsgs_tke, - # } - - # l_to_e = LagrangianToEulerian_GEOS( - # self.stencil_factory, - # self.quantity_factory, - # DynamicalCoreConfig.from_namelist(self.namelist).remapping, - # communicator, - # self.grid.grid_data, - # state_namespace.nq, - # state_namespace.pfull, - # tracers, - # state_namespace.adiabatic, - # ) + tracers = { + "vapor": state_namespace.qvapor, + "liquid": state_namespace.qliquid, + "ice": state_namespace.qice, + "rain": state_namespace.qrain, + "snow": state_namespace.qsnow, + "graupel": state_namespace.qgraupel, + "cloud": state_namespace.qcld, + "qo3mr": state_namespace.qo3mr, + "qsgs_tke": state_namespace.qsgs_tke, + } - # print("type(state_namespace.ptop): ", type(state_namespace.ptop)) - # print("type(state_namespace.akap): ", type(state_namespace.akap)) - # print("type(state_namespace.zvir): ", type(state_namespace.zvir)) - # print("type(state_namespace.last_step): ", type(state_namespace.last_step)) - # print("type(state_namespace.consv_te): ", type(state_namespace.consv_te)) - # print("type(state_namespace.mdt): ", type(state_namespace.mdt)) - # print("type(nq): ", type(state_namespace.nq)) + l_to_e = LagrangianToEulerian_GEOS( + self.stencil_factory, + self.quantity_factory, + DynamicalCoreConfig.from_namelist(self.namelist).remapping, + communicator, + self.grid.grid_data, + state_namespace.nq, + state_namespace.pfull, + tracers, + DynamicalCoreConfig.adiabatic, + ) - # l_to_e( - # tracers, - # state_namespace.pt, - # state_namespace.delp, - # state_namespace.delz, - # state_namespace.peln, - # state_namespace.u, - # state_namespace.v, - # state_namespace.w, - # state_namespace.mfx, - # state_namespace.mfy, - # state_namespace.cx, - # state_namespace.cy, - # state_namespace.cappa, - # state_namespace.q_con, - # state_namespace.pkz, - # state_namespace.pk, - # state_namespace.pe, - # state_namespace.hs, - # state_namespace.te0_2d, - # state_namespace.ps, - # state_namespace.wsd, - # state_namespace.ak, - # state_namespace.bk, - # state_namespace.dp1, - # Float(state_namespace.ptop), - # Float(state_namespace.akap), - # state_namespace.zvir, - # state_namespace.last_step, - # state_namespace.consv_te, - # state_namespace.mdt, - # ) + l_to_e( + tracers, + state_namespace.pt, + state_namespace.delp, + state_namespace.delz, + state_namespace.peln, + state_namespace.u, + state_namespace.v, + state_namespace.w, + state_namespace.mfx, + state_namespace.mfy, + state_namespace.cx, + state_namespace.cy, + state_namespace.cappa, + state_namespace.q_con, + state_namespace.pkz, + state_namespace.pk, + state_namespace.pe, + state_namespace.hs, + state_namespace.te0_2d, + state_namespace.ps, + state_namespace.wsd, + state_namespace.ak, + state_namespace.bk, + state_namespace.dp1, + Float(state_namespace.ptop), + Float(state_namespace.akap), + state_namespace.zvir, + state_namespace.last_step, + state_namespace.consv_te, + state_namespace.mdt, + ) return self.outputs_from_state(state) From 4817f621b5cbbc00d8e7d41ae9a6f838e950f31c Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Fri, 7 Feb 2025 07:47:37 -0800 Subject: [PATCH 164/252] Re-incorporated Tracer class into GEOS Remapping translate test --- pyFV3/stencils/fillz.py | 8 +-- pyFV3/stencils/mapn_tracer.py | 10 +-- pyFV3/stencils/remapping.py | 6 +- pyFV3/stencils/remapping_GEOS.py | 72 +++++++++---------- .../translate/translate_remapping_GEOS_v2.py | 46 +++++++++--- 5 files changed, 83 insertions(+), 59 deletions(-) diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index 5cd5c239..bc8ef7c7 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -1,12 +1,12 @@ import typing -from typing import Dict from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval import ndsl.dsl.gt4py_utils as utils -from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntFieldIJ +from pyFV3.tracers import Tracers @typing.no_type_check @@ -118,7 +118,7 @@ def __init__( stencil_factory: StencilFactory, quantity_factory: QuantityFactory, nq: int, - tracers: Dict[str, Quantity], + tracers: Tracers, ): orchestrate( obj=self, @@ -152,7 +152,7 @@ def __init__( def __call__( self, dp2: FloatField, - tracers: Dict[str, Quantity], + tracers: Tracers, ): """ Args: diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index 050e5610..79e598c9 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -1,13 +1,10 @@ -# from typing import Dict - import ndsl.dsl.gt4py_utils as utils - -# from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.stencils.map_single import MapSingle +from pyFV3.tracers import Tracers class MapNTracer: @@ -22,8 +19,7 @@ def __init__( kord: int, nq: int, fill: bool, - # tracers: Dict[str, Quantity], - tracers, + tracers: Tracers, ): orchestrate( obj=self, @@ -67,7 +63,7 @@ def __call__( pe1: FloatField, pe2: FloatField, dp2: FloatField, - # tracers: Dict[str, Quantity], + # tracers: Tracers, tracers, ): """ diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 8e5a1dd6..4174493b 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -32,6 +32,9 @@ from pyFV3.stencils.saturation_adjustment import SatAdjust3d +# from pyFV3.tracers import Tracers + + # TODO: Should this be set here or in global_constants? CONSV_MIN = 0.001 @@ -361,7 +364,8 @@ def __init__( area_64, nq, pfull, - tracers: Dict[str, Quantity], + # tracers: Tracers, + tracers, checkpointer: Optional[Checkpointer] = None, ): orchestrate( diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 6fca74c6..3d6c3629 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -413,12 +413,12 @@ def __call__( # Build remapping profiles self._init_pe(pe, self._pe1, self._pe2, ptop) self._moist_cv_pt_pressure( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers["qvapor"], + qliquid=tracers["qliquid"], + qrain=tracers["qrain"], + qsnow=tracers["qsnow"], + qice=tracers["qice"], + qgraupel=tracers["qgraupel"], q_con=q_con, pt=pt, cappa=cappa, @@ -499,12 +499,12 @@ def __call__( ) self._moist_cv_pkz( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers["qvapor"], + qliquid=tracers["qliquid"], + qrain=tracers["qrain"], + qsnow=tracers["qsnow"], + qice=tracers["qice"], + qgraupel=tracers["qgraupel"], pkz=pkz, pt=pt, cappa=cappa, @@ -517,12 +517,12 @@ def __call__( if last_step: if consv_te > CONSV_MIN: self._moist_cv_te( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers["qvapor"], + qliquid=tracers["qliquid"], + qrain=tracers["qrain"], + qsnow=tracers["qsnow"], + qice=tracers["qice"], + qgraupel=tracers["qgraupel"], u=u, v=v, w=w, @@ -570,13 +570,13 @@ def __call__( fast_mp_consv = consv_te > CONSV_MIN self._saturation_adjustment( dp1, - tracers["vapor"], - tracers["liquid"], - tracers["ice"], - tracers["rain"], - tracers["snow"], - tracers["graupel"], - tracers["cloud"], + tracers["qvapor"], + tracers["qliquid"], + tracers["qice"], + tracers["qrain"], + tracers["qsnow"], + tracers["qgraupel"], + tracers["qcld"], hs, peln, delp, @@ -598,12 +598,12 @@ def __call__( # to the physics, but if we're staying in dynamics we need # to keep it as the virtual potential temperature self._moist_cv_last_step_stencil( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers["qvapor"], + qliquid=tracers["qliquid"], + qrain=tracers["qrain"], + qsnow=tracers["qsnow"], + qice=tracers["qice"], + qgraupel=tracers["qgraupel"], pt=pt, pkz=pkz, dtmp=Float(dtmp), @@ -611,11 +611,11 @@ def __call__( ) self._fill_cond( q_con=q_con, - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qliquid=tracers["qliquid"], + qrain=tracers["qrain"], + qsnow=tracers["qsnow"], + qice=tracers["qice"], + qgraupel=tracers["qgraupel"], ) else: # converts virtual temperature back to virtual potential temperature diff --git a/tests/savepoint/translate/translate_remapping_GEOS_v2.py b/tests/savepoint/translate/translate_remapping_GEOS_v2.py index 56df0259..d1b236fc 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS_v2.py +++ b/tests/savepoint/translate/translate_remapping_GEOS_v2.py @@ -13,6 +13,7 @@ from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyFV3.tracers import Tracers # from pyFV3._config import RemappingConfig @@ -497,21 +498,34 @@ def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") def compute_parallel(self, inputs, communicator): + + tracers = Tracers(self.quantity_factory) + + tracers.copy_tracer_data("qvapor", inputs["qvapor"]) + tracers.copy_tracer_data("qliquid", inputs["qliquid"]) + tracers.copy_tracer_data("qice", inputs["qice"]) + tracers.copy_tracer_data("qrain", inputs["qrain"]) + tracers.copy_tracer_data("qsnow", inputs["qsnow"]) + tracers.copy_tracer_data("qgraupel", inputs["qgraupel"]) + tracers.copy_tracer_data("qcld", inputs["qcld"]) + tracers.copy_tracer_data("qo3mr", inputs["qo3mr"]) + tracers.copy_tracer_data("qsgs_tke", inputs["qsgs_tke"]) + inputs["te0_2d"] = inputs["te0_2d"].astype(Float) state = self.state_from_inputs(inputs) state_namespace = SimpleNamespace(**state) - tracers = { - "vapor": state_namespace.qvapor, - "liquid": state_namespace.qliquid, - "ice": state_namespace.qice, - "rain": state_namespace.qrain, - "snow": state_namespace.qsnow, - "graupel": state_namespace.qgraupel, - "cloud": state_namespace.qcld, - "qo3mr": state_namespace.qo3mr, - "qsgs_tke": state_namespace.qsgs_tke, - } + # tracers = { + # "qvapor": state_namespace.qvapor, + # "qliquid": state_namespace.qliquid, + # "qice": state_namespace.qice, + # "qrain": state_namespace.qrain, + # "qsnow": state_namespace.qsnow, + # "qgraupel": state_namespace.qgraupel, + # "qcld": state_namespace.qcld, + # "qo3mr": state_namespace.qo3mr, + # "qsgs_tke": state_namespace.qsgs_tke, + # } l_to_e = LagrangianToEulerian_GEOS( self.stencil_factory, @@ -558,4 +572,14 @@ def compute_parallel(self, inputs, communicator): state_namespace.mdt, ) + state_namespace.qvapor.data = tracers["qvapor"].data + state_namespace.qliquid.data = tracers["qliquid"].data + state_namespace.qice.data = tracers["qice"].data + state_namespace.qrain.data = tracers["qrain"].data + state_namespace.qsnow.data = tracers["qsnow"].data + state_namespace.qgraupel.data = tracers["qgraupel"].data + state_namespace.qcld.data = tracers["qcld"].data + state_namespace.qo3mr.data = tracers["qo3mr"].data + state_namespace.qsgs_tke.data = tracers["qsgs_tke"].data + return self.outputs_from_state(state) From 9d4c9fbd1836c660a8edabf2a1bf1b9afc84e65c Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 12 Feb 2025 04:39:02 -0800 Subject: [PATCH 165/252] updated fv_dynamics serialization file --- fv_dynamics.F90.SER | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/fv_dynamics.F90.SER b/fv_dynamics.F90.SER index 1f8cae76..9b65f06b 100644 --- a/fv_dynamics.F90.SER +++ b/fv_dynamics.F90.SER @@ -796,12 +796,11 @@ contains !$ser data iq=iq last_step=last_step consv_te=consv_te ps=ps pe=pe delp=delp pkz=pkz pk=pk mdt=mdt bdt=bdt sphum=sphum q_con=q_con u=u v=v w=w delz=delz pt=pt tracers=q phis=phis zvir=zvir akap=akap cappa=cappa kord_tracer=kord_tracer peln=peln te_2d=te_2d ua=ua va=va omga=omga dp1=dp1 wsd=ws reproduce_sum=reproduce_sum ptop=ptop ak=ak bk=bk pfull=pfull hybrid_z=hybrid_z do_adiabatic_init=do_adiabatic_init nq=nq !$ser savepoint Remapping_GEOS_v2-In !$ser data last_step=last_step consv_te=consv_te ps=ps pe=pe delp=delp pkz=pkz mdt=mdt pdt=bdt - !$ser data km=npz nq=nq nwat=nwat sphum=sphum q_con=q_con u=u v=v w=w delz=delz pt=pt tracers=q - !$ser data hs=phis zvir=zvir cp=cp_air akap=akap cappa=cappa kord_tr=kord_tracer - !$ser data peln=peln te0_2d=te_2d ng=ng ua=ua va=va omga=omga dp1=dp1 wsd=ws fill=fill - !$ser data reproduce_sum=reproduce_sum dtdt=dtdt_m ptop=ptop ak=ak bk=bk pfull=pfull - !$ser data hydrostatic=hydrostatic hybrid_z=hybrid_z do_omega=do_omega - !$ser data do_adiabatic_init=do_adiabatic_init + !$ser data q_con=q_con nq=nq u=u v=v w=w delz=delz pt=pt tracers=q + !$ser data hs=phis zvir=zvir + !$ser data akap=akap cappa=cappa + !$ser data peln=peln te0_2d=te_2d dp1=dp1 wsd=ws + !$ser data ptop=ptop ak=ak bk=bk pfull=pfull !$ser data mfx=mfxL !$ser data mfy=mfyL !$ser data cx=cxL @@ -831,7 +830,8 @@ contains !$ser data te_2d=te_2d pk=pk tracers=q delp=delp pe=pe ps=ps u=u v=v w=w pt=pt delz=delz q_con=q_con cappa=cappa ua=ua va=va omga=omga peln=peln pkz=pkz dp1=dp1 !$ser savepoint Remapping_GEOS_v2-Out !$ser data ps=ps pe=pe pkz=pkz q_con=q_con u=u v=v w=w delz=delz pt=pt tracers=q cappa=cappa - !$ser data peln=peln te0_2d=te_2d ua=ua va=va omga=omga dp1=dp1 dtdt=dtdt_m + !$ser data peln=peln + !$ser data dp1=dp1 !$ser data mfx=mfxL !$ser data mfy=mfyL !$ser data cx=cxL From a4f89d9b34436c536db4be56cc3b52fa3993f31a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 12 Feb 2025 12:08:37 -0500 Subject: [PATCH 166/252] Fix Rayleigh Friction trigger configuration (raise) --- pyFV3/stencils/fv_dynamics.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index d04d7dfa..cb969a3b 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -587,7 +587,8 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool): # Rayleigh fast if ( - not self.config.acoustic_dynamics.rf_fast + not self.config.hydrostatic + and not self.config.acoustic_dynamics.rf_fast and self.config.acoustic_dynamics.tau > 0 ): raise NotImplementedError( From de77b0ee12cc8eb10cb66d7488d9ed67cc904ad6 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Thu, 27 Feb 2025 09:58:23 -0800 Subject: [PATCH 167/252] Added temporary calculation to flux calculation in fxadv.py that will reduce error rates of dpx values in D_SW --- pyFV3/stencils/fxadv.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/pyFV3/stencils/fxadv.py b/pyFV3/stencils/fxadv.py index c1bf805d..511cdf00 100644 --- a/pyFV3/stencils/fxadv.py +++ b/pyFV3/stencils/fxadv.py @@ -489,19 +489,24 @@ def fxadv_fluxes_stencil( with computation(PARALLEL), interval(...): with horizontal(region[local_is : local_ie + 2, :]): + # Including the temporary (tmp) calculation enables x_area_flux and y_area_flux + # to more closely precision match the respective Fortran calculation + # since Fortran also performs this temporary calcuation + tmp = dt * uc_contra if uc_contra > 0: - crx = dt * uc_contra * rdxa[-1, 0] - x_area_flux = dy * dt * uc_contra * sin_sg3[-1, 0] + crx = tmp * rdxa[-1, 0] + x_area_flux = dy * tmp * sin_sg3[-1, 0] else: - crx = dt * uc_contra * rdxa - x_area_flux = dy * dt * uc_contra * sin_sg1 + crx = tmp * rdxa + x_area_flux = dy * tmp * sin_sg1 with horizontal(region[:, local_js : local_je + 2]): + tmp = dt * vc_contra if vc_contra > 0: - cry = dt * vc_contra * rdya[0, -1] - y_area_flux = dx * dt * vc_contra * sin_sg4[0, -1] + cry = tmp * rdya[0, -1] + y_area_flux = dx * tmp * sin_sg4[0, -1] else: - cry = dt * vc_contra * rdya - y_area_flux = dx * dt * vc_contra * sin_sg2 + cry = tmp * rdya + y_area_flux = dx * tmp * sin_sg2 class FiniteVolumeFluxPrep: From 03206b088cf27b3bab20f6ce3cef0af8f9423149 Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Tue, 4 Mar 2025 13:40:44 -0800 Subject: [PATCH 168/252] Added fix to improve result matching of 'delpc' with GEOS codebase --- pyFV3/stencils/divergence_damping.py | 37 +++++++++++++++++----------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/pyFV3/stencils/divergence_damping.py b/pyFV3/stencils/divergence_damping.py index 46a949ce..be94a6ac 100644 --- a/pyFV3/stencils/divergence_damping.py +++ b/pyFV3/stencils/divergence_damping.py @@ -1,4 +1,5 @@ import gt4py.cartesian.gtscript as gtscript +import numpy as np from gt4py.cartesian.gtscript import ( __INLINED, PARALLEL, @@ -27,6 +28,12 @@ def damp_tmp(q, da_min_c, d2_bg, dddmp): damp = da_min_c * max(d2_bg, mintmp) return damp +@gtscript.function +def damp_tmp2(q, da_min_c, d2_bg, dddmp): + mintmp = min(0.2, dddmp * q) + damp = da_min_c * max(d2_bg, mintmp) + return damp + def compute_u_contra_dyc( u: FloatField, @@ -58,10 +65,10 @@ def compute_u_contra_dyc( # TODO: why does vc_from_va sometimes have different sign than vc? vc_from_va = 0.5 * (va[0, -1, 0] + va) # TODO: why do we use vc_from_va and not just vc? - u_contra = contravariant(u, vc_from_va, cosa_v, sina_v) + u_contra_dyc = contravariant(u, vc_from_va, cosa_v, dyc) + u_contra_dyc = u_contra_dyc * sina_v with horizontal(region[:, j_start], region[:, j_end + 1]): - u_contra = u * sin_sg4[0, -1] if vc > 0 else u * sin_sg2 - u_contra_dyc = u_contra * dyc + u_contra_dyc = u * dyc* sin_sg4[0, -1] if vc > 0 else u * dyc*sin_sg2 def compute_v_contra_dxc( @@ -93,10 +100,10 @@ def compute_v_contra_dxc( # TODO: why does uc_from_ua sometimes have different sign than uc? uc_from_ua = 0.5 * (ua[-1, 0, 0] + ua) # TODO: why do we use uc_from_ua and not just uc? - v_contra = contravariant(v, uc_from_ua, cosa_u, sina_u) + v_contra_dxc = contravariant(v, uc_from_ua, cosa_u, dxc) + v_contra_dxc = v_contra_dxc * sina_u with horizontal(region[i_start, :], region[i_end + 1, :]): - v_contra = v * sin_sg3[-1, 0] if uc > 0 else v * sin_sg1 - v_contra_dxc = v_contra * dxc + v_contra_dxc = v * dxc*sin_sg3[-1, 0] if uc > 0 else v * dxc*sin_sg1 def delpc_computation( @@ -324,11 +331,11 @@ def __init__( if nested: raise NotImplementedError("Divergence Dampoing: nested not implemented.") # TODO: make dddmp a compile-time external, instead of runtime scalar - self._dddmp = dddmp + self._dddmp = Float(dddmp) # TODO: make da_min_c a compile-time external, instead of runtime scalar self._damping_coefficients = damping_coefficients self._stretched_grid = stretched_grid - self._d4_bg = d4_bg + self._d4_bg = Float(d4_bg) self._grid_type = grid_type self._nord_column = nord_col self._d2_bg_column = d2_bg @@ -539,12 +546,12 @@ def __init__( # odd and adds a lot of boilerplate throughout the model code. @dace_inhibitor - def _get_da_min_c(self) -> float: - return self._damping_coefficients.da_min_c + def _get_da_min_c(self) -> Float: + return Float(self._damping_coefficients.da_min_c) @dace_inhibitor - def _get_da_min(self) -> float: - return self._damping_coefficients.da_min + def _get_da_min(self) -> Float: + return Float(self._damping_coefficients.da_min) def __call__( self, @@ -697,9 +704,11 @@ def __call__( da_min: Float = self._get_da_min() if self._stretched_grid: # reference https://github.com/NOAA-GFDL/GFDL_atmos_cubed_sphere/blob/main/model/sw_core.F90#L1422 # noqa: E501 - dd8 = da_min * self._d4_bg ** (self._nonzero_nord + 1) + dd8 = da_min * np.power(self._d4_bg, (self._nonzero_nord + 1), dtype=Float) else: - dd8 = (da_min_c * self._d4_bg) ** (self._nonzero_nord + 1) + dd8 = np.power( + (da_min_c * self._d4_bg), (self._nonzero_nord + 1), dtype=Float + ) self._damping_nord_highorder_stencil( damped_rel_vort_bgrid, From 639b1d0792ae5edc2d56b37e522fe28930376b8c Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 5 Mar 2025 08:43:22 -0800 Subject: [PATCH 169/252] Added notes on porting --- pyFV3/stencils/divergence_damping.py | 34 ++++++++++++++++++++++++++++ pyFV3/stencils/fxadv.py | 15 ++++++++++++ 2 files changed, 49 insertions(+) diff --git a/pyFV3/stencils/divergence_damping.py b/pyFV3/stencils/divergence_damping.py index be94a6ac..30c47a59 100644 --- a/pyFV3/stencils/divergence_damping.py +++ b/pyFV3/stencils/divergence_damping.py @@ -58,6 +58,20 @@ def compute_u_contra_dyc( sin_sg2 (in): sin_sg4 (in): u_contra_dyc (out): contravariant u-wind on d-grid + + Porting Notes + * The compute_u_contra_dyc and compute_v_contra_dxc functions have the dyc and dxc values + incorporated earlier in the calcuation rather than later, and this enables the u_contra_dyc + and v_contra_dxc values to match with the Fortran. As a result, the delpc computation + matches the Fortran value of delpc. + + Ex : Previous implementation of compute_u_contra_dyc + ================================================================= + u_contra = contravariant(u, vc_from_va, cosa_v, sina_v) + with horizontal(region[:, j_start], region[:, j_end + 1]): + u_contra = u * sin_sg4[0, -1] if vc > 0 else u * sin_sg2 + u_contra_dyc = u_contra * dyc + ================================================================= """ from __externals__ import j_end, j_start @@ -93,6 +107,19 @@ def compute_v_contra_dxc( uc (in): sin_sg3 (in): sin_sg1 (in): + + Porting Notes + * The compute_u_contra_dyc and compute_v_contra_dxc functions have the dyc and dxc values + incorporated earlier in the calcuation rather than later, and this enables the u_contra_dyc + and v_contra_dxc values to match with the Fortran. As a result, the delpc computation + matches the Fortran value of delpc. + + Ex : Previous implementation of compute_v_contra_dxc + ================================================================= + v_contra = contravariant(v, uc_from_ua, cosa_u, sina_u) + with horizontal(region[i_start, :], region[i_end + 1, :]): + v_contra = v * sin_sg3[-1, 0] if uc > 0 else v * sin_sg1 + v_contra_dxc = v_contra * dxc """ from __externals__ import i_end, i_start @@ -580,6 +607,13 @@ def __call__( Applies both a background second-order diffusion (with strength controlled by d2_bg passed on init) and a higher-order hyperdiffusion. + Porting Notes + * The dd8 computation has different results when comparing between Fortran and Python, + which is likely due to the user of the power function. The difference in dd8 results in + the ke value having on the order of 10,000 total error difference when running the translate + test. + + Args: u (in): x-velocity on d-grid v (in): y-velocity on d-grid diff --git a/pyFV3/stencils/fxadv.py b/pyFV3/stencils/fxadv.py index 511cdf00..626a6bfc 100644 --- a/pyFV3/stencils/fxadv.py +++ b/pyFV3/stencils/fxadv.py @@ -484,6 +484,21 @@ def fxadv_fluxes_stencil( y_area_flux (out): uc_contra (in): vc_contra (in): + + Porting Note + * The tmp introduced in the computation allows fxadv_fluxes_stencil to closely match the Fortran order + of computation, which allows the x_area_flux and y_area_flux match the + respective Fortran values. + + Example of previous stencil looked as follows: + ========================================================== + if uc_contra > 0: + crx = dt * uc_contra * rdxa[-1, 0] + x_area_flux = dy * dt * uc_contra * sin_sg3[-1, 0] + else: + crx = dt * uc_contra * rdxa + x_area_flux = dy * dt * uc_contra * sin_sg1 + ========================================================== """ from __externals__ import local_ie, local_is, local_je, local_js From e4d65df40f45fb714b923785b6754edf774ebc8b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 5 Mar 2025 15:54:52 -0500 Subject: [PATCH 170/252] Move damping increment to own K stencils Trigger damping increment calculation exactly once --- pyFV3/stencils/ray_fast.py | 81 +++++++++++++++++++++++++++++--------- 1 file changed, 63 insertions(+), 18 deletions(-) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index 3a65c417..a28e8946 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -15,7 +15,15 @@ import ndsl.constants as constants from ndsl import StencilFactory, orchestrate -from ndsl.constants import SECONDS_PER_DAY, X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.boilerplate import get_factories_single_tile +from ndsl.constants import ( + SECONDS_PER_DAY, + X_INTERFACE_DIM, + Y_INTERFACE_DIM, + Z_DIM, + X_DIM, + Y_DIM, +) from ndsl.dsl.typing import Float, FloatField, FloatFieldK @@ -43,14 +51,31 @@ def dm_layer(rf, dp, wind): return (1.0 - rf) * dp * wind +def ray_fast_damping_increment( + pfull: FloatFieldK, # type:ignore + dt: Float, # type:ignore + ptop: Float, # type:ignore + rf: FloatField, # type:ignore +): + """rf is rayleigh damping increment, fraction of vertical velocity + left after doing rayleigh damping (w -> w * rf) + """ + from __externals__ import rf_cutoff, tau + + with computation(PARALLEL), interval(...): + if pfull < rf_cutoff: + # rf is rayleigh damping increment, fraction of vertical velocity + # left after doing rayleigh damping (w -> w * rf) + rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SECONDS_PER_DAY, ptop) + + def ray_fast_wind_compute( u: FloatField, v: FloatField, w: FloatField, delta_p_ref: FloatFieldK, # reference delta pressure pfull: FloatFieldK, # input layer pressure reference? - dt: Float, - ptop: Float, + rf: FloatFieldK, rf_cutoff_nudge: Float, ): """ @@ -68,13 +93,6 @@ def ray_fast_wind_compute( from __externals__ import hydrostatic, local_ie, local_je, rf_cutoff, tau # dm_stencil - with computation(PARALLEL), interval(...): - # TODO -- in the fortran model rf is only computed once, repeating - # the computation every time ray_fast is run is inefficient - if pfull < rf_cutoff: - # rf is rayleigh damping increment, fraction of vertical velocity - # left after doing rayleigh damping (w -> w * rf) - rf = compute_rff_vals(pfull, dt, rf_cutoff, tau * SECONDS_PER_DAY, ptop) with computation(FORWARD): with interval(0, 1): if pfull < rf_cutoff_nudge: @@ -191,6 +209,29 @@ def __init__( }, ) + # We compute the damping increment once using a trick to write a + # FloatFieldK as a (1, 1, K) 3D writable Field + K_stencil_factory, K_quantity_factory = get_factories_single_tile( + 1, + 1, + domain[2], + 0, + stencil_factory.backend, + ) + self._ray_fast_damping_increment = K_stencil_factory.from_origin_domain( + ray_fast_damping_increment, + origin=(0, 0, origin[2]), + domain=(1, 1, domain[2]), + externals={ + "rf_cutoff": self._rf_cutoff, + "tau": tau, + }, + ) + self._damping_increment = K_quantity_factory.ones( + [X_DIM, Y_DIM, Z_DIM], units="n/a" + ) + self._initialize_damping_increment = False + def __call__( self, u: FloatField, @@ -203,13 +244,17 @@ def __call__( ): rf_cutoff_nudge = self._rf_cutoff + min(Float(100.0), Float(10.0) * ptop) + if not self._initialize_damping_increment: + self._ray_fast_damping_increment( + pfull=pfull, dt=dt, ptop=ptop, rf=self._damping_increment + ) + self._initialize_damping_increment = True self._ray_fast_wind_compute( - u, - v, - w, - dp, - pfull, - dt, - ptop, - rf_cutoff_nudge, + u=u, + v=v, + w=w, + delta_p_ref=dp, + pfull=pfull, + rf=self._damping_increment.view[0, 0, :], + rf_cutoff_nudge=rf_cutoff_nudge, ) From 4b45e2943f838923bb433d0dadbdf4eb8dec181b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 5 Mar 2025 17:01:07 -0500 Subject: [PATCH 171/252] Update to new NDSL API --- pyFV3/stencils/fv_dynamics.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index cb969a3b..500064f1 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -25,7 +25,7 @@ FloatField, FloatField64, FloatFieldIJ64, - global_set_floating_point_precision, + get_precision, NDSL_32BIT_FLOAT_TYPE, NDSL_64BIT_FLOAT_TYPE, ) @@ -458,9 +458,7 @@ def __init__( self._timestep = timestep.total_seconds() # At 32-bit precision we still need - self._f32_correction = ( - global_set_floating_point_precision() == NDSL_32BIT_FLOAT_TYPE - ) + self._f32_correction = get_precision() == NDSL_32BIT_FLOAT_TYPE if self._f32_correction: self._mfx_f64 = quantity_factory.zeros( dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], From 00f2f593f8760cafaf0341de58ebf77392021f4a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 10 Mar 2025 12:48:33 -0400 Subject: [PATCH 172/252] Fix `f32_correction` + lint --- pyFV3/stencils/fv_dynamics.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 500064f1..4f2df9ee 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -2,7 +2,7 @@ from typing import Mapping, Optional from dace.frontend.python.interface import nounroll as dace_no_unroll -from gt4py.cartesian.gtscript import PARALLEL, FORWARD, computation, interval +from gt4py.cartesian.gtscript import FORWARD, PARALLEL, computation, interval import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.moist_cv as moist_cv @@ -12,22 +12,21 @@ KAPPA, NQ, X_DIM, + X_INTERFACE_DIM, Y_DIM, + Y_INTERFACE_DIM, Z_DIM, Z_INTERFACE_DIM, ZVIR, - Y_INTERFACE_DIM, - X_INTERFACE_DIM, ) from ndsl.dsl.dace.orchestration import dace_inhibitor, orchestrate from ndsl.dsl.typing import ( + NDSL_64BIT_FLOAT_TYPE, Float, FloatField, FloatField64, FloatFieldIJ64, get_precision, - NDSL_32BIT_FLOAT_TYPE, - NDSL_64BIT_FLOAT_TYPE, ) from ndsl.grid import DampingCoefficients, GridData from ndsl.logging import ndsl_log @@ -38,12 +37,12 @@ from pyFV3._config import DynamicalCoreConfig from pyFV3.dycore_state import DycoreState from pyFV3.stencils import fvtp2d, tracer_2d_1l +from pyFV3.stencils.compute_total_energy import ComputeTotalEnergy from pyFV3.stencils.del2cubed import HyperdiffusionDamping from pyFV3.stencils.dyn_core import AcousticDynamics from pyFV3.stencils.neg_adj3 import AdjustNegativeTracerMixingRatio from pyFV3.stencils.remapping import LagrangianToEulerian from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS -from pyFV3.stencils.compute_total_energy import ComputeTotalEnergy from pyFV3.version import IS_GEOS @@ -458,7 +457,7 @@ def __init__( self._timestep = timestep.total_seconds() # At 32-bit precision we still need - self._f32_correction = get_precision() == NDSL_32BIT_FLOAT_TYPE + self._f32_correction = get_precision() == 32 if self._f32_correction: self._mfx_f64 = quantity_factory.zeros( dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], From 1c2be5d032ada3c6de52b790a12753f37286fe67 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 11 Mar 2025 14:58:32 -0400 Subject: [PATCH 173/252] Fix inner power calculation with a mixed precision factor --- pyFV3/stencils/delnflux.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/pyFV3/stencils/delnflux.py b/pyFV3/stencils/delnflux.py index ed4304f1..36f5ea02 100644 --- a/pyFV3/stencils/delnflux.py +++ b/pyFV3/stencils/delnflux.py @@ -17,7 +17,12 @@ def calc_damp(damp_c: Quantity, da_min: Float, nord: Quantity) -> Quantity: "current implementation requires damp_c and nord to have " "identical data shape and dims" ) - data = np.power((damp_c.data * da_min), (nord.data + 1), dtype=Float) + # `da_min` is a 64 bit float and we have to cast the array to deal + # with downcasting behavior of array * scalar in numpy + # We then reproduce the proper casting so `calc_damp` is a 32-bit float + data = np.power( + (damp_c.data.astype(np.float64) * da_min), (nord.data + 1), dtype=np.float64 + ).astype(Float) return Quantity( data=data, dims=damp_c.dims, @@ -371,6 +376,7 @@ def __init__( rarea: Quantity, nord_col: Quantity, damp_c: Quantity, + damp_coeff: Quantity | None = None, ): """ nord sets the order of damping to apply: @@ -419,9 +425,14 @@ def __init__( func=diffusive_damp, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM] ) - self._damp = calc_damp( - damp_c=damp_c, da_min=damping_coefficients.da_min, nord=nord_col - ) + damp_c.to_netcdf("damp_c.nc4") + nord_col.to_netcdf("nord_col.nc4") + if damp_coeff is None: + self._damp = calc_damp( + damp_c=damp_c, da_min=damping_coefficients.da_min, nord=nord_col + ) + else: + self._damp = damp_coeff self.delnflux_nosg = DelnFluxNoSG( stencil_factory, damping_coefficients, rarea, nord_col, nk=nk From 4bb5cffea4ad0d71e83264e592d6535493b76aa1 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 11 Mar 2025 15:47:38 -0400 Subject: [PATCH 174/252] Restore `cd` to it's original 64-bit factor --- pyFV3/stencils/del2cubed.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/del2cubed.py b/pyFV3/stencils/del2cubed.py index 24e1976e..6a887469 100644 --- a/pyFV3/stencils/del2cubed.py +++ b/pyFV3/stencils/del2cubed.py @@ -1,3 +1,4 @@ +import numpy as np from gt4py.cartesian.gtscript import PARALLEL, computation, horizontal, interval, region import ndsl.stencils.corners as corners @@ -68,7 +69,7 @@ def corner_fill(q_in: FloatField, q_out: FloatField): # Q update stencil # ------------------ def update_q( - q: FloatField, rarea: FloatFieldIJ, fx: FloatField, fy: FloatField, cd: Float + q: FloatField, rarea: FloatFieldIJ, fx: FloatField, fy: FloatField, cd: np.float64 ): with computation(PARALLEL), interval(...): q += cd * rarea * (fx - fx[1, 0, 0] + fy - fy[0, 1, 0]) @@ -169,7 +170,7 @@ def __init__( update_q, origins, domains, stencil_factory=stencil_factory ) - def __call__(self, qdel: FloatField, cd: Float): + def __call__(self, qdel: FloatField, cd: np.float64): """ Perform hyperdiffusion damping/filtering. From cf020438f242b1d2dbc8d9751392b9d1c8d405cb Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 11 Mar 2025 15:49:05 -0400 Subject: [PATCH 175/252] Update README.md --- README.md | 3 +++ 1 file changed, 3 insertions(+) diff --git a/README.md b/README.md index c01225c4..b100bfb3 100644 --- a/README.md +++ b/README.md @@ -31,12 +31,14 @@ Branches: - 🔶 `fix/RayleighDamping_mixed_precision`@Florian: fix the Ray_Fast test - 🔶 `GEOS_update/yppm_xppm`@Florian: fix the YPPM/XPPM with `hord = -6` - 🔶 `fix/DelnFlux_f32_support`@Florian: Fix for f32 support for DelnFlux (partial pass) +- 🔶 `fix/GEOSv11_4_2/HyperDiffusionDamping`@Florian: fix the Hyperdiffusion Damping by restoring factor to be 64-bit float - ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation, column calculation and new `dpx` accumulation (partial pass) - ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit A2B_Ord4 - ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit - ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit and `dpx` calculation - MERGE ORDER: after `fix/GEOS/D_SW` + - MERGE ORDER: after `fix/GEOSv11_4_2/HyperDiffusionDamping` - ⚙️ `feature/tracer_rework_part1` @Florian: Allow for update of N Tracers - ⚙️ `fix/GEOS/TracerAdvection` @Florian: Allow for non-update of mass fluxes and courant number, f32 fixes, correct computation of `cmax` and `nsplit`, overcomputation into the algorithm - BASED ON `tracer_rework_part1` @@ -45,3 +47,4 @@ Branches: - ⚙️ `fix/GEOSv11_4_2/Dynamics`@Florian: Fix for the f32 & GEOS version of dynamics - REQUIRES: `ndsl` with tracer rework - REQUIRES: `tracer_rework_part1`, `fix/GEOSv11_4_2/Dyncore`, `fix/GEOS/TracerAdvection` + - MERGE ORDER: after `fix/GEOSv11_4_2/HyperDiffusionDamping` From eca6ee22aa130499146e1f8e3b5373a78150927b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 11 Mar 2025 15:52:13 -0400 Subject: [PATCH 176/252] Restore `da_min` true 64-bit type --- pyFV3/stencils/dyn_core.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 426f11bc..2e235f28 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -1,4 +1,4 @@ -from typing import Dict, Mapping, Optional +from typing import Dict, Mapping import numpy as np from dace.frontend.python.interface import nounroll as dace_nounroll @@ -620,8 +620,8 @@ def __init__( # See divergence_damping.py, _get_da_min for explanation of this function @dace_inhibitor - def _get_da_min(self) -> Float: - return Float(self._da_min) + def _get_da_min(self) -> np.float64: + return self._da_min # TODO: fix me - we shouldn't need a function here, Dace is fudging the types # See https://github.com/GEOS-ESM/pace/issues/9 From c11348608db885347ed331d459c50750c29545f9 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 11 Mar 2025 15:54:25 -0400 Subject: [PATCH 177/252] Remove debuggin artefacts --- pyFV3/stencils/delnflux.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/pyFV3/stencils/delnflux.py b/pyFV3/stencils/delnflux.py index 36f5ea02..2b96032f 100644 --- a/pyFV3/stencils/delnflux.py +++ b/pyFV3/stencils/delnflux.py @@ -376,7 +376,6 @@ def __init__( rarea: Quantity, nord_col: Quantity, damp_c: Quantity, - damp_coeff: Quantity | None = None, ): """ nord sets the order of damping to apply: @@ -425,14 +424,9 @@ def __init__( func=diffusive_damp, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM] ) - damp_c.to_netcdf("damp_c.nc4") - nord_col.to_netcdf("nord_col.nc4") - if damp_coeff is None: - self._damp = calc_damp( - damp_c=damp_c, da_min=damping_coefficients.da_min, nord=nord_col - ) - else: - self._damp = damp_coeff + self._damp = calc_damp( + damp_c=damp_c, da_min=damping_coefficients.da_min, nord=nord_col + ) self.delnflux_nosg = DelnFluxNoSG( stencil_factory, damping_coefficients, rarea, nord_col, nk=nk From 82ba97c6853b6781f7b1aa9fae1fc53dd9790df4 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 12 Mar 2025 12:06:13 -0400 Subject: [PATCH 178/252] Fix `dd8` power function Fix using proper precision for `da_min` Fix: use proper `damp_tmp2` for `ke` --- pyFV3/stencils/divergence_damping.py | 61 +++++++++++++--------------- 1 file changed, 28 insertions(+), 33 deletions(-) diff --git a/pyFV3/stencils/divergence_damping.py b/pyFV3/stencils/divergence_damping.py index 30c47a59..e42d86c6 100644 --- a/pyFV3/stencils/divergence_damping.py +++ b/pyFV3/stencils/divergence_damping.py @@ -4,6 +4,7 @@ __INLINED, PARALLEL, computation, + f32, horizontal, interval, region, @@ -24,14 +25,13 @@ @gtscript.function def damp_tmp(q, da_min_c, d2_bg, dddmp): - mintmp = min(0.2, dddmp * abs(q)) - damp = da_min_c * max(d2_bg, mintmp) + damp: f32 = da_min_c * max(d2_bg, min(0.2, dddmp * abs(q))) return damp + @gtscript.function def damp_tmp2(q, da_min_c, d2_bg, dddmp): - mintmp = min(0.2, dddmp * q) - damp = da_min_c * max(d2_bg, mintmp) + damp: f32 = da_min_c * max(d2_bg, min(0.2, dddmp * q)) return damp @@ -60,10 +60,11 @@ def compute_u_contra_dyc( u_contra_dyc (out): contravariant u-wind on d-grid Porting Notes - * The compute_u_contra_dyc and compute_v_contra_dxc functions have the dyc and dxc values - incorporated earlier in the calcuation rather than later, and this enables the u_contra_dyc - and v_contra_dxc values to match with the Fortran. As a result, the delpc computation - matches the Fortran value of delpc. + * The compute_u_contra_dyc and compute_v_contra_dxc functions have + the dyc and dxc values incorporated earlier in the calcuation rather than later, + and this enables the u_contra_dyc and v_contra_dxc values + to match with the Fortran. + As a result, the delpc computation matches the Fortran value of delpc. Ex : Previous implementation of compute_u_contra_dyc ================================================================= @@ -82,7 +83,7 @@ def compute_u_contra_dyc( u_contra_dyc = contravariant(u, vc_from_va, cosa_v, dyc) u_contra_dyc = u_contra_dyc * sina_v with horizontal(region[:, j_start], region[:, j_end + 1]): - u_contra_dyc = u * dyc* sin_sg4[0, -1] if vc > 0 else u * dyc*sin_sg2 + u_contra_dyc = u * dyc * sin_sg4[0, -1] if vc > 0 else u * dyc * sin_sg2 def compute_v_contra_dxc( @@ -109,9 +110,10 @@ def compute_v_contra_dxc( sin_sg1 (in): Porting Notes - * The compute_u_contra_dyc and compute_v_contra_dxc functions have the dyc and dxc values - incorporated earlier in the calcuation rather than later, and this enables the u_contra_dyc - and v_contra_dxc values to match with the Fortran. As a result, the delpc computation + * The compute_u_contra_dyc and compute_v_contra_dxc functions + have the dyc and dxc values incorporated earlier in the calcuation + rather than later, and this enables the u_contra_dyc and v_contra_dxc + values to match with the Fortran. As a result, the delpc computation matches the Fortran value of delpc. Ex : Previous implementation of compute_v_contra_dxc @@ -130,7 +132,7 @@ def compute_v_contra_dxc( v_contra_dxc = contravariant(v, uc_from_ua, cosa_u, dxc) v_contra_dxc = v_contra_dxc * sina_u with horizontal(region[i_start, :], region[i_end + 1, :]): - v_contra_dxc = v * dxc*sin_sg3[-1, 0] if uc > 0 else v * dxc*sin_sg1 + v_contra_dxc = v * dxc * sin_sg3[-1, 0] if uc > 0 else v * dxc * sin_sg1 def delpc_computation( @@ -175,7 +177,7 @@ def damping( vort: FloatField, ke: FloatField, d2_bg: FloatFieldK, - da_min_c: Float, + da_min_c: np.float64, dddmp: Float, dt: Float, ): @@ -199,7 +201,7 @@ def damping_nord_highorder_stencil( delpc: FloatField, divg_d: FloatField, d2_bg: FloatFieldK, - da_min_c: Float, + da_min_c: np.float64, dddmp: Float, dd8: Float, ): @@ -215,7 +217,7 @@ def damping_nord_highorder_stencil( """ # TODO: propagate variable renaming into this routine with computation(PARALLEL), interval(...): - damp = damp_tmp(vort, da_min_c, d2_bg, dddmp) + damp = damp_tmp2(vort, da_min_c, d2_bg, dddmp) vort = damp * delpc + dd8 * divg_d ke = ke + vort @@ -283,7 +285,7 @@ def smagorinsky_diffusion_approx(delpc: FloatField, vort: FloatField, absdt: Flo absdt (in): abs(dt) """ with computation(PARALLEL), interval(...): - vort = absdt * (delpc ** 2.0 + vort ** 2.0) ** 0.5 + vort = absdt * sqrt(delpc ** 2 + vort ** 2) def smag_corner( @@ -573,12 +575,12 @@ def __init__( # odd and adds a lot of boilerplate throughout the model code. @dace_inhibitor - def _get_da_min_c(self) -> Float: - return Float(self._damping_coefficients.da_min_c) + def _get_da_min_c(self) -> np.float64: + return self._damping_coefficients.da_min_c @dace_inhibitor - def _get_da_min(self) -> Float: - return Float(self._damping_coefficients.da_min) + def _get_da_min(self) -> np.float64: + return self._damping_coefficients.da_min def __call__( self, @@ -607,13 +609,6 @@ def __call__( Applies both a background second-order diffusion (with strength controlled by d2_bg passed on init) and a higher-order hyperdiffusion. - Porting Notes - * The dd8 computation has different results when comparing between Fortran and Python, - which is likely due to the user of the power function. The difference in dd8 results in - the ke value having on the order of 10,000 total error difference when running the translate - test. - - Args: u (in): x-velocity on d-grid v (in): y-velocity on d-grid @@ -673,7 +668,7 @@ def __call__( self.v_contra_dxc, ) - da_min_c: Float = self._get_da_min_c() + da_min_c = self._get_da_min_c() self._damping( delpc, damped_rel_vort_bgrid, @@ -735,13 +730,13 @@ def __call__( abs(dt), ) - da_min: Float = self._get_da_min() + da_min = self._get_da_min() if self._stretched_grid: # reference https://github.com/NOAA-GFDL/GFDL_atmos_cubed_sphere/blob/main/model/sw_core.F90#L1422 # noqa: E501 - dd8 = da_min * np.power(self._d4_bg, (self._nonzero_nord + 1), dtype=Float) + dd8 = Float(da_min * np.power(self._d4_bg, (self._nonzero_nord + 1))) else: - dd8 = np.power( - (da_min_c * self._d4_bg), (self._nonzero_nord + 1), dtype=Float + dd8 = np.power((da_min_c * self._d4_bg), (self._nonzero_nord + 1)).astype( + Float ) self._damping_nord_highorder_stencil( From 11788b5abede75392e87c4b6ae5a1c20a177980a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 17 Mar 2025 09:10:30 -0400 Subject: [PATCH 179/252] Match Fortran by multiplyin by 1/GRAV instead of divide --- pyFV3/stencils/riem_solver_c.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/riem_solver_c.py b/pyFV3/stencils/riem_solver_c.py index e498b676..7390405a 100644 --- a/pyFV3/stencils/riem_solver_c.py +++ b/pyFV3/stencils/riem_solver_c.py @@ -65,7 +65,7 @@ def precompute( dz = gz[0, 0, 1] - gz with computation(PARALLEL), interval(...): gm = 1.0 / (1.0 - cappa) - dm /= constants.GRAV + dm *= constants.RGRAV with computation(PARALLEL), interval(0, -1): # (1) From \partial p*/\partial z = -\rho g, we can separate and integrate # over a layer to get From 72465a6d02e21d78cdf35c5ead7edae21e636968 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:26:43 -0400 Subject: [PATCH 180/252] Fix --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b100bfb3..728dd95c 100644 --- a/README.md +++ b/README.md @@ -34,7 +34,7 @@ Branches: - 🔶 `fix/GEOSv11_4_2/HyperDiffusionDamping`@Florian: fix the Hyperdiffusion Damping by restoring factor to be 64-bit float - ⚙️ `fix/GEOS/D_SW`@Florian: Fix D_SW heat dissipation, column calculation and new `dpx` accumulation (partial pass) - ⚙️ `fix/GEOSv11_4_2/A2B_Ord4`@Florian: Fix for 32-bit A2B_Ord4 -- ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit A2B_Ord4 +- ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit RiemanSolver - ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit - ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit and `dpx` calculation - MERGE ORDER: after `fix/GEOS/D_SW` From e1bc27b1ce4a8fce32f17cbfab0a54344a3ed8f3 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:27:59 -0400 Subject: [PATCH 181/252] Undo `dt` refactoring in interpolation --- pyFV3/stencils/d_sw.py | 66 ++++++++++++++++++++---------------------- 1 file changed, 32 insertions(+), 34 deletions(-) diff --git a/pyFV3/stencils/d_sw.py b/pyFV3/stencils/d_sw.py index 754bdd54..d70276e7 100644 --- a/pyFV3/stencils/d_sw.py +++ b/pyFV3/stencils/d_sw.py @@ -8,6 +8,8 @@ horizontal, interval, region, + I, + J, ) from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate @@ -16,7 +18,6 @@ from ndsl.grid import DampingCoefficients, GridData from pyFV3._config import DGridShallowWaterLagrangianDynamicsConfig from pyFV3.stencils import delnflux -from pyFV3.stencils.d2a2c_vect import contravariant from pyFV3.stencils.delnflux import DelnFluxNoSG from pyFV3.stencils.divergence_damping import DivergenceDamping from pyFV3.stencils.fvtp2d import FiniteVolumeTransport @@ -254,20 +255,22 @@ def compute_kinetic_energy( from __externals__ import grid_type with computation(PARALLEL), interval(...): + dt4 = 0.25 * dt + dt5 = 0.5 * dt if __INLINED(grid_type < 3): ub_contra, vb_contra = interpolate_uc_vc_to_cell_corners( - uc, vc, cosa, rsina, uc_contra, vc_contra + uc, vc, cosa, rsina, uc_contra, vc_contra, dt4, dt5 ) else: - ub_contra = 0.5 * (uc[0, -1, 0] + uc) - vb_contra = 0.5 * (vc[-1, 0, 0] + vc) + ub_contra = dt5 * (uc[0, -1, 0] + uc) + vb_contra = dt5 * (vc[-1, 0, 0] + vc) advected_v = advect_v_along_y(v, vb_contra, rdy=rdy, dy=dy, dya=dya, dt=dt) advected_u = advect_u_along_x(u, ub_contra, rdx=rdx, dx=dx, dxa=dxa, dt=dt) # makes sure the kinetic energy part of the governing equation is computed # the same way as the vorticity flux part (in terms of time splitting) # to avoid a Hollingsworth-Kallberg instability - dt_kinetic_energy_on_cell_corners = ( - 0.5 * dt * (ub_contra * advected_u + vb_contra * advected_v) + dt_kinetic_energy_on_cell_corners = 0.5 * ( + ub_contra * advected_u + vb_contra * advected_v ) dt_kinetic_energy_on_cell_corners = all_corners_ke( dt_kinetic_energy_on_cell_corners, u, v, uc_contra, vc_contra, dt @@ -339,11 +342,9 @@ def compute_vorticity( # cell-mean vorticity is equal to the circulation around the gridcell # divided by the area of the gridcell. It isn't exactly true that # area = dx * dy, so the form below is necessary to get an exact result. - rdy_tmp = rarea * dx - rdx_tmp = rarea * dy - vorticity = (u - u[0, 1, 0] * dx[0, 1] / dx) * rdy_tmp + ( - v[1, 0, 0] * dy[1, 0] / dy - v - ) * rdx_tmp + ut = v * dy + vt = u * dx + vorticity = rarea * (vt - vt[J + 1] - ut + ut[I + 1]) def adjust_w_and_qcon( @@ -742,42 +743,39 @@ def get_column_namelist( @gtscript.function def interpolate_uc_vc_to_cell_corners( - uc_cov, vc_cov, cosa, rsina, uc_contra, vc_contra + uc_cov, vc_cov, cosa, rsina, uc_contra, vc_contra, dt4, dt5 ): """ Convert covariant C-grid winds to contravariant B-grid (cell-corner) winds. """ from __externals__ import i_end, i_start, j_end, j_start - # In the original Fortran, this routine was given dt4 (0.25 * dt) - # and dt5 (0.5 * dt), and its outputs were wind times timestep. This has - # been refactored so the timestep is later explicitly multiplied, when - # the wind is integrated forward in time. - # TODO: ask Lucas why we interpolate then convert to contravariant in tile center, - # but convert to contravariant and then interpolate on tile edges. - ub_cov = 0.5 * (uc_cov[0, -1, 0] + uc_cov) - vb_cov = 0.5 * (vc_cov[-1, 0, 0] + vc_cov) - ub_contra = contravariant(ub_cov, vb_cov, cosa, rsina) - vb_contra = contravariant(vb_cov, ub_cov, cosa, rsina) - # ASSUME : if __INLINED(namelist.grid_type < 3): + # Orders matter because corners take the last edge computation values + + # Center domain + ub = dt5 * (uc_cov[J - 1] + uc_cov - (vc_cov[I - 1] + vc_cov) * cosa) * rsina + vb = dt5 * (vc_cov[I - 1] + vc_cov - (uc_cov[J - 1] + uc_cov) * cosa) * rsina + # UB - Orders matter because corners take the last edge computation values + # North/South edge with horizontal(region[:, j_start], region[:, j_end + 1]): - ub_contra = 0.25 * ( - -uc_contra[0, -2, 0] - + 3.0 * (uc_contra[0, -1, 0] + uc_contra) - - uc_contra[0, 1, 0] + ub = dt4 * ( + -uc_contra[J - 2] + 3.0 * (uc_contra[J - 1] + uc_contra) - uc_contra[J + 1] ) + # East/West with horizontal(region[i_start, :], region[i_end + 1, :]): - ub_contra = 0.5 * (uc_contra[0, -1, 0] + uc_contra) + ub = dt5 * (uc_contra[J - 1] + uc_contra) + + # VB - Orders matter because corners take the last edge computation values + # North/South edge with horizontal(region[i_start, :], region[i_end + 1, :]): - vb_contra = 0.25 * ( - -vc_contra[-2, 0, 0] - + 3.0 * (vc_contra[-1, 0, 0] + vc_contra) - - vc_contra[1, 0, 0] + vb = dt4 * ( + -vc_contra[I - 2] + 3.0 * (vc_contra[I - 1] + vc_contra) - vc_contra[I + 1] ) + # East/West with horizontal(region[:, j_start], region[:, j_end + 1]): - vb_contra = 0.5 * (vc_contra[-1, 0, 0] + vc_contra) + vb = dt5 * (vc_contra[I - 1] + vc_contra) - return ub_contra, vb_contra + return ub, vb class DGridShallowWaterLagrangianDynamics: From 3450b0d6be8563b3386fd3054092ab29e75e0c39 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:28:15 -0400 Subject: [PATCH 182/252] Adjust contravriant functions --- pyFV3/stencils/xtp_u.py | 2 +- pyFV3/stencils/ytp_v.py | 2 +- tests/savepoint/translate/translate_d_sw.py | 12 ++++++------ 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pyFV3/stencils/xtp_u.py b/pyFV3/stencils/xtp_u.py index 9b9bd0ac..94e7e34f 100644 --- a/pyFV3/stencils/xtp_u.py +++ b/pyFV3/stencils/xtp_u.py @@ -85,7 +85,7 @@ def advect_u_along_x( bl, br = get_bl_br(u, dx, dxa) b0 = bl + br - cfl = ub_contra * dt * rdx[-1, 0] if ub_contra > 0 else ub_contra * dt * rdx + cfl = ub_contra * rdx[-1, 0] if ub_contra > 0 else ub_contra * rdx fx0 = xppm.fx1_fn(cfl, br, b0, bl) if __INLINED(iord < 8): diff --git a/pyFV3/stencils/ytp_v.py b/pyFV3/stencils/ytp_v.py index 4c1e4af9..a2ebc6ee 100644 --- a/pyFV3/stencils/ytp_v.py +++ b/pyFV3/stencils/ytp_v.py @@ -84,7 +84,7 @@ def advect_v_along_y( bl, br = get_bl_br(v, dy, dya) b0 = bl + br - cfl = vb_contra * dt * rdy[0, -1] if vb_contra > 0 else vb_contra * dt * rdy + cfl = vb_contra * rdy[0, -1] if vb_contra > 0 else vb_contra * rdy fx0 = yppm.fx1_fn(cfl, br, b0, bl) if __INLINED(jord < 8): diff --git a/tests/savepoint/translate/translate_d_sw.py b/tests/savepoint/translate/translate_d_sw.py index ea042712..34f34cd7 100644 --- a/tests/savepoint/translate/translate_d_sw.py +++ b/tests/savepoint/translate/translate_d_sw.py @@ -77,9 +77,9 @@ def ubke( dt5: Float, ): with computation(PARALLEL), interval(...): - dt = 2.0 * dt5 - ub, _ = d_sw.interpolate_uc_vc_to_cell_corners(uc, vc, cosa, rsina, ut, ut) - ub = ub * dt + ub, _ = d_sw.interpolate_uc_vc_to_cell_corners( + uc, vc, cosa, rsina, ut, ut, dt4, dt5 + ) class TranslateUbKE(TranslateDycoreFortranData2Py): @@ -124,9 +124,9 @@ def vbke( dt5: Float, ): with computation(PARALLEL), interval(...): - dt = 2.0 * dt5 - _, vb = d_sw.interpolate_uc_vc_to_cell_corners(uc, vc, cosa, rsina, vt, vt) - vb = vb * dt + _, vb = d_sw.interpolate_uc_vc_to_cell_corners( + uc, vc, cosa, rsina, vt, vt, dt4, dt5 + ) class TranslateVbKE(TranslateDycoreFortranData2Py): From e077ab454340a476075396bc146f757a70071b55 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:29:08 -0400 Subject: [PATCH 183/252] Docs --- pyFV3/stencils/ray_fast.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index a28e8946..e1dea877 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -242,6 +242,16 @@ def __call__( dt: Float, ptop: Float, ): + """ + Args: + u (inout) + v (inout) + w (inout) + dp (in) + pfull (in) + dt (in) + ptop (in) + """ rf_cutoff_nudge = self._rf_cutoff + min(Float(100.0), Float(10.0) * ptop) if not self._initialize_damping_increment: From 8e631e09d98a64a34df2b66c15816852a505f15c Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:29:54 -0400 Subject: [PATCH 184/252] Rename a2b_ord4 use for clarity --- pyFV3/stencils/nh_p_grad.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/nh_p_grad.py b/pyFV3/stencils/nh_p_grad.py index 03988d0d..bf540255 100644 --- a/pyFV3/stencils/nh_p_grad.py +++ b/pyFV3/stencils/nh_p_grad.py @@ -172,7 +172,7 @@ def __init__( z_dim=Z_INTERFACE_DIM, replace=True, ) - self.a2b_kbuffer = AGrid2BGridFourthOrder( + self.a2b_kinterface = AGrid2BGridFourthOrder( stencil_factory, quantity_factory=quantity_factory, grid_data=grid_data, @@ -245,7 +245,7 @@ def __call__( self.a2b_k1(pp, self._tmp_wk1) self.a2b_k1(pk3, self._tmp_wk1) - self.a2b_kbuffer(gz, self._tmp_wk1) + self.a2b_kinterface(gz, self._tmp_wk1) self.a2b_kstandard(delp, self._tmp_wk1) self._set_k0_and_calc_wk_stencil(pp, pk3, self._tmp_wk, top_value) From 5083d0d327a0e02f7127d219fa48f6ab021e3808 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:31:29 -0400 Subject: [PATCH 185/252] Match Fortan circulation c grid operation factor Do not inline sub-dt calculation --- pyFV3/stencils/c_sw.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/pyFV3/stencils/c_sw.py b/pyFV3/stencils/c_sw.py index 9208e472..b43c4c0f 100644 --- a/pyFV3/stencils/c_sw.py +++ b/pyFV3/stencils/c_sw.py @@ -5,6 +5,8 @@ horizontal, interval, region, + I, + J, ) from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate @@ -378,7 +380,8 @@ def transportdelp_update_vorticity_and_kineticenergy( with horizontal(region[i_end + 1, :], region[i_start, :]): ke = ke * sin_sg1 + v * cos_sg1 if ua > 0.0 else ke - ke = 0.5 * dt2 * (ua * ke + va * vort) + dt4 = 0.5 * dt2 + ke = dt4 * (ua * ke + va * vort) def circulation_cgrid( @@ -400,18 +403,18 @@ def circulation_cgrid( from __externals__ import i_end, i_start, j_end, j_start with computation(PARALLEL), interval(...): - fx = dxc * uc - fy = dyc * vc - # fx1 and fy1 are the shifted versions of fx and fy and are defined - # because temporaries are not allowed to be accessed with offsets in regions. - fx1 = dxc[0, -1] * uc[0, -1, 0] - fy1 = dyc[-1, 0] * vc[-1, 0, 0] - - vort_c = fx1 - fx - fy1 + fy + fx = uc * dxc + fy = vc * dyc + + vort_c = fx[J - 1] - fx - fy[I - 1] + fy + + # Remove the extra term at the corners + # WEST with horizontal(region[i_start, j_start], region[i_start, j_end + 1]): - vort_c = fx1 - fx + fy + vort_c = vort_c + (vc[I - 1] * dyc[I - 1]) + # EAST with horizontal(region[i_end + 1, j_start], region[i_end + 1, j_end + 1]): - vort_c = fx1 - fx - fy1 + vort_c = vort_c - fy def absolute_vorticity(vort: FloatField, fC: FloatFieldIJ, rarea_c: FloatFieldIJ): From 26635057d55b6e68a1f35a358cef913bd380416a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:32:15 -0400 Subject: [PATCH 186/252] Match fortran on cubic interpolation & edge interpolation Fix order of op for UT update --- pyFV3/stencils/d2a2c_vect.py | 45 +++++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/pyFV3/stencils/d2a2c_vect.py b/pyFV3/stencils/d2a2c_vect.py index 292b407a..6436fdf9 100644 --- a/pyFV3/stencils/d2a2c_vect.py +++ b/pyFV3/stencils/d2a2c_vect.py @@ -3,7 +3,7 @@ from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, get_precision from ndsl.grid import GridData from ndsl.stencils import corners from pyFV3.stencils.a2b_ord4 import a1, a2, lagrange_x_func, lagrange_y_func @@ -140,7 +140,7 @@ def east_west_edges( uc = utc * sin_sg3[-1, 0] if utc > 0 else utc * sin_sg1 with horizontal(region[i_end + 2, local_js - 1 : local_je + 2]): - uc = vol_conserv_cubic_interp_func_x_rev(utmp) + uc = vol_conserv_cubic_interp_func_x_rev_2(utmp) with horizontal(region[i_end, local_js - 1 : local_je + 2]): utc = contravariant(uc, v, cosa_u, rsin_u) @@ -300,6 +300,13 @@ def vol_conserv_cubic_interp_func_x_rev(u): return c1 * u[1, 0, 0] + c2 * u + c3 * u[-1, 0, 0] +@gtscript.function +def vol_conserv_cubic_interp_func_x_rev_2(u): + """Series order is reversed compared to original + vol_conserv_cubic_interp_func_x_rev to match Fortran""" + return c3 * u[-1, 0, 0] + c2 * u + c1 * u[1, 0, 0] + + @gtscript.function def vol_conserv_cubic_interp_func_y(v): return c1 * v[0, -2, 0] + c2 * v[0, -1, 0] + c3 * v @@ -361,18 +368,20 @@ def vc_y_edge1( def edge_interpolate4_x(ua, dxa): t1 = dxa[-2, 0] + dxa[-1, 0] t2 = dxa[0, 0] + dxa[1, 0] - n1 = (t1 + dxa[-1, 0]) * ua[-1, 0, 0] - dxa[-1, 0] * ua[-2, 0, 0] - n2 = (t1 + dxa[0, 0]) * ua[0, 0, 0] - dxa[0, 0] * ua[1, 0, 0] - return 0.5 * (n1 / t1 + n2 / t2) + return 0.5 * ( + ((t1 + dxa[-1, 0]) * ua[-1, 0, 0] - dxa[-1, 0] * ua[-2, 0, 0]) / t1 + + ((t1 + dxa[0, 0]) * ua[0, 0, 0] - dxa[0, 0] * ua[1, 0, 0]) / t2 + ) @gtscript.function def edge_interpolate4_y(va, dya): t1 = dya[0, -2] + dya[0, -1] t2 = dya[0, 0] + dya[0, 1] - n1 = (t1 + dya[0, -1]) * va[0, -1, 0] - dya[0, -1] * va[0, -2, 0] - n2 = (t1 + dya[0, 0]) * va[0, 0, 0] - dya[0, 0] * va[0, 1, 0] - return 0.5 * (n1 / t1 + n2 / t2) + return 0.5 * ( + ((t1 + dya[0, -1]) * va[0, -1, 0] - dya[0, -1] * va[0, -2, 0]) / t1 + + ((t1 + dya[0, 0]) * va[0, 0, 0] - dya[0, 0] * va[0, 1, 0]) / t2 + ) class DGrid2AGrid2CGridVectors: @@ -409,7 +418,7 @@ def __init__( self._sin_sg4 = grid_data.sin_sg4 self._grid_type = grid_type - self._big_number = Float(1e30) # 1e8 if 32 bit + self._big_number = Float(1e30) if get_precision() == 64 else Float(1e8) nx = grid_indexing.iec + 1 # grid.npx + 2 ny = grid_indexing.jec + 1 # grid.npy + 2 i1 = grid_indexing.isc - 1 @@ -640,15 +649,6 @@ def __call__(self, uc, vc, u, v, ua, va, utc, vtc): va, ) - self._ut_main( - self._utmp, - uc, - v, - self._cosa_u, - self._rsin_u, - utc, - ) - if self._grid_type < 3: self._east_west_edges( u, @@ -664,6 +664,15 @@ def __call__(self, uc, vc, u, v, ua, va, utc, vtc): self._dxa, ) + self._ut_main( + self._utmp, + uc, + v, + self._cosa_u, + self._rsin_u, + utc, + ) + # Ydir: self._fill_corners_y( self._utmp, From 9d17b56a3769f51ba2b905dfc90369233eb1397b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:34:30 -0400 Subject: [PATCH 187/252] PK3 Halo match the fortran --- pyFV3/stencils/pk3_halo.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/pk3_halo.py b/pyFV3/stencils/pk3_halo.py index 6daf945a..20436ac9 100644 --- a/pyFV3/stencils/pk3_halo.py +++ b/pyFV3/stencils/pk3_halo.py @@ -1,4 +1,12 @@ -from gt4py.cartesian.gtscript import FORWARD, computation, horizontal, interval, region +from gt4py.cartesian.gtscript import ( + FORWARD, + computation, + horizontal, + interval, + region, + log, + exp, +) from ndsl import QuantityFactory, StencilFactory from ndsl.constants import X_DIM, Y_DIM @@ -29,7 +37,7 @@ def edge_pe_update( region[local_is - 2 : local_ie + 3, local_je + 1 : local_je + 3], ): pe = pe + delp[0, 0, -1] - pk3 = pe ** akap + pk3 = exp(akap * log(pe)) class PK3Halo: From 648e548d42d5146297edbb47683af14fb8646787 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:34:50 -0400 Subject: [PATCH 188/252] Make sure to use `dz_min` as parameters to allow for better GEOS integration --- pyFV3/stencils/updatedzc.py | 30 ++++++++++++++++++------------ pyFV3/stencils/updatedzd.py | 28 +++++++++++++++++++++------- 2 files changed, 39 insertions(+), 19 deletions(-) diff --git a/pyFV3/stencils/updatedzc.py b/pyFV3/stencils/updatedzc.py index 3b3a76b8..dff2d5bd 100644 --- a/pyFV3/stencils/updatedzc.py +++ b/pyFV3/stencils/updatedzc.py @@ -1,16 +1,12 @@ import gt4py.cartesian.gtscript as gtscript from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval -import ndsl.constants as constants from ndsl import Quantity, QuantityFactory, StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK from ndsl.stencils import corners -DZ_MIN = constants.DZ_MIN - - @gtscript.function def p_weighted_average_top(vel, dp0): # TODO: ratio is a constant, where should this be placed? @@ -69,23 +65,25 @@ def update_dz_c( ws: FloatFieldIJ, *, dt: Float, + dz_min: Float, ): """ Step dz forward on c-grid Eusures gz is monotonically increasing in z at the end Args: - dp_ref: - zs: + dp_ref: layer thickness in Pa + zs: surface height in m area: - ut: - vt: - gz: + ut: horizontal wind (TODO: covariant or contravariant?) + vt: horizontal wind (TODO: covariant or contravariant?) + gz: geopotential height on model interfaces gz_x: gz with corners copied to perform derivatives in x-direction gz_y: gz with corners copied to perform derivatives in y-direction ws: lagrangian (parcel-following) surface vertical wind implied by lowest-level gz change note that a parcel moving horizontally across terrain will be moving in the vertical (eqn 5.5 in documentation) - dt: + dt: timestep over which to evolve the geopotential height, in seconds + dz_min: Controls minimum thickness in NH solver """ # there's some complexity due to gz being defined on interfaces @@ -112,7 +110,7 @@ def update_dz_c( rdt = 1.0 / dt ws = (zs - gz) * rdt with computation(BACKWARD), interval(0, -1): - gz_kp1 = gz[0, 0, 1] + DZ_MIN + gz_kp1 = gz[0, 0, 1] + dz_min gz = gz if gz > gz_kp1 else gz_kp1 @@ -123,11 +121,18 @@ def __init__( quantity_factory: QuantityFactory, area: Quantity, dp_ref: Quantity, - grid_type, + grid_type: int, + dz_min: Float, ): + """ + Args: + dz_min: controls minimum thickness in NH solver + """ + grid_indexing = stencil_factory.grid_indexing self._area = area self._grid_type = grid_type + self._dz_min = dz_min # TODO: this is needed because GridData.dp_ref does not have access # to a QuantityFactory, we should add a way to perform operations on # Quantity and persist the QuantityFactory choices @@ -221,4 +226,5 @@ def __call__( self._gz_y, ws, dt=dt, + dz_min=self._dz_min, ) diff --git a/pyFV3/stencils/updatedzd.py b/pyFV3/stencils/updatedzd.py index abbdbeb7..385aea3f 100644 --- a/pyFV3/stencils/updatedzd.py +++ b/pyFV3/stencils/updatedzd.py @@ -3,7 +3,6 @@ import gt4py.cartesian.gtscript as gtscript from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval -import ndsl.constants as constants from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate from ndsl.constants import ( X_DIM, @@ -19,9 +18,6 @@ from pyFV3.stencils.fvtp2d import FiniteVolumeTransport -DZ_MIN = constants.DZ_MIN - - @gtscript.function def _apply_height_advective_flux( height: FloatField, @@ -77,6 +73,7 @@ def apply_height_fluxes( surface_height: FloatFieldIJ, ws: FloatFieldIJ, dt: Float, + dz_min: Float, ): """ Apply all computed fluxes to height profile. @@ -100,6 +97,7 @@ def apply_height_fluxes( surface_height (in): surface height ws (out): vertical velocity of the lowest level (to keep it at the surface) dt (in): acoustic timestep (seconds) + dz_min(in): controls minimum thickness in NH solver Grid variable inputs: area """ @@ -120,7 +118,7 @@ def apply_height_fluxes( ws = (surface_height - height) / dt with interval(0, -1): # ensure layer thickness exceeds minimum - other = height[0, 0, 1] + DZ_MIN + other = height[0, 0, 1] + dz_min height = height if height > other else other @@ -219,8 +217,20 @@ def __init__( grid_data: GridData, grid_type: int, hord_tm: int, + dz_min: Float, column_namelist, ): + """ + Args: + stencil_factory + quantity_factory + damping_coefficients + grid_data + grid_type + hord_tm + dz_min (in): controls minimum thickness in NH solver + column_namelist + """ orchestrate( obj=self, config=stencil_factory.config.dace_config, @@ -229,8 +239,11 @@ def __init__( self.grid_indexing = grid_indexing self._area = grid_data.area self._column_namelist = column_namelist - if any(column_namelist["damp_vt"].view[:] <= 1e-5): - raise NotImplementedError("damp <= 1e-5 in column_namelist is untested") + if any(column_namelist["damp_vt"].view[:] <= Float(1e-5)): + raise NotImplementedError( + "damp <= 1e-5 in column_namelist is not implemented" + ) + self._dz_min = dz_min self._dp_ref = grid_data.dp_ref self._allocate_temporary_storages(quantity_factory) self._gk, self._beta, self._gamma = cubic_spline_interpolation_constants( @@ -383,4 +396,5 @@ def __call__( surface_height, ws, dt, + self._dz_min, ) From b494884a29a61849a91eb4d0af9de3758f83f1ca Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 25 Mar 2025 14:37:28 -0400 Subject: [PATCH 189/252] Match Fortran great circle distance Fix doublu periodic order of op to match Fortran --- pyFV3/stencils/a2b_ord4.py | 50 +++++++++++++++++++++++++------------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/pyFV3/stencils/a2b_ord4.py b/pyFV3/stencils/a2b_ord4.py index 514d7a56..0744796f 100644 --- a/pyFV3/stencils/a2b_ord4.py +++ b/pyFV3/stencils/a2b_ord4.py @@ -1,3 +1,4 @@ +import xarray as xr import gt4py.cartesian.gtscript as gtscript from gt4py.cartesian.gtscript import ( PARALLEL, @@ -9,11 +10,18 @@ region, sin, sqrt, + f64, ) from ndsl import GridIndexing, QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM -from ndsl.dsl.typing import Float, FloatField, FloatFieldI, FloatFieldIJ +from ndsl.dsl.typing import ( + Float, + FloatField, + FloatFieldI64, + FloatFieldIJ, + FloatFieldIJ64, +) from ndsl.grid import GridData from ndsl.stencils.basic_operations import copy_defn @@ -24,18 +32,26 @@ d1 = Float(0.375) d2 = Float(-1.0) / Float(24.0) # PPM volume mean form -b1 = Float(7.0) / Float(12.0) +b1 = Float(7.0) / Float(12.0) # 0.58333333 b2 = Float(-1.0) / Float(12.0) # 4-pt Lagrange interpolation -a1 = Float(9.0) / Float(16.0) -a2 = Float(-1.0) / Float(16.0) +a1 = Float(0.5625) # 9/16 +a2 = Float(-0.0625) # -1/16 + +r3 = Float(1.0 / 3.0) @gtscript.function def great_circle_dist(p1a, p1b, p2a, p2b): - tb = sin((p1b - p2b) / 2.0) ** 2.0 - ta = sin((p1a - p2a) / 2.0) ** 2.0 - return asin(sqrt(tb + cos(p1b) * cos(p2b) * ta)) * 2.0 + return ( + asin( + sqrt( + sin((p1b - p2b) / 2.0) ** 2.0 + + cos(p1b) * cos(p2b) * sin((p1a - p2a) / 2.0) ** 2.0 + ) + ) + * 2.0 + ) @gtscript.function @@ -105,7 +121,7 @@ def _sw_corner( qin[1, -2, 0], ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + qout = (ec1 + ec2 + ec3) * r3 tmp_qout_edges = qout @@ -159,7 +175,7 @@ def _nw_corner( qin[0, 0, 0], qin[1, 1, 0], ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + qout = (ec1 + ec2 + ec3) * r3 tmp_qout_edges = qout @@ -213,7 +229,7 @@ def _ne_corner( qin[-1, 0, 0], qin[-2, 1, 0], ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + qout = (ec1 + ec2 + ec3) * r3 tmp_qout_edges = qout @@ -267,7 +283,7 @@ def _se_corner( qin[0, 0, 0], qin[1, 1, 0], ) - qout = (ec1 + ec2 + ec3) * (1.0 / 3.0) + qout = (ec1 + ec2 + ec3) * r3 tmp_qout_edges = qout @@ -284,7 +300,7 @@ def lagrange_x_func(qy): def qout_x_edge( qin: FloatField, dxa: FloatFieldIJ, - edge_w: FloatFieldIJ, + edge_w: FloatFieldIJ64, qout: FloatField, tmp_qout_edges: FloatField, ): @@ -305,7 +321,7 @@ def qout_x_edge( def qout_y_edge( qin: FloatField, dya: FloatFieldIJ, - edge_s: FloatFieldI, + edge_s: FloatFieldI64, qout: FloatField, tmp_qout_edges: FloatField, ): @@ -510,11 +526,11 @@ def doubly_periodic_a2b_ord4(qin): Grid conversion is much simpler on a doubly-periodic, orthogonal grid so we can bypass most of the above code """ - qx = b1 * (qin[-1, 0, 0] + qin) + b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) - qy = b1 * (qin[0, -1, 0] + qin) + b2 * (qin[0, -2, 0] + qin[0, 1, 0]) + qx = b2 * (qin[-2, 0, 0] + qin[1, 0, 0]) * b1 * (qin[-1, 0, 0] + qin) + qy = b2 * (qin[0, -2, 0] + qin[0, 1, 0]) * b1 * (qin[0, -1, 0] + qin) qout = 0.5 * ( - a1 * (qx[0, -1, 0] + qx + qy[-1, 0, 0] + qy) - + a2 * (qx[0, -2, 0] + qx[0, 1, 0] + qy[-2, 0, 0] + qy[1, 0, 0]) + a2 * (qx[0, -2, 0] + qx[0, 1, 0] + qy[-2, 0, 0] + qy[1, 0, 0]) + + a1 * (qx[0, -1, 0] + qx + qy[-1, 0, 0] + qy) ) return qout From 327d050784d408d5c16de4d52208e26cef657382 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 26 Mar 2025 11:51:19 -0400 Subject: [PATCH 190/252] Pass `dz_min` from `dyn_core` to the updater --- pyFV3/stencils/dyn_core.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 2e235f28..4bf79c86 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -484,6 +484,7 @@ def __init__( grid_data=grid_data, grid_type=grid_type, hord_tm=config.hord_tm, + dz_min=constants.DZ_MIN, column_namelist=column_namelist, ) self.vertical_solver = NonhydrostaticVerticalSolver( @@ -551,6 +552,7 @@ def __init__( area=grid_data.area, dp_ref=grid_data.dp_ref, grid_type=config.grid_type, + dz_min=constants.DZ_MIN, ) ) From 38b6581948a3348604831b6f1da2e21fe763fcd9 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 27 Mar 2025 14:15:11 -0400 Subject: [PATCH 191/252] Update `int` dtypes to precision capable ones --- pyFV3/stencils/fillz.py | 4 ++-- pyFV3/stencils/map_single.py | 32 ++++++++++++++++---------------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index bc8ef7c7..0426956e 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -5,7 +5,7 @@ import ndsl.dsl.gt4py_utils as utils from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntFieldIJ +from ndsl.dsl.typing import Float, Int, FloatField, FloatFieldIJ, IntFieldIJ from pyFV3.tracers import Tracers @@ -133,7 +133,7 @@ def __init__( # Setting initial value of upper_fix to zero is only needed for validation. # The values in the compute domain are set to zero in the stencil. - self._zfix = quantity_factory.zeros([X_DIM, Y_DIM], units="unknown", dtype=int) + self._zfix = quantity_factory.zeros([X_DIM, Y_DIM], units="unknown", dtype=Int) self._sum0 = quantity_factory.zeros( [X_DIM, Y_DIM], units="unknown", diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index 2f4c5edf..32fde431 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -257,7 +257,7 @@ def make_quantity(): units="unknown", dtype=Float, ) - self._lev = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=int) + self._lev = quantity_factory.zeros([X_DIM, Y_DIM], units="", dtype=Int) self._copy_stencil = stencil_factory.from_dims_halo( copy_defn, @@ -290,13 +290,13 @@ def make_quantity(): self._INDEX_LM1 = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="", - dtype=int, + dtype=Int, ) self._INDEX_LP0 = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="", - dtype=int, + dtype=Int, ) self._km = grid_indexing.domain[2] self._not_exit_loop = quantity_factory.zeros( @@ -369,17 +369,17 @@ def __call__( ) else: self._lagrangian_contributions_interp( - self._km, - self._not_exit_loop, - self._INDEX_LM1, - self._INDEX_LP0, - q1, - pe1, - pe2, - self._q4_1, - self._q4_2, - self._q4_3, - self._q4_4, - self._dp1, - self._lev, + km=self._km, + not_exit_loop=self._not_exit_loop, + INDEX_LM1=self._INDEX_LM1, + INDEX_LP0=self._INDEX_LP0, + q=q1, + pe1=pe1, + pe2=pe2, + q4_1=self._q4_1, + q4_2=self._q4_2, + q4_3=self._q4_3, + q4_4=self._q4_4, + dp1=self._dp1, + lev=self._lev, ) From 54d5619ee4bc6e173aee9f34633e7bd308e0a746 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 27 Mar 2025 14:39:26 -0400 Subject: [PATCH 192/252] Fix remap after tracer merge --- pyFV3/stencils/map_single.py | 1 + pyFV3/stencils/remapping_GEOS.py | 80 ++++++++++++++++---------------- 2 files changed, 40 insertions(+), 41 deletions(-) diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index 32fde431..fc09e817 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -5,6 +5,7 @@ from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import ( # noqa: F401 + Int, BoolFieldIJ, Float, FloatField, diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 3d6c3629..fcbd9e41 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -85,8 +85,7 @@ def __init__( self._adiabatic = adiabatic self.kmp = grid_indexing.domain[2] - 1 for k in range(pfull.shape[0]): - # if pfull.view[k] > 10.0e2: - if pfull[k] > 10.0e2: + if pfull.view[k] > 10.0e2: self.kmp = k break # do_omega = hydrostatic and last_step # TODO pull into inputs @@ -208,11 +207,10 @@ def __init__( self._mapn_tracer = MapNTracer( stencil_factory, quantity_factory, - abs(config.kord_tr), - self._nq, + kord=abs(config.kord_tr), fill=config.fill, tracers=tracers, - # exclude_tracers=[], + exclude_tracers=[], ) self._map_single_w = MapSingle( @@ -413,12 +411,12 @@ def __call__( # Build remapping profiles self._init_pe(pe, self._pe1, self._pe2, ptop) self._moist_cv_pt_pressure( - qvapor=tracers["qvapor"], - qliquid=tracers["qliquid"], - qrain=tracers["qrain"], - qsnow=tracers["qsnow"], - qice=tracers["qice"], - qgraupel=tracers["qgraupel"], + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], q_con=q_con, pt=pt, cappa=cappa, @@ -499,12 +497,12 @@ def __call__( ) self._moist_cv_pkz( - qvapor=tracers["qvapor"], - qliquid=tracers["qliquid"], - qrain=tracers["qrain"], - qsnow=tracers["qsnow"], - qice=tracers["qice"], - qgraupel=tracers["qgraupel"], + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], pkz=pkz, pt=pt, cappa=cappa, @@ -517,12 +515,12 @@ def __call__( if last_step: if consv_te > CONSV_MIN: self._moist_cv_te( - qvapor=tracers["qvapor"], - qliquid=tracers["qliquid"], - qrain=tracers["qrain"], - qsnow=tracers["qsnow"], - qice=tracers["qice"], - qgraupel=tracers["qgraupel"], + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], u=u, v=v, w=w, @@ -570,13 +568,13 @@ def __call__( fast_mp_consv = consv_te > CONSV_MIN self._saturation_adjustment( dp1, - tracers["qvapor"], - tracers["qliquid"], - tracers["qice"], - tracers["qrain"], - tracers["qsnow"], - tracers["qgraupel"], - tracers["qcld"], + tracers["vapor"], + tracers["liquid"], + tracers["ice"], + tracers["rain"], + tracers["snow"], + tracers["graupel"], + tracers["cloud"], hs, peln, delp, @@ -598,12 +596,12 @@ def __call__( # to the physics, but if we're staying in dynamics we need # to keep it as the virtual potential temperature self._moist_cv_last_step_stencil( - qvapor=tracers["qvapor"], - qliquid=tracers["qliquid"], - qrain=tracers["qrain"], - qsnow=tracers["qsnow"], - qice=tracers["qice"], - qgraupel=tracers["qgraupel"], + qvapor=tracers["vapor"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], pt=pt, pkz=pkz, dtmp=Float(dtmp), @@ -611,11 +609,11 @@ def __call__( ) self._fill_cond( q_con=q_con, - qliquid=tracers["qliquid"], - qrain=tracers["qrain"], - qsnow=tracers["qsnow"], - qice=tracers["qice"], - qgraupel=tracers["qgraupel"], + qliquid=tracers["liquid"], + qrain=tracers["rain"], + qsnow=tracers["snow"], + qice=tracers["ice"], + qgraupel=tracers["graupel"], ) else: # converts virtual temperature back to virtual potential temperature From b5626b7b03c9eda2a5be58a3303547c2e029b128 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 27 Mar 2025 14:40:57 -0400 Subject: [PATCH 193/252] Proper type for da_min --- pyFV3/stencils/fv_dynamics.py | 4 ++-- pyFV3/testing/translate_fvdynamics.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 4f2df9ee..ed95d93b 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -521,8 +521,8 @@ def __init__( # See divergence_damping.py, _get_da_min for explanation of this function @dace_inhibitor - def _get_da_min(self) -> Float: # type: ignore - return Float(self._da_min) + def _get_da_min(self) -> NDSL_64BIT_FLOAT_TYPE: # type: ignore + return self._da_min def step_dynamics( self, diff --git a/pyFV3/testing/translate_fvdynamics.py b/pyFV3/testing/translate_fvdynamics.py index af4f5ac0..8572a279 100644 --- a/pyFV3/testing/translate_fvdynamics.py +++ b/pyFV3/testing/translate_fvdynamics.py @@ -367,7 +367,8 @@ def compute_parallel(self, inputs, communicator): config=DynamicalCoreConfig.from_namelist(self.namelist), phis=state.phis, state=state, - timestep=timedelta(seconds=inputs["bdt"]), + exclude_tracers=["cloud"], + timestep=timedelta(seconds=float(inputs["bdt"])), ) self.dycore.step_dynamics(state, NullTimer()) outputs = self.outputs_from_state(state) From 593da4a55e268c85486e4a29092c78667b4308ec Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 28 Mar 2025 09:34:55 -0400 Subject: [PATCH 194/252] Fix bad merge --- pyFV3/stencils/fillz.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index 212b08ee..9d6088be 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -4,7 +4,7 @@ from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, IntFieldIJ +from ndsl.dsl.typing import Float, Int, FloatField, FloatFieldIJ, IntFieldIJ from pyFV3.tracers import Tracers From 9a98368161c97d3d869e4c3361885974f39637ef Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 14 Apr 2025 09:01:12 -0400 Subject: [PATCH 195/252] Heat source from vorticity on the correct grid --- pyFV3/stencils/d_sw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/d_sw.py b/pyFV3/stencils/d_sw.py index d70276e7..24f73591 100644 --- a/pyFV3/stencils/d_sw.py +++ b/pyFV3/stencils/d_sw.py @@ -999,7 +999,7 @@ def make_quantity(): self._heat_source_from_vorticity_damping_stencil = ( stencil_factory.from_dims_halo( func=heat_source_from_vorticity_damping, - compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], + compute_dims=[X_DIM, Y_DIM, Z_DIM], externals={ "do_stochastic_ke_backscatter": config.do_skeb, "d_con": config.d_con, From 23f986275f6959b4d8101484befb3fd7af6fcdda Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 14 Apr 2025 09:02:32 -0400 Subject: [PATCH 196/252] Insure abs on proper precision --- pyFV3/stencils/dyn_core.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 4bf79c86..677b548b 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -923,8 +923,9 @@ def __call__( # to damp self._hyperdiffusion(self._heat_source, cd) if not self.config.hydrostatic: - delt_time_factor = abs( - dt_acoustic_substep * Float(self.config.delt_max) + delt_time_factor = np.abs( + dt_acoustic_substep * Float(self.config.delt_max), + dtype=Float, ) # TODO: it looks like state.pkz is being used as a temporary here, # and overwritten at the start of remapping. See if we can make it From 8e59620ad9b4ed1289cd8933e29749952afe968b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 14 Apr 2025 16:03:45 -0400 Subject: [PATCH 197/252] Add `dz_min` as a namelist input --- pyFV3/_config.py | 3 +++ pyFV3/stencils/dyn_core.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/pyFV3/_config.py b/pyFV3/_config.py index ce7fb2a6..fcb9a630 100644 --- a/pyFV3/_config.py +++ b/pyFV3/_config.py @@ -112,6 +112,7 @@ class AcousticDynamicsConfig: """ mainly for backwards compatibility, not really used anymore """ + dz_min: float riemann: RiemannConfig d_grid_shallow_water: DGridShallowWaterLagrangianDynamicsConfig @@ -274,6 +275,7 @@ class DynamicalCoreConfig: nf_omega: int = NamelistDefaults.nf_omega fv_sg_adj: int = NamelistDefaults.fv_sg_adj n_sponge: int = NamelistDefaults.n_sponge + dz_min: float = Namelist.dz_min namelist_override: Optional[str] = None def __post_init__(self): @@ -488,6 +490,7 @@ def acoustic_dynamics(self) -> AcousticDynamicsConfig: use_old_omega=self.use_old_omega, riemann=self.riemann, d_grid_shallow_water=self.d_grid_shallow_water, + dz_min=self.dz_min, ) @property diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 530d7f26..524b3696 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -497,7 +497,7 @@ def __init__( dtype=Float, ) self._zs.data[:] = self._zs.np.asarray( - phis.data / constants.GRAV, dtype=self._zs.data.dtype + phis.data * constants.RGRAV, dtype=self._zs.data.dtype ) self.update_height_on_d_grid = updatedzd.UpdateHeightOnDGrid( @@ -508,6 +508,7 @@ def __init__( grid_type=grid_type, hord_tm=config.hord_tm, column_namelist=column_namelist, + dz_min=Float(config.dz_min), ) self.vertical_solver = NonhydrostaticVerticalSolver( stencil_factory, @@ -574,6 +575,7 @@ def __init__( area=grid_data.area, dp_ref=grid_data.dp_ref, grid_type=config.grid_type, + dz_min=Float(config.dz_min), ) ) From c7f56bcbd6896b63f6431dfdb7863788e9418bc2 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 14 Apr 2025 16:05:37 -0400 Subject: [PATCH 198/252] Hotfix for dyncore --- pyFV3/stencils/dyn_core.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 6326d02a..2c2678e9 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -496,7 +496,6 @@ def __init__( grid_data=grid_data, grid_type=grid_type, hord_tm=config.hord_tm, - dz_min=constants.DZ_MIN, column_namelist=column_namelist, dz_min=Float(config.dz_min), ) From 9624c3f05c31d64e8d0a9c38da9f2cffe3f838e7 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 15 Apr 2025 09:16:18 -0400 Subject: [PATCH 199/252] Add defaiult to `dz_min` --- pyFV3/_config.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyFV3/_config.py b/pyFV3/_config.py index fcb9a630..dcb2e8de 100644 --- a/pyFV3/_config.py +++ b/pyFV3/_config.py @@ -275,7 +275,8 @@ class DynamicalCoreConfig: nf_omega: int = NamelistDefaults.nf_omega fv_sg_adj: int = NamelistDefaults.fv_sg_adj n_sponge: int = NamelistDefaults.n_sponge - dz_min: float = Namelist.dz_min + dz_min: float = 2.0 + """Controls minimum thickness in NH solver""" namelist_override: Optional[str] = None def __post_init__(self): From d8c688183d4a68e0e2842bcf80493b235a80e413 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 18 Apr 2025 08:02:35 -0400 Subject: [PATCH 200/252] Fix proper use of `gz` corner extension --- pyFV3/stencils/updatedzc.py | 62 +++++++++++++++++++++++++------------ 1 file changed, 42 insertions(+), 20 deletions(-) diff --git a/pyFV3/stencils/updatedzc.py b/pyFV3/stencils/updatedzc.py index 3b3a76b8..b88740b1 100644 --- a/pyFV3/stencils/updatedzc.py +++ b/pyFV3/stencils/updatedzc.py @@ -57,6 +57,11 @@ def double_copy(q_in: FloatField, copy_1: FloatField, copy_2: FloatField): copy_2 = q_in +def copy(q_in: FloatField, q_copy: FloatField): + with computation(PARALLEL), interval(...): + q_copy = q_in + + def update_dz_c( dp_ref: FloatFieldK, zs: FloatFieldIJ, @@ -66,6 +71,7 @@ def update_dz_c( gz: FloatField, gz_x: FloatField, gz_y: FloatField, + gz_filled: FloatField, ws: FloatFieldIJ, *, dt: Float, @@ -74,18 +80,19 @@ def update_dz_c( Step dz forward on c-grid Eusures gz is monotonically increasing in z at the end Args: - dp_ref: - zs: - area: - ut: - vt: - gz: - gz_x: gz with corners copied to perform derivatives in x-direction - gz_y: gz with corners copied to perform derivatives in y-direction - ws: lagrangian (parcel-following) surface vertical wind implied by + dp_ref(in): layer thickness in Pa + zs(in): surface height in m + area(in): + ut(in): horizontal wind (TODO: covariant or contravariant?) + vt(in): horizontal wind (TODO: covariant or contravariant?) + gz(inout): geopotential height on model interfaces + gz_x(in): gz with corners copied to perform derivatives in x-direction + gz_y(in): gz with corners copied to perform derivatives in y-direction + ws(out): lagrangian (parcel-following) surface vertical wind implied by lowest-level gz change note that a parcel moving horizontally across terrain will be moving in the vertical (eqn 5.5 in documentation) - dt: + dt(in): timestep over which to evolve the geopotential height, in seconds + dz_min(in): Controls minimum thickness in NH solver """ # there's some complexity due to gz being defined on interfaces @@ -105,7 +112,7 @@ def update_dz_c( # xfx/yfx are now ut/vt interpolated to layer interfaces with computation(PARALLEL), interval(...): fx, fy = xy_flux(gz_x, gz_y, xfx, yfx) - gz = (gz * area + (fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) / ( + gz = (gz_filled * area + (fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) / ( area + (xfx - xfx[1, 0, 0]) + (yfx - yfx[0, 1, 0]) ) with computation(FORWARD), interval(-1, None): @@ -150,6 +157,11 @@ def __init__( units="m**2/s**2", dtype=Float, ) + self._gz_filled = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="m**2/s**2", + dtype=Float, + ) full_origin = grid_indexing.origin_full() full_domain = grid_indexing.domain_full(add=(0, 0, 1)) self._double_copy_stencil = stencil_factory.from_origin_domain( @@ -157,6 +169,11 @@ def __init__( origin=full_origin, domain=full_domain, ) + self._copy_stencil = stencil_factory.from_origin_domain( + copy, + origin=full_origin, + domain=full_domain, + ) ax_offsets = grid_indexing.axis_offsets(full_origin, full_domain) @@ -203,6 +220,10 @@ def __call__( # TODO: is this advecting gz, and if so can we name it that? # Can we reduce duplication of advection logic with other stencils? + self._copy_stencil(gz, self._gz_filled) + self._fill_corners_x_stencil(self._gz_filled, self._gz_filled) + self._fill_corners_y_stencil(self._gz_filled, self._gz_filled) + self._double_copy_stencil(gz, self._gz_x, self._gz_y) # TODO(eddied): We pass the same fields 2x to avoid GTC validation errors @@ -211,14 +232,15 @@ def __call__( self._fill_corners_y_stencil(self._gz_y, self._gz_y) self._update_dz_c( - self._dp_ref, - zs, - self._area, - ut, - vt, - gz, - self._gz_x, - self._gz_y, - ws, + dp_ref=self._dp_ref, + zs=zs, + area=self._area, + ut=ut, + vt=vt, + gz=gz, + gz_x=self._gz_x, + gz_y=self._gz_y, + gz_filled=self._gz_filled, + ws=ws, dt=dt, ) From 53bbfd8295b953585dc56ee26d1ff211d07533ef Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 18 Apr 2025 08:03:32 -0400 Subject: [PATCH 201/252] Match divg_d Fortran computation for divergence_corner --- pyFV3/stencils/c_sw.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pyFV3/stencils/c_sw.py b/pyFV3/stencils/c_sw.py index b43c4c0f..39b58969 100644 --- a/pyFV3/stencils/c_sw.py +++ b/pyFV3/stencils/c_sw.py @@ -161,7 +161,8 @@ def divergence_corner( ) vf0 = v * dxc * 0.5 * (sin_sg3[-1, 0] + sin_sg1) uf0 = u * dyc * 0.5 * (sin_sg4[0, -1] + sin_sg2) - divg_d = (-vf0 + uf1 - uf0) * rarea_c + divg_d = vf1 - vf0 + uf1 - uf0 + divg_d = rarea_c * (divg_d - vf1) with horizontal(region[i_end + 1, j_end + 1], region[i_start, j_end + 1]): vf1 = ( @@ -171,8 +172,8 @@ def divergence_corner( u[-1, 0, 0] * dyc[-1, 0] * 0.5 * (sin_sg4[-1, -1] + sin_sg2[-1, 0]) ) uf0 = u * dyc * 0.5 * (sin_sg4[0, -1] + sin_sg2) - divg_d = (vf1 + uf1 - uf0) * rarea_c - + divg_d = vf1 - vf0 + uf1 - uf0 + divg_d = rarea_c * (divg_d + vf0) # --------- From f89fadfac381bf55fa450f9d59d42d45436134eb Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 18 Apr 2025 08:04:03 -0400 Subject: [PATCH 202/252] Initialize `gz` to HUGE_R --- pyFV3/stencils/dyn_core.py | 98 ++++++++++++++++++++------------------ 1 file changed, 52 insertions(+), 46 deletions(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 677b548b..07fd5f27 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -460,6 +460,7 @@ def __init__( if not config.hydrostatic: self._pk3.data[:] = HUGE_R + self._gz.data[:] = HUGE_R column_namelist = d_sw.get_column_namelist( config.d_grid_shallow_water, quantity_factory=quantity_factory @@ -752,7 +753,12 @@ def __call__( ) if not self.config.hydrostatic: self.update_geopotential_height_on_c_grid( - self._zs, self._ut, self._vt, self._gz, self._ws3, dt2 + zs=self._zs, + ut=self._ut, + vt=self._vt, + gz=self._gz, + ws=self._ws3, + dt=dt2, ) # TODO (floriand): Due to DaCe VRAM pooling creating a memory # leak with the usage pattern of those two fields @@ -762,28 +768,28 @@ def __call__( # DaCe has already a fix on their side and it awaits release # issue self.vertical_solver_cgrid( - dt2, - self.cappa, - self._ptop, - state.phis, - self._ws3, - self.cgrid_shallow_water_lagrangian_dynamics.ptc, - state.q_con, - self.cgrid_shallow_water_lagrangian_dynamics.delpc, - self._gz, - self._pkc, - state.omga, + dt2=dt2, + cappa=self.cappa, + ptop=self._ptop, + hs=state.phis, + ws=self._ws3, + ptc=self.cgrid_shallow_water_lagrangian_dynamics.ptc, + q_con=state.q_con, + delpc=self.cgrid_shallow_water_lagrangian_dynamics.delpc, + gz=self._gz, + pef=self._pkc, + w3=state.omga, ) self._p_grad_c( - self.grid_data.rdxc, - self.grid_data.rdyc, - state.uc, - state.vc, - self.cgrid_shallow_water_lagrangian_dynamics.delpc, - self._pkc, - self._gz, - dt2, + rdxc=self.grid_data.rdxc, + rdyc=self.grid_data.rdyc, + uc=state.uc, + vc=state.vc, + delpc=self.cgrid_shallow_water_lagrangian_dynamics.delpc, + pkc=self._pkc, + gz=self._gz, + dt2=dt2, ) self._halo_updaters.uc__vc.start() if self.config.nord > 0: @@ -842,23 +848,23 @@ def __call__( dt=dt_acoustic_substep, ) self.vertical_solver( - remap_step, - dt_acoustic_substep, - self.cappa, - self._ptop, - self._zs, - self._wsd, - state.delz, - state.q_con, - state.delp, - state.pt, - self._zh, - state.pe, - self._pkc, - self._pk3, - state.pk, - state.peln, - state.w, + last_call=remap_step, + dt=dt_acoustic_substep, + cappa=self.cappa, + ptop=self._ptop, + zs=self._zs, + ws=self._wsd, + delz=state.delz, + q_con=state.q_con, + delp=state.delp, + pt=state.pt, + zh=self._zh, + p=state.pe, + ppe=self._pkc, + pk3=self._pk3, + pk=state.pk, + log_p_interface=state.peln, + w=state.w, ) self._halo_updaters.zh.start() @@ -881,15 +887,15 @@ def __call__( self._halo_updaters.pkc.wait() self.nonhydrostatic_pressure_gradient( - state.u, - state.v, - self._pkc, - self._gz, - self._pk3, - state.delp, - dt_acoustic_substep, - self._ptop, - self._akap, + u=state.u, + v=state.v, + pp=self._pkc, + gz=self._gz, + pk3=self._pk3, + delp=state.delp, + dt=dt_acoustic_substep, + ptop=self._ptop, + akap=self._akap, ) if self.config.rf_fast: From f63892702086d66844ba4e951ac455de6a1efdbd Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 18 Apr 2025 08:05:32 -0400 Subject: [PATCH 203/252] Fix translate test halo reach & pass `dz_min` --- .../translate/translate_updatedzc.py | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/tests/savepoint/translate/translate_updatedzc.py b/tests/savepoint/translate/translate_updatedzc.py index 9553a512..ed8d82bc 100644 --- a/tests/savepoint/translate/translate_updatedzc.py +++ b/tests/savepoint/translate/translate_updatedzc.py @@ -16,15 +16,17 @@ def __init__( ): super().__init__(grid, namelist, stencil_factory) self.stencil_factory = stencil_factory - update_gz_on_c_grid = UpdateGeopotentialHeightOnCGrid( - self.stencil_factory, - quantity_factory=self.grid.quantity_factory, - area=grid.grid_data.area, - dp_ref=grid.grid_data.dp_ref, - grid_type=namelist.grid_type, - ) def compute(**kwargs): + update_gz_on_c_grid = UpdateGeopotentialHeightOnCGrid( + self.stencil_factory, + quantity_factory=self.grid.quantity_factory, + area=grid.grid_data.area, + dp_ref=grid.grid_data.dp_ref, + grid_type=namelist.grid_type, + dz_min=kwargs.pop("dz_min"), + ) + kwargs["dt"] = kwargs.pop("dt2") update_gz_on_c_grid(**kwargs) @@ -36,7 +38,7 @@ def compute(**kwargs): "gz": {}, "ws": {}, } - self.in_vars["parameters"] = ["dt2"] + self.in_vars["parameters"] = ["dt2", "dz_min"] self.out_vars = { "gz": grid.default_buffer_k_dict(), "ws": {"kstart": -1, "kend": None}, @@ -44,12 +46,12 @@ def compute(**kwargs): self._subset = get_subset_func( self.grid.grid_indexing, dims=[X_DIM, Y_DIM, Z_DIM], - n_halo=((0, 0), (0, 0)), + n_halo=((3, 3), (3, 3)), ) self._subset_2d = get_subset_func( self.grid.grid_indexing, dims=[X_DIM, Y_DIM], - n_halo=((0, 0), (0, 0)), + n_halo=((3, 3), (3, 3)), ) def compute(self, inputs): From e58d4169c011b30f7ff7f162496056f2869dd880 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 18 Apr 2025 08:06:20 -0400 Subject: [PATCH 204/252] Pass `dz_min` in translate test --- tests/savepoint/translate/translate_updatedzd.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/savepoint/translate/translate_updatedzd.py b/tests/savepoint/translate/translate_updatedzd.py index 75800895..a340cfc0 100644 --- a/tests/savepoint/translate/translate_updatedzd.py +++ b/tests/savepoint/translate/translate_updatedzd.py @@ -33,7 +33,7 @@ def __init__( self.in_vars["data_vars"]["height"]["serialname"] = "zh" self.in_vars["data_vars"]["ws"]["serialname"] = "wsd" - self.in_vars["parameters"] = ["dt"] + self.in_vars["parameters"] = ["dt", "dz_min"] out_vars = [ "height", "courant_number_x", @@ -66,6 +66,7 @@ def compute(self, inputs): self.grid.grid_data, self.grid.grid_type, self.namelist.hord_tm, + dz_min=inputs.pop("dz_min"), column_namelist=d_sw.get_column_namelist( self.namelist, quantity_factory=self.grid.quantity_factory ), From 8cf966b335584a2d9be8880b5951599a3eff7f37 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 30 Apr 2025 15:24:59 -0400 Subject: [PATCH 205/252] Clean up Translate test --- tests/savepoint/translate/__init__.py | 1 - .../translate/translate_remapping_GEOS.py | 1032 ++++------------- .../translate/translate_remapping_GEOS_v2.py | 585 ---------- 3 files changed, 202 insertions(+), 1416 deletions(-) delete mode 100644 tests/savepoint/translate/translate_remapping_GEOS_v2.py diff --git a/tests/savepoint/translate/__init__.py b/tests/savepoint/translate/__init__.py index 55412ce6..1e11f4af 100644 --- a/tests/savepoint/translate/__init__.py +++ b/tests/savepoint/translate/__init__.py @@ -92,7 +92,6 @@ from .translate_ray_fast import TranslateRay_Fast from .translate_remapping import TranslateRemapping from .translate_remapping_GEOS import TranslateRemapping_GEOS -from .translate_remapping_GEOS_v2 import TranslateRemapping_GEOS_v2 from .translate_riem_solver3 import TranslateRiem_Solver3 from .translate_riem_solver_c import TranslateRiem_Solver_C from .translate_satadjust3d import TranslateSatAdjust3d diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index c92f519f..ed2a8947 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,9 +1,7 @@ from types import SimpleNamespace -from ndsl import Namelist, StencilFactory +from ndsl import Namelist, StencilFactory, Quantity from ndsl.constants import ( - CV_AIR, - GRAV, X_DIM, X_INTERFACE_DIM, Y_DIM, @@ -14,78 +12,17 @@ from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig - -# from pyFV3._config import RemappingConfig -from pyFV3.stencils import moist_cv -from pyFV3.stencils.map_single import MapSingle -from pyFV3.stencils.mapn_tracer import MapNTracer -from pyFV3.stencils.mpp_global_sum import mpp_global_sum -from pyFV3.stencils.remapping import ( - CONSV_MIN, - init_pe, - moist_cv_pt_pressure, - pe0_ptop_xmax, - pe_pk_delp_peln, - pn2_pk_delp, - pressures_mapu, - pressures_mapv, -) -from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 -from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment +from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyFV3.tracers import Tracers class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): inputs = { - "pe_": { - "name": "pe_", + "pe": { + "name": "pe", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, - "qvapor": { - "name": "qvapor", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qliquid": { - "name": "qliquid", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qice": { - "name": "qice", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qrain": { - "name": "qrain", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsnow": { - "name": "qsnow", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qgraupel": { - "name": "qgraupel", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qcld": { - "name": "qcld", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qo3mr": { - "name": "qo3mr", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsgs_tke": { - "name": "qsgs_tke", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, "delp": { "name": "delp", "dims": [X_DIM, Y_DIM, Z_DIM], @@ -116,8 +53,8 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "peln_3d": { - "name": "peln_3d", + "peln": { + "name": "peln", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, @@ -133,7 +70,7 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, "pkz": { @@ -146,11 +83,6 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "ws_": { - "name": "ws_", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, "u": { "name": "u", "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], @@ -161,36 +93,51 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "mfy": { + "mfy_R4": { "name": "mfy", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "cy": { + "cy_R4": { "name": "cy", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "mfx_": { - "name": "mfx_", + "mfx_R4": { + "name": "mfx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "cx_": { - "name": "cx_", + "cx_R4": { + "name": "cx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "hs": { - "name": "hs", + "phis": { + "name": "phis", "dims": [X_DIM, Y_DIM], "units": "No Units", }, - "te0_2d_": { - "name": "te0_2d_", + "te_2d": { + "name": "te_2d", "dims": [X_DIM, Y_DIM], "units": "No Units", }, + "wsd": { + "name": "wsd", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, + "dp1": { + "name": "dp1", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "pfull": { + "name": "pfull", + "dims": [Z_DIM], + "units": "No Units", + }, } outputs = { "pt": { @@ -213,51 +160,6 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "qvapor": { - "name": "qvapor", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qliquid": { - "name": "qliquid", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qice": { - "name": "qice", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qrain": { - "name": "qrain", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsnow": { - "name": "qsnow", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qgraupel": { - "name": "qgraupel", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qcld": { - "name": "qcld", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qo3mr": { - "name": "qo3mr", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsgs_tke": { - "name": "qsgs_tke", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, "w": { "name": "w", "dims": [X_DIM, Y_DIM, Z_DIM], @@ -273,39 +175,39 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "mfy": { + "mfy_R4": { "name": "mfy", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "cy": { + "cy_R4": { "name": "cy", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "mfx_": { - "name": "mfx_", + "mfx_R4": { + "name": "mfx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "cx_": { - "name": "cx_", + "cx_R4": { + "name": "cx", "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, - "peln_3d": { - "name": "peln_3d", + "peln": { + "name": "peln", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, - "pe_": { - "name": "pe_", + "pe": { + "name": "pe", "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], "units": "No Units", }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, "pkz": { @@ -318,6 +220,16 @@ class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): "dims": [X_DIM, Y_DIM, Z_DIM], "units": "No Units", }, + "dp1": { + "name": "dp1", + "dims": [X_DIM, Y_DIM, Z_DIM], + "units": "No Units", + }, + "ps": { + "name": "ps", + "dims": [X_DIM, Y_DIM], + "units": "No Units", + }, } def __init__( @@ -329,55 +241,19 @@ def __init__( super().__init__(grid, namelist, stencil_factory) self._base.in_vars["data_vars"] = { - "pe_": { - "istart": grid.is_ - 1, - "iend": grid.ie + 1, - "jstart": grid.js - 1, - "jend": grid.je + 1, - "kend": grid.npz + 1, - }, - "qvapor": { - "kend": grid.npz - 1, - }, - "qliquid": { - "kend": grid.npz - 1, - }, - "qice": { - "kend": grid.npz - 1, - }, - "qrain": { - "kend": grid.npz - 1, - }, - "qsnow": { - "kend": grid.npz - 1, - }, - "qgraupel": { - "kend": grid.npz - 1, - }, - "qcld": { - "kend": grid.npz - 1, - }, - "qo3mr": { - "kend": grid.npz - 1, - }, - "qsgs_tke": { + "tracers": {}, + "w": { "kend": grid.npz - 1, }, - "delp": {}, + "u": grid.y3d_domain_dict(), + "v": grid.x3d_domain_dict(), "delz": {}, - "q_con": {}, "pt": {}, + "dp1": {}, + "delp": {}, "cappa": {}, - "ps": {}, - "peln_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "ak": {}, - "bk": {}, + "q_con": {}, + "pkz": grid.compute_dict(), "pk": { "istart": grid.is_, "iend": grid.ie, @@ -385,713 +261,209 @@ def __init__( "jend": grid.je, "kend": grid.npz + 1, }, - "pkz": { + "peln": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, + "kaxis": 1, + "kend": grid.npz, }, - "w": { - "kend": grid.npz - 1, + "pe": { + "istart": grid.is_ - 1, + "iend": grid.ie + 1, + "jstart": grid.js - 1, + "jend": grid.je + 1, + "kend": grid.npz + 1, + "kaxis": 1, }, - "ws_": { + "ps": {}, + "wsd": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, }, - "u": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.jsd, - "jend": grid.jed + 1, - "kend": grid.npz, - }, - "v": { - "istart": grid.isd, - "iend": grid.ied + 1, - "jstart": grid.jsd, - "jend": grid.jed, - "kend": grid.npz, - }, - "mfy": { + "mfy_R4": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je + 1, "kend": grid.npz - 1, }, - "cy": { + "cy_R4": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.js, "jend": grid.je + 1, "kend": grid.npz - 1, }, - "mfx_": { + "mfx_R4": { "istart": grid.is_, "iend": grid.ie + 1, "jstart": grid.js, "jend": grid.je, "kend": grid.npz - 1, }, - "cx_": { + "cx_R4": { "istart": grid.is_, "iend": grid.ie + 1, "jstart": grid.jsd, "jend": grid.jed, "kend": grid.npz - 1, }, - "hs": { + "phis": { "istart": grid.isd, "iend": grid.ied, "jstart": grid.jsd, "jend": grid.jed, }, - "te0_2d_": { + "te_2d": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, }, + # column variables... + "ak": {}, + "bk": {}, + "pfull": grid.compute_buffer_k_dict(), } self._base.in_vars["parameters"] = [ "ptop", - "r_vir", "akap", + "zvir", "last_step", - "do_adiabatic_init", - "consv", - "adiabatic", + "consv_te", + "mdt", "nq", ] - self._base.out_vars = { - "pt": {}, - "cappa": {}, - "q_con": {}, - "delp": {}, - "delz": {}, - "qvapor": { - "kend": grid.npz - 1, - }, - "qliquid": { - "kend": grid.npz - 1, - }, - "qice": { - "kend": grid.npz - 1, - }, - "qrain": { - "kend": grid.npz - 1, - }, - "qsnow": { - "kend": grid.npz - 1, - }, - "qgraupel": { - "kend": grid.npz - 1, - }, - "qcld": { - "kend": grid.npz - 1, - }, - "qo3mr": { - "kend": grid.npz - 1, - }, - "qsgs_tke": { - "kend": grid.npz - 1, - }, - "w": { - "kend": grid.npz - 1, - }, - "u": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.jsd, - "jend": grid.jed + 1, - "kend": grid.npz - 1, - }, - "v": { - "istart": grid.isd, - "iend": grid.ied + 1, - "jstart": grid.jsd, - "jend": grid.jed, - "kend": grid.npz - 1, - }, - "mfy": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je + 1, - "kend": grid.npz - 1, - }, - "cy": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.js, - "jend": grid.je + 1, - "kend": grid.npz - 1, - }, - "mfx_": { - "istart": grid.is_, - "iend": grid.ie + 1, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz - 1, - }, - "cx_": { - "istart": grid.is_, - "iend": grid.ie + 1, - "jstart": grid.jsd, - "jend": grid.jed, - "kend": grid.npz - 1, - }, - "peln_3d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pe_": { - "istart": grid.is_ - 1, - "iend": grid.ie + 1, - "jstart": grid.js - 1, - "jend": grid.je + 1, - "kend": grid.npz + 1, - }, - "pk": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "pkz": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - }, - } + self._base.out_vars = {} + for k in [ + "tracers", + "pe", + "pkz", + "pk", + "peln", + "pt", + "cappa", + "delp", + "delz", + "q_con", + "u", + "v", + "w", + "ps", + "dp1", + "mfy_R4", + "cy_R4", + "mfx_R4", + "cx_R4", + ]: + self._base.out_vars[k] = self._base.in_vars["data_vars"][k] self.stencil_factory = stencil_factory self.quantity_factory = grid.quantity_factory - # self.namelist found in ParallelTranslateBaseSlicing - config = DynamicalCoreConfig.from_namelist(self.namelist).remapping - - hydrostatic = config.hydrostatic - if hydrostatic: - raise NotImplementedError("Hydrostatic is not implemented") - - grid_indexing = stencil_factory.grid_indexing - - self._t_min = Float(184.0) - self._w_max = Float(90.0) - self._w_min = Float(-60.0) - self._area_64 = self.grid.grid_data.area_64 - self._cosa_s = self.grid.grid_data.cosa_s - self._rsin2 = self.grid.grid_data.rsin2 - self._kord_tr = config.kord_tr - self._kord_tm = abs(config.kord_tm) - self._kord_wz = config.kord_wz - self._kord_mt = config.kord_mt - self._do_sat_adjust = config.do_sat_adj - - self.fill = True - - self._gz = self.quantity_factory.zeros( - [X_DIM, Y_DIM], - units="m^2 s^-2", - dtype=Float, - ) - - self._w2 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="temp W", - dtype=Float, - ) - - self._zsum1 = self.quantity_factory.zeros( - [X_DIM, Y_DIM], - units="Pa", - dtype=Float, - ) - - self._compute_performed = self.quantity_factory.zeros( - [X_DIM, Y_DIM], - units="mask", - dtype=bool, - ) - - self._ps = self.quantity_factory.zeros( - [X_DIM, Y_DIM], - units="Pa", - dtype=Float, - ) - - self._pe0 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], - units="Pa", - dtype=Float, - ) - - self._pe1 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], - units="Pa", - dtype=Float, - ) - - self._pe2 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], - units="Pa", - dtype=Float, - ) - - self._pe3 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], - units="Pa", - dtype=Float, - ) - - self._pn1 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="Pa", - dtype=Float, - ) - - self._pn2 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="Pa", - dtype=Float, - ) - - self._dp2 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="Pa", - dtype=Float, - ) - - self._pk2 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="Pa", - dtype=Float, - ) - - self._phis = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], - units="n/a", - dtype=Float, - ) - - self._te_2d = self.quantity_factory.zeros( - [X_DIM, Y_DIM], - units="Pa", - dtype=Float, - ) - - self._init_pe = stencil_factory.from_origin_domain( - init_pe, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(0, 1, 1)), - ) - - self._moist_cv_pt_pressure = stencil_factory.from_origin_domain( - moist_cv_pt_pressure, - externals={"hydrostatic": hydrostatic}, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(0, 0, 1)), - ) - - self._pn2_pk_delp = stencil_factory.from_origin_domain( - pn2_pk_delp, - origin=grid_indexing.origin_compute(add=(0, 0, 1)), - domain=grid_indexing.domain_compute(add=(0, 0, -1)), - ) - - self._map_scalar = MapSingle( - self.stencil_factory, - self.quantity_factory, - self._kord_tm, - 1, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - - self._rescale_delz_1 = stencil_factory.from_origin_domain( - rescale_delz_1, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), - ) - - self._rescale_delz_2 = stencil_factory.from_origin_domain( - rescale_delz_2, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), - ) - - self._w_fix_consrv_moment = stencil_factory.from_origin_domain( - func=W_fix_consrv_moment, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), - ) - - self._pressures_mapu = stencil_factory.from_origin_domain( - pressures_mapu, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(0, 1, 1)), - ) - - self._pe0_ptop_xmax = stencil_factory.from_origin_domain( - pe0_ptop_xmax, - origin=( - grid_indexing.n_halo + grid_indexing.domain[0], - grid_indexing.n_halo, - 0, - ), - domain=(1, grid_indexing.domain[1] + 1, 1), - ) - - self._pressures_mapv = stencil_factory.from_origin_domain( - pressures_mapv, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(1, 0, 1)), - ) - - self._pe_pk_delp_peln = stencil_factory.from_origin_domain( - pe_pk_delp_peln, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(0, 0, 1)), - ) - - self._moist_cv_pkz = stencil_factory.from_origin_domain( - moist_cv.moist_pkz, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), - ) - - self._moist_cv_te = stencil_factory.from_origin_domain( - moist_cv.moist_te, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(0, 0, 1)), - ) - - self._te_zsum = stencil_factory.from_origin_domain( - moist_cv.te_zsum, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), - ) - - self._moist_cv_pt_last_step = stencil_factory.from_origin_domain( - moist_cv.moist_pt_last_step, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(0, 0, 1)), - ) - - self._fill_cond = stencil_factory.from_origin_domain( - moist_cv.cond_output, - origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), - ) - - self._map1_ppm_u = MapSingle( - self.stencil_factory, - self.quantity_factory, - self._kord_mt, - -1, - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], - ) - - self._map1_ppm_v = MapSingle( - self.stencil_factory, - self.quantity_factory, - self._kord_mt, - -1, - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], - ) - - self._map_single_w = MapSingle( - self.stencil_factory, - self.quantity_factory, - self._kord_wz, - -2, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - - self._map_single_delz = MapSingle( - self.stencil_factory, - self.quantity_factory, - self._kord_wz, - 1, - dims=[X_DIM, Y_DIM, Z_DIM], - ) + self.stencil_factory = stencil_factory + self.namelist = DynamicalCoreConfig.from_namelist(namelist) + self.grid = grid def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") - def compute_parallel(self, inputs, communicator): - state = self.state_from_inputs(inputs) - state_namespace = SimpleNamespace(**state) - - tracers = { - "qvapor": state_namespace.qvapor, - "qliquid": state_namespace.qliquid, - "qice": state_namespace.qice, - "qrain": state_namespace.qrain, - "qsnow": state_namespace.qsnow, - "qgraupel": state_namespace.qgraupel, - "qcld": state_namespace.qcld, - "qo3mr": state_namespace.qo3mr, - "qsgs_tke": state_namespace.qsgs_tke, - } + def state_from_inputs(self, inputs: dict, tracers: Tracers) -> SimpleNamespace: + input_storages = super().state_from_inputs(inputs) + # Rename fluxes and courant numbers + input_storages["mfx"] = input_storages.pop("mfx_R4") + input_storages["mfy"] = input_storages.pop("mfy_R4") + input_storages["cx"] = input_storages.pop("cx_R4") + input_storages["cy"] = input_storages.pop("cy_R4") + # Make tracers + input_storages["tracers"] = tracers + return SimpleNamespace(**input_storages) + + def outputs_from_state(self, state: dict): + if len(self.outputs) == 0: + return {} + outputs = {} + storages = {} + for name, properties in self.outputs.items(): + if name in ["mfx_R4", "mfy_R4", "cx_R4", "cy_R4"]: + storages[name] = state[name[:-3]] + elif isinstance(state[name], Quantity): + storages[name] = state[name].data + elif len(self.outputs[name]["dims"]) > 0: + storages[name] = state[name] # assume it's a storage + else: + outputs[name] = state[name] # scalar + # Put tracers + storages["tracers"] = state["tracers"].as_4D_array() + outputs.update(self._base.slice_output(storages)) + return outputs - self._init_pe( - state_namespace.pe_, - self._pe1, - self._pe2, - state_namespace.ptop, - ) - - self._moist_cv_pt_pressure( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - state_namespace.q_con, - state_namespace.pt, - state_namespace.cappa, - state_namespace.delp, - state_namespace.delz, - state_namespace.pe_, - self._pe2, - state_namespace.ak, - state_namespace.bk, - self._dp2, - self._ps, - self._pn1, - self._pn2, - state_namespace.peln_3d, - True, - Float(state_namespace.r_vir), - ) - - self._pn2_pk_delp( - self._pe2, - self._pn2, - self._pk2, - Float(state_namespace.akap), + def compute_parallel(self, inputs, communicator): + tracers_mapping = Tracers.blind_mapping_from_data(inputs["tracers"]) + tracers_mapping[0] = "vapor" + tracers_mapping[1] = "liquid" + tracers_mapping[2] = "rain" + tracers_mapping[3] = "snow" + tracers_mapping[4] = "ice" + tracers_mapping[5] = "graupel" + tracers_mapping[6] = "cloud" + tracers = Tracers.make_from_4D_array( + self.quantity_factory, + tracers_mapping, + inputs["tracers"], ) - self._map_scalar( - state_namespace.pt, - self._pn1, - self._pn2, - self._t_min, - interp=True, - ) + inputs["te_2d"] = inputs["te_2d"].astype(Float) + state = self.state_from_inputs(inputs, tracers) - self._mapn_tracer = MapNTracer( + l_to_e = LagrangianToEulerian_GEOS( self.stencil_factory, self.quantity_factory, - abs(self._kord_tr), - state_namespace.nq, - fill=self.fill, - tracers=tracers, - ) - - self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) - - self._map_single_w( - state_namespace.w, - self._pe1, - self._pe2, - qs=state_namespace.ws_, - interp=False, - ) - - self._rescale_delz_1( - state_namespace.delz, - state_namespace.delp, - ) - - self._map_single_delz(state_namespace.delz, self._pe1, self._pe2) - - self._rescale_delz_2( - state_namespace.delz, - self._dp2, - ) - - self._w_fix_consrv_moment( - state_namespace.w, - self._w2, - self._dp2, - self._gz, - self._w_max, - self._w_min, - self._compute_performed, - ) - - self._pressures_mapu( - state_namespace.pe_, - state_namespace.ak, - state_namespace.bk, - self._pe0, - self._pe3, - state_namespace.ptop, - ) - - self._pe0_ptop_xmax( - self._pe0, - state_namespace.ptop, - ) - - self._map1_ppm_u( - state_namespace.u, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_u( - state_namespace.mfy, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_u( - state_namespace.cy, - self._pe0, - self._pe3, - interp=False, - ) - - self._pressures_mapv( - state_namespace.pe_, - state_namespace.ak, - state_namespace.bk, - self._pe0, - self._pe3, - ) - - self._map1_ppm_v( - state_namespace.v, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_v( - state_namespace.mfx_, - self._pe0, - self._pe3, - interp=False, - ) - - self._map1_ppm_v( - state_namespace.cx_, - self._pe0, - self._pe3, - interp=False, - ) - - self._pe_pk_delp_peln( - state_namespace.pe_, - state_namespace.pk, - state_namespace.delp, - state_namespace.peln_3d, - self._pe2, - self._pk2, - self._pn2, - state_namespace.ak, - state_namespace.bk, - state_namespace.akap, - state_namespace.ptop, - ) - - self._moist_cv_pkz( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - state_namespace.pkz, - state_namespace.pt, - state_namespace.cappa, - state_namespace.delp, - state_namespace.delz, - Float(state_namespace.r_vir), - ) - - if state_namespace.last_step and not state_namespace.do_adiabatic_init: - - if state_namespace.consv > CONSV_MIN: - - self._moist_cv_te( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - state_namespace.u, - state_namespace.v, - state_namespace.w, - self._te_2d, - state_namespace.pt, - self._phis, - state_namespace.delp, - self._rsin2, - self._cosa_s, - state_namespace.hs, - state_namespace.delz, - GRAV, - ) - - self._te_zsum( - self._te_2d, - state_namespace.te0_2d_, - state_namespace.delp, - state_namespace.pkz, - self._zsum1, - ) - - # Note, mpp_global_sum is currently set up for the C24 TBC setup - inputArray = self._te_2d.view[:] * self._area_64.view[:] - tesum = mpp_global_sum(inputArray, communicator, self.stencil_factory) - # print("tesum: ", tesum) - # print("type(self._zsum1) = ",type(self._zsum1)) - # print("type(self._area_64) = ",type(self._area_64)) - inputArray = self._zsum1.view[:] * self._area_64.view[:] - zsum = mpp_global_sum(inputArray, communicator, self.stencil_factory) - # print("zsum: ", zsum) - dtmp = tesum / (CV_AIR * zsum) - print("dtmp: ", dtmp) - # I ignore the E_flux computation since it's not used elsewhere - # in our current setup once it's computed - - if state_namespace.last_step and not state_namespace.adiabatic: - - self._moist_cv_pt_last_step( - tracers["qvapor"], - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - state_namespace.pt, - state_namespace.pkz, - Float(dtmp), - state_namespace.r_vir, - ) - - self._fill_cond( - state_namespace.q_con, - tracers["qliquid"], - tracers["qrain"], - tracers["qsnow"], - tracers["qice"], - tracers["qgraupel"], - ) - - return self.outputs_from_state(state) + DynamicalCoreConfig.from_namelist(self.namelist).remapping, + communicator, + self.grid.grid_data, + state.nq, + state.pfull, + state.tracers, + DynamicalCoreConfig.adiabatic, + ) + + l_to_e( + state.tracers, + state.pt, + state.delp, + state.delz, + state.peln, + state.u, + state.v, + state.w, + state.mfx, + state.mfy, + state.cx, + state.cy, + state.cappa, + state.q_con, + state.pkz, + state.pk, + state.pe, + state.phis, + state.te_2d, + state.ps, + state.wsd, + state.ak, + state.bk, + state.dp1, + state.ptop, + state.akap, + state.zvir, + state.last_step, + state.consv_te, + state.mdt, + ) + + outputs = self.outputs_from_state(vars(state)) + return outputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS_v2.py b/tests/savepoint/translate/translate_remapping_GEOS_v2.py deleted file mode 100644 index d1b236fc..00000000 --- a/tests/savepoint/translate/translate_remapping_GEOS_v2.py +++ /dev/null @@ -1,585 +0,0 @@ -from types import SimpleNamespace - -from ndsl import Namelist, StencilFactory -from ndsl.constants import ( - X_DIM, - X_INTERFACE_DIM, - Y_DIM, - Y_INTERFACE_DIM, - Z_DIM, - Z_INTERFACE_DIM, -) -from ndsl.dsl.typing import Float -from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing -from pyFV3 import DynamicalCoreConfig -from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS -from pyFV3.tracers import Tracers - - -# from pyFV3._config import RemappingConfig - - -class TranslateRemapping_GEOS_v2(ParallelTranslateBaseSlicing): - inputs = { - "pe": { - "name": "pe", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], - "units": "No Units", - }, - "qvapor": { - "name": "qvapor", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qliquid": { - "name": "qliquid", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qice": { - "name": "qice", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qrain": { - "name": "qrain", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsnow": { - "name": "qsnow", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qgraupel": { - "name": "qgraupel", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qcld": { - "name": "qcld", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qo3mr": { - "name": "qo3mr", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsgs_tke": { - "name": "qsgs_tke", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "delp": { - "name": "delp", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "delz": { - "name": "delz", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "q_con": { - "name": "q_con", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "pt": { - "name": "pt", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "cappa": { - "name": "cappa", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "ps": { - "name": "ps", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, - "peln": { - "name": "peln", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], - "units": "No Units", - }, - "ak": { - "name": "ak", - "dims": [Z_INTERFACE_DIM], - "units": "No Units", - }, - "bk": { - "name": "bk", - "dims": [Z_INTERFACE_DIM], - "units": "No Units", - }, - "pk": { - "name": "pk", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "pkz": { - "name": "pkz", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "w": { - "name": "w", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "u": { - "name": "u", - "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], - "units": "No Units", - }, - "v": { - "name": "v", - "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "mfy": { - "name": "mfy", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "cy": { - "name": "cy", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "mfx": { - "name": "mfx", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "cx": { - "name": "cx", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "hs": { - "name": "hs", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, - "te0_2d": { - "name": "te0_2d", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, - "wsd": { - "name": "wsd", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, - "dp1": { - "name": "dp1", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - } - outputs = { - "pt": { - "name": "pt", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "cappa": { - "name": "cappa", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "delp": { - "name": "delp", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "delz": { - "name": "delz", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qvapor": { - "name": "qvapor", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qliquid": { - "name": "qliquid", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qice": { - "name": "qice", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qrain": { - "name": "qrain", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsnow": { - "name": "qsnow", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qgraupel": { - "name": "qgraupel", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qcld": { - "name": "qcld", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qo3mr": { - "name": "qo3mr", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "qsgs_tke": { - "name": "qsgs_tke", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "w": { - "name": "w", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "u": { - "name": "u", - "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], - "units": "No Units", - }, - "v": { - "name": "v", - "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "mfy": { - "name": "mfy", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "cy": { - "name": "cy", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "mfx": { - "name": "mfx", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "cx": { - "name": "cx", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "peln": { - "name": "peln", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], - "units": "No Units", - }, - "pe": { - "name": "pe", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], - "units": "No Units", - }, - "pk": { - "name": "pk", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "pkz": { - "name": "pkz", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "q_con": { - "name": "q_con", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "dp1": { - "name": "dp1", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "No Units", - }, - "ps": { - "name": "ps", - "dims": [X_DIM, Y_DIM], - "units": "No Units", - }, - } - - def __init__( - self, - grid: Grid, - namelist: Namelist, - stencil_factory: StencilFactory, - ): - super().__init__(grid, namelist, stencil_factory) - - self._base.in_vars["data_vars"] = { - # "tracers": {}, - "qvapor": { - "kend": grid.npz - 1, - }, - "qliquid": { - "kend": grid.npz - 1, - }, - "qice": { - "kend": grid.npz - 1, - }, - "qrain": { - "kend": grid.npz - 1, - }, - "qsnow": { - "kend": grid.npz - 1, - }, - "qgraupel": { - "kend": grid.npz - 1, - }, - "qcld": { - "kend": grid.npz - 1, - }, - "qo3mr": { - "kend": grid.npz - 1, - }, - "qsgs_tke": { - "kend": grid.npz - 1, - }, - "w": { - "kend": grid.npz - 1, - }, - "u": grid.y3d_domain_dict(), - "v": grid.x3d_domain_dict(), - "delz": {}, - "pt": {}, - "dp1": {}, - "delp": {}, - "cappa": {}, - "q_con": {}, - "pkz": grid.compute_dict(), - "pk": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz + 1, - }, - "peln": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - "kaxis": 1, - "kend": grid.npz, - }, - "pe": { - "istart": grid.is_ - 1, - "iend": grid.ie + 1, - "jstart": grid.js - 1, - "jend": grid.je + 1, - "kend": grid.npz + 1, - "kaxis": 1, - }, - # "hs": {"serialname": "phis"}, - "ps": {}, - "wsd": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - }, - "mfy": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je + 1, - "kend": grid.npz - 1, - }, - "cy": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.js, - "jend": grid.je + 1, - "kend": grid.npz - 1, - }, - "mfx": { - "istart": grid.is_, - "iend": grid.ie + 1, - "jstart": grid.js, - "jend": grid.je, - "kend": grid.npz - 1, - }, - "cx": { - "istart": grid.is_, - "iend": grid.ie + 1, - "jstart": grid.jsd, - "jend": grid.jed, - "kend": grid.npz - 1, - }, - "hs": { - "istart": grid.isd, - "iend": grid.ied, - "jstart": grid.jsd, - "jend": grid.jed, - }, - "te0_2d": { - "istart": grid.is_, - "iend": grid.ie, - "jstart": grid.js, - "jend": grid.je, - }, - # column variables... - "ak": {}, - "bk": {}, - "pfull": {}, - } - self._base.in_vars["parameters"] = [ - "ptop", - "akap", - "zvir", - "last_step", - "consv_te", - "mdt", - "nq", - ] - self._base.out_vars = {} - for k in [ - "pe", - "pkz", - "pk", - "peln", - "pt", - # "tracers", - "qvapor", - "qliquid", - "qice", - "qrain", - "qsnow", - "qgraupel", - "qcld", - "qo3mr", - "qsgs_tke", - "cappa", - "delp", - "delz", - "q_con", - "u", - "v", - "w", - "ps", - "dp1", - "mfy", - "cy", - "mfx", - "cx", - ]: - self._base.out_vars[k] = self._base.in_vars["data_vars"][k] - - self.stencil_factory = stencil_factory - self.quantity_factory = grid.quantity_factory - - self.stencil_factory = stencil_factory - self.namelist = DynamicalCoreConfig.from_namelist(namelist) - self.grid = grid - - def compute_sequential(self, inputs_list, communicator_list): - print("No serial test available") - - def compute_parallel(self, inputs, communicator): - - tracers = Tracers(self.quantity_factory) - - tracers.copy_tracer_data("qvapor", inputs["qvapor"]) - tracers.copy_tracer_data("qliquid", inputs["qliquid"]) - tracers.copy_tracer_data("qice", inputs["qice"]) - tracers.copy_tracer_data("qrain", inputs["qrain"]) - tracers.copy_tracer_data("qsnow", inputs["qsnow"]) - tracers.copy_tracer_data("qgraupel", inputs["qgraupel"]) - tracers.copy_tracer_data("qcld", inputs["qcld"]) - tracers.copy_tracer_data("qo3mr", inputs["qo3mr"]) - tracers.copy_tracer_data("qsgs_tke", inputs["qsgs_tke"]) - - inputs["te0_2d"] = inputs["te0_2d"].astype(Float) - state = self.state_from_inputs(inputs) - state_namespace = SimpleNamespace(**state) - - # tracers = { - # "qvapor": state_namespace.qvapor, - # "qliquid": state_namespace.qliquid, - # "qice": state_namespace.qice, - # "qrain": state_namespace.qrain, - # "qsnow": state_namespace.qsnow, - # "qgraupel": state_namespace.qgraupel, - # "qcld": state_namespace.qcld, - # "qo3mr": state_namespace.qo3mr, - # "qsgs_tke": state_namespace.qsgs_tke, - # } - - l_to_e = LagrangianToEulerian_GEOS( - self.stencil_factory, - self.quantity_factory, - DynamicalCoreConfig.from_namelist(self.namelist).remapping, - communicator, - self.grid.grid_data, - state_namespace.nq, - state_namespace.pfull, - tracers, - DynamicalCoreConfig.adiabatic, - ) - - l_to_e( - tracers, - state_namespace.pt, - state_namespace.delp, - state_namespace.delz, - state_namespace.peln, - state_namespace.u, - state_namespace.v, - state_namespace.w, - state_namespace.mfx, - state_namespace.mfy, - state_namespace.cx, - state_namespace.cy, - state_namespace.cappa, - state_namespace.q_con, - state_namespace.pkz, - state_namespace.pk, - state_namespace.pe, - state_namespace.hs, - state_namespace.te0_2d, - state_namespace.ps, - state_namespace.wsd, - state_namespace.ak, - state_namespace.bk, - state_namespace.dp1, - Float(state_namespace.ptop), - Float(state_namespace.akap), - state_namespace.zvir, - state_namespace.last_step, - state_namespace.consv_te, - state_namespace.mdt, - ) - - state_namespace.qvapor.data = tracers["qvapor"].data - state_namespace.qliquid.data = tracers["qliquid"].data - state_namespace.qice.data = tracers["qice"].data - state_namespace.qrain.data = tracers["qrain"].data - state_namespace.qsnow.data = tracers["qsnow"].data - state_namespace.qgraupel.data = tracers["qgraupel"].data - state_namespace.qcld.data = tracers["qcld"].data - state_namespace.qo3mr.data = tracers["qo3mr"].data - state_namespace.qsgs_tke.data = tracers["qsgs_tke"].data - - return self.outputs_from_state(state) From c0d8685ff5fe2b51d98eab6059da0e9f37607108 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 1 May 2025 16:25:27 -0400 Subject: [PATCH 206/252] Fix copy to cover interface dimensions --- pyFV3/stencils/fv_dynamics.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index ed95d93b..8758d8a2 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -1,5 +1,5 @@ from datetime import timedelta -from typing import Mapping, Optional +from typing import Mapping from dace.frontend.python.interface import nounroll as dace_no_unroll from gt4py.cartesian.gtscript import FORWARD, PARALLEL, computation, interval @@ -297,6 +297,29 @@ def __init__( f" nwat=={config.nwat} is not implemented." " Only nwat=6 has been implemented." ) + + # Implemented dynamics options require those tracers to be present at minima + # this is a more granular list than carried by the `nwat` single integer + # but cover the same topic + required_tracers = [ + "vapor", + "liquid", + "rain", + "snow", + "ice", + "graupel", + "cloud", + ] + if not all(n in state.tracers.names() for n in required_tracers): + raise NotImplementedError( + "Dynamical core (fv_dynamics):" + " missing required tracers. Dynamics requires:\n" + f" {required_tracers}\n" + "but only the following where given:\n" + f" {state.tracers.names()}" + ) + + self._comm = comm self.comm_rank = comm.rank self.grid_data = grid_data self.grid_indexing = grid_indexing @@ -516,7 +539,7 @@ def __init__( self._copy_cast = stencil_factory.from_origin_domain( func=_copy_cast_defn, origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), + domain=grid_indexing.domain_compute(add=(1, 1, 0)), ) # See divergence_damping.py, _get_da_min for explanation of this function From 0484df8b8b2c1de67db8f3e128dd545a783f3687 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 1 May 2025 16:38:46 -0400 Subject: [PATCH 207/252] Fix increment of fluxes for interface dims --- pyFV3/stencils/fv_dynamics.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 8758d8a2..02dc278f 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -1,5 +1,5 @@ from datetime import timedelta -from typing import Mapping +from typing import List, Mapping from dace.frontend.python.interface import nounroll as dace_no_unroll from gt4py.cartesian.gtscript import FORWARD, PARALLEL, computation, interval @@ -534,7 +534,7 @@ def __init__( self._increment = stencil_factory.from_origin_domain( func=_increment_stencil, origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), + domain=grid_indexing.domain_compute(add=(1, 1, 0)), ) self._copy_cast = stencil_factory.from_origin_domain( func=_copy_cast_defn, From 31acc83b81dcd5c320dcf0f779c5f6ee1320be39 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:11:34 -0400 Subject: [PATCH 208/252] Fix set value to cover interface --- pyFV3/stencils/fv_dynamics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index 02dc278f..70df1b53 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -529,7 +529,7 @@ def __init__( self._set_value = stencil_factory.from_origin_domain( func=set_value_defn, origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(), + domain=grid_indexing.domain_compute(add=(1, 1, 0)), ) self._increment = stencil_factory.from_origin_domain( func=_increment_stencil, From 03bc3e31d8fd6f0aa3daafbc07684123835fa20b Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:12:56 -0400 Subject: [PATCH 209/252] Flip refactor division to match fortran --- pyFV3/stencils/updatedzd.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyFV3/stencils/updatedzd.py b/pyFV3/stencils/updatedzd.py index b9c1ccb6..60c02b37 100644 --- a/pyFV3/stencils/updatedzd.py +++ b/pyFV3/stencils/updatedzd.py @@ -113,7 +113,8 @@ def apply_height_fluxes( with computation(BACKWARD): with interval(-1, None): - ws = (surface_height - height) / dt + rdt = 1 / dt + ws = (surface_height - height) * rdt with interval(0, -1): # ensure layer thickness exceeds minimum other = height[0, 0, 1] + DZ_MIN From 0ca04d79ee65996af05e80f10cb6234deb075edd Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:13:25 -0400 Subject: [PATCH 210/252] Localize ratio to match fortran --- pyFV3/stencils/updatedzc.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/pyFV3/stencils/updatedzc.py b/pyFV3/stencils/updatedzc.py index b88740b1..605b1779 100644 --- a/pyFV3/stencils/updatedzc.py +++ b/pyFV3/stencils/updatedzc.py @@ -1,5 +1,11 @@ import gt4py.cartesian.gtscript as gtscript -from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval +from gt4py.cartesian.gtscript import ( + BACKWARD, + FORWARD, + PARALLEL, + computation, + interval, +) import ndsl.constants as constants from ndsl import Quantity, QuantityFactory, StencilFactory @@ -12,16 +18,14 @@ @gtscript.function -def p_weighted_average_top(vel, dp0): +def p_weighted_average_top(vel, top_dp_ratio): # TODO: ratio is a constant, where should this be placed? - ratio = dp0 / (dp0 + dp0[1]) - return vel + (vel - vel[0, 0, 1]) * ratio + return vel + (vel - vel[0, 0, 1]) * top_dp_ratio @gtscript.function -def p_weighted_average_bottom(vel, dp0): - ratio = dp0[-1] / (dp0[-2] + dp0[-1]) - return vel[0, 0, -1] + (vel[0, 0, -1] - vel[0, 0, -2]) * ratio +def p_weighted_average_bottom(vel, bottom_dp_ratio): + return vel[0, 0, -1] + (vel[0, 0, -1] - vel[0, 0, -2]) * bottom_dp_ratio @gtscript.function From 87d3bbaf06cd5b6105d32195cfbf413ad2cb305e Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:14:35 -0400 Subject: [PATCH 211/252] Split p_grad by dimensions to forbid halos bleed --- pyFV3/stencils/dyn_core.py | 52 +++++++++++++++++++++++++++++--------- 1 file changed, 40 insertions(+), 12 deletions(-) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 07fd5f27..2df88a30 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -126,11 +126,9 @@ def compute_geopotential(zh: FloatField, gz: FloatField): gz = zh * constants.GRAV -def p_grad_c_stencil( +def p_grad_c_stencil_x( rdxc: FloatFieldIJ, - rdyc: FloatFieldIJ, uc: FloatField, - vc: FloatField, delpc: FloatField, pkc: FloatField, gz: FloatField, @@ -147,11 +145,8 @@ def p_grad_c_stencil( Args: rdxc (in): - rdyc (in): uc (inout): x-velocity on the C-grid, has been updated due to advection but not yet due to pressure gradient force - vc (inout): y-velocity on the C-grid, has been updated due to advection - but not yet due to pressure gradient force delpc (in): vertical delta in pressure pkc (in): pressure if non-hydrostatic, (edge pressure)**(moist kappa) if hydrostatic @@ -174,6 +169,26 @@ def p_grad_c_stencil( + (gz[-1, 0, 0] - gz[0, 0, 1]) * (pkc[-1, 0, 1] - pkc) ) + +def p_grad_c_stencil_y( + rdyc: FloatFieldIJ, + vc: FloatField, + delpc: FloatField, + pkc: FloatField, + gz: FloatField, + dt2: Float, +): + """ + See p_grad_c_stencil_y + """ + from __externals__ import hydrostatic + + with computation(PARALLEL), interval(...): + if __INLINED(hydrostatic): + wk = pkc[0, 0, 1] - pkc + else: + wk = delpc + # wk is pressure gradient vc = vc + dt2 * rdyc / (wk[0, -1, 0] + wk) * ( (gz[0, -1, 1] - gz) * (pkc[0, 0, 1] - pkc[0, -1, 0]) + (gz[0, -1, 0] - gz[0, 0, 1]) * (pkc[0, -1, 1] - pkc) @@ -539,10 +554,16 @@ def __init__( ) ) - self._p_grad_c = stencil_factory.from_origin_domain( - p_grad_c_stencil, + self._p_grad_c_x = stencil_factory.from_origin_domain( + p_grad_c_stencil_x, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(add=(1, 0, 0)), + externals={"hydrostatic": config.hydrostatic}, + ) + self._p_grad_c_y = stencil_factory.from_origin_domain( + p_grad_c_stencil_y, origin=grid_indexing.origin_compute(), - domain=grid_indexing.domain_compute(add=(1, 1, 0)), + domain=grid_indexing.domain_compute(add=(0, 1, 0)), externals={"hydrostatic": config.hydrostatic}, ) @@ -560,7 +581,7 @@ def __init__( self._zero_data = stencil_factory.from_origin_domain( zero_data, origin=grid_indexing.origin_full(), - domain=grid_indexing.domain_full(), + domain=grid_indexing.domain_full(add=(1, 1, 0)), ) ax_offsets_pe = grid_indexing.axis_offsets( grid_indexing.origin_full(), @@ -781,16 +802,23 @@ def __call__( w3=state.omga, ) - self._p_grad_c( + self._p_grad_c_x( rdxc=self.grid_data.rdxc, - rdyc=self.grid_data.rdyc, uc=state.uc, + delpc=self.cgrid_shallow_water_lagrangian_dynamics.delpc, + pkc=self._pkc, + gz=self._gz, + dt2=dt2, + ) + self._p_grad_c_y( + rdyc=self.grid_data.rdyc, vc=state.vc, delpc=self.cgrid_shallow_water_lagrangian_dynamics.delpc, pkc=self._pkc, gz=self._gz, dt2=dt2, ) + self._halo_updaters.uc__vc.start() if self.config.nord > 0: self._halo_updaters.divgd.wait() From cd605abd5884cee83b442be537ebaf0e2a073295 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:15:12 -0400 Subject: [PATCH 212/252] Correct docs --- pyFV3/stencils/fvtp2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/stencils/fvtp2d.py b/pyFV3/stencils/fvtp2d.py index 77aad226..4e5fe3e2 100644 --- a/pyFV3/stencils/fvtp2d.py +++ b/pyFV3/stencils/fvtp2d.py @@ -274,7 +274,7 @@ def __call__( by contrast are area weighted. Args: - q (in): scalar to be transported + q (inout): scalar to be transported (corners are copied in halo) crx (in): Courant number in x-direction cry (in): Courant number in y-direction x_area_flux (in): flux of area in x-direction, in units of m^2 From 9ae1b062f47f15aeff0fcb0ac694f559baeebbc1 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:15:37 -0400 Subject: [PATCH 213/252] Restore `UnreachableStmtPruning` to fix dace:cpu layout transparency issue --- pyFV3/stencils/divergence_damping.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pyFV3/stencils/divergence_damping.py b/pyFV3/stencils/divergence_damping.py index e42d86c6..6afd5d09 100644 --- a/pyFV3/stencils/divergence_damping.py +++ b/pyFV3/stencils/divergence_damping.py @@ -501,7 +501,6 @@ def __init__( func=corners.fill_corners_dgrid_defn, compute_dims=[X_INTERFACE_DIM, Y_INTERFACE_DIM, Z_DIM], compute_halos=(self.grid_indexing.n_halo, self.grid_indexing.n_halo), - skip_passes=("UnreachableStmtPruning",), ) self._redo_divg_d_stencils = get_stencils_with_varied_bounds( From 64063e0a593402887f53e37bb763ac86f1628df5 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:16:09 -0400 Subject: [PATCH 214/252] Extend translate test to halo --- .../savepoint/translate/translate_a2b_ord4.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tests/savepoint/translate/translate_a2b_ord4.py b/tests/savepoint/translate/translate_a2b_ord4.py index acc4808f..dfba1ed0 100644 --- a/tests/savepoint/translate/translate_a2b_ord4.py +++ b/tests/savepoint/translate/translate_a2b_ord4.py @@ -1,9 +1,11 @@ +import numpy as np from typing import Any, Dict from ndsl import Namelist, StencilFactory, orchestrate -from ndsl.constants import Z_DIM +from ndsl.constants import X_DIM, Y_DIM, Z_DIM from pyFV3.stencils import DivergenceDamping from pyFV3.testing import TranslateDycoreFortranData2Py +from pyFV3.utils.functional_validation import get_subset_func class A2B_Ord4Compute: @@ -58,6 +60,11 @@ def __init__( self.namelist = namelist # type: ignore self.stencil_factory = stencil_factory self.compute_obj = A2B_Ord4Compute(stencil_factory) + self._subset = get_subset_func( + self.grid.grid_indexing, + dims=[X_DIM, Y_DIM, Z_DIM], + n_halo=((3, 3), (3, 3)), + ) def compute_from_storage(self, inputs): nord_col = self.grid.quantity_factory.zeros(dims=[Z_DIM], units="unknown") @@ -80,3 +87,13 @@ def compute_from_storage(self, inputs): inputs["grid_type"] = 0 self.compute_obj(divdamp, **inputs) return inputs + + def subset_output(self, varname: str, output: np.ndarray) -> np.ndarray: + """ + Given an output array, return the slice of the array which we'd + like to validate against reference data + """ + if varname in ["wk"]: + return self._subset(output) + else: + return output From 3ef943cb7b55ab208911692f7f433764613dbbd1 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:16:41 -0400 Subject: [PATCH 215/252] Fix translate sizes --- tests/savepoint/translate/translate_updatedzd.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/savepoint/translate/translate_updatedzd.py b/tests/savepoint/translate/translate_updatedzd.py index a340cfc0..358e4dd0 100644 --- a/tests/savepoint/translate/translate_updatedzd.py +++ b/tests/savepoint/translate/translate_updatedzd.py @@ -52,7 +52,7 @@ def __init__( self._subset = get_subset_func( self.grid.grid_indexing, dims=[X_DIM, Y_DIM, Z_DIM], - n_halo=((0, 0), (0, 0)), + n_halo=((3, 3), (3, 3)), ) self.ignore_near_zero_errors = {"zh": True, "wsd": True} self.near_zero = 1e-30 From 7d5987fb5a30feae237c5c01a077e530dcb01db6 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:17:16 -0400 Subject: [PATCH 216/252] Extend rieman solver test to halos --- .../translate/translate_riem_solver_c.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/savepoint/translate/translate_riem_solver_c.py b/tests/savepoint/translate/translate_riem_solver_c.py index aa193d67..9b5f0d11 100644 --- a/tests/savepoint/translate/translate_riem_solver_c.py +++ b/tests/savepoint/translate/translate_riem_solver_c.py @@ -1,6 +1,9 @@ +import numpy as np from ndsl import Namelist, StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM from pyFV3.stencils import NonhydrostaticVerticalSolverCGrid from pyFV3.testing import TranslateDycoreFortranData2Py +from pyFV3.utils.functional_validation import get_subset_func class TranslateRiem_Solver_C(TranslateDycoreFortranData2Py): @@ -31,3 +34,23 @@ def __init__( self.out_vars = {"pef": {"kend": grid.npz}, "gz": {"kend": grid.npz}} self.max_error = 5e-14 self.stencil_factory = stencil_factory + self._subset = get_subset_func( + self.grid.grid_indexing, + dims=[X_DIM, Y_DIM, Z_DIM], + n_halo=((3, 3), (3, 3)), + ) + + def compute(self, inputs): + outputs = super().compute(inputs) + outputs["gz"] = self.subset_output("gz", outputs["gz"]) + return outputs + + def subset_output(self, varname: str, output: np.ndarray) -> np.ndarray: + """ + Given an output array, return the slice of the array which we'd + like to validate against reference data + """ + if varname in ["gz", "pef"]: + return self._subset(output) + else: + return output From a1c113217371c9fa2b1c873d5116caa833dbc74d Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:17:55 -0400 Subject: [PATCH 217/252] Docs and verbose parameters --- pyFV3/stencils/delnflux.py | 81 +++++++++++++++++++++++++++++--------- 1 file changed, 63 insertions(+), 18 deletions(-) diff --git a/pyFV3/stencils/delnflux.py b/pyFV3/stencils/delnflux.py index 2b96032f..8b89f8e4 100644 --- a/pyFV3/stencils/delnflux.py +++ b/pyFV3/stencils/delnflux.py @@ -429,7 +429,11 @@ def __init__( ) self.delnflux_nosg = DelnFluxNoSG( - stencil_factory, damping_coefficients, rarea, nord_col, nk=nk + stencil_factory, + damping_coefficients, + rarea, + nord_col, + nk=nk, ) def __call__( @@ -443,11 +447,11 @@ def __call__( """ Del-n damping for fluxes, where n = 2 * nord + 2 Args: - q: Field for which to calculate damped fluxes (in) - fx: x-flux on A-grid (inout) - fy: y-flux on A-grid (inout) - d2: A damped copy of the q field (in) - mass: Mass to weight the diffusive flux by (in) + q (in): Field for which to calculate damped fluxes + fx (inout): x-flux on A-grid + fy (inout): y-flux on A-grid + d2 (in): A damped copy of the q field + mass (in): Mass to weight the diffusive flux by """ if self._no_compute is True: return fx, fy @@ -617,14 +621,12 @@ def __init__( externals={**corner_axis_offsets}, origin=corner_origin, domain=corner_domain, - skip_passes=("UnreachableStmtPruning",), ) self._copy_corners_y_nord = stencil_factory.from_origin_domain( copy_corners_y_nord, externals={**corner_axis_offsets}, origin=corner_origin, domain=corner_domain, - skip_passes=("UnreachableStmtPruning",), ) def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): @@ -645,17 +647,46 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): """ if mass is None: - self._d2_damp(q=q, d2=d2, damp=damp_c, nord=self._nord) + self._d2_damp( + q=q, + d2=d2, + damp=damp_c, + nord=self._nord, + ) else: - self._copy_stencil_interval(q_in=q, q_out=d2, nord=self._nord) + self._copy_stencil_interval( + q_in=q, + q_out=d2, + nord=self._nord, + ) - self._copy_corners_x_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + self._copy_corners_x_nord( + q_in=d2, + q_out=d2, + nord=self._nord, + current_nord=0, + ) - self._fx_calc_stencil(q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord) + self._fx_calc_stencil( + q=d2, + del6_v=self._del6_v, + fx=fx2, + nord=self._nord, + ) - self._copy_corners_y_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + self._copy_corners_y_nord( + q_in=d2, + q_out=d2, + nord=self._nord, + current_nord=0, + ) - self._fy_calc_stencil(q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord) + self._fy_calc_stencil( + q=d2, + del6_u=self._del6_u, + fy=fy2, + nord=self._nord, + ) for n in range(self._nmax): self._d2_stencil[n]( @@ -668,17 +699,31 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): ) self._copy_corners_x_nord( - q_in=d2, q_out=d2, nord=self._nord, current_nord=n + q_in=d2, + q_out=d2, + nord=self._nord, + current_nord=n, ) self._column_conditional_fx_calculation[n]( - q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord, current_nord=n + q=d2, + del6_v=self._del6_v, + fx=fx2, + nord=self._nord, + current_nord=n, ) self._copy_corners_y_nord( - q_in=d2, q_out=d2, nord=self._nord, current_nord=n + q_in=d2, + q_out=d2, + nord=self._nord, + current_nord=n, ) self._column_conditional_fy_calculation[n]( - q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord, current_nord=n + q=d2, + del6_u=self._del6_u, + fy=fy2, + nord=self._nord, + current_nord=n, ) From 4764ff5c12e2c8d4709ef11f0ed4e188dd1a888d Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:18:23 -0400 Subject: [PATCH 218/252] Fix bad copy --- pyFV3/stencils/tracer_2d_1l.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index f2a0ba19..61dbe61c 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -367,7 +367,7 @@ def __call__( if self._update_mass_courant: working_x_mass_flux = x_mass_flux - working_y_mass_flux = x_mass_flux + working_y_mass_flux = y_mass_flux working_x_courant = x_courant working_y_courant = y_courant else: @@ -402,13 +402,13 @@ def __call__( ) self._divide_fluxes_by_n_substeps( - working_x_courant, - self._x_area_flux, - working_x_mass_flux, - working_y_courant, - self._y_area_flux, - working_y_mass_flux, - self._cmax, + cxd=working_x_courant, + xfx=self._x_area_flux, + mfxd=working_x_mass_flux, + cyd=working_y_courant, + yfx=self._y_area_flux, + mfyd=working_y_mass_flux, + cmax=self._cmax, ) self._tracers_halo_updater.update() From a853b9f483c6c178ac0672d3c4729c85c40ced03 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 2 May 2025 14:19:36 -0400 Subject: [PATCH 219/252] Fix precision of PressureAdjustedTemperature test --- .../translate_pressureadjustedtemperature_nonhydrostatic.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py b/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py index 6ed5e02d..7f9afeac 100644 --- a/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py +++ b/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py @@ -1,6 +1,8 @@ +import numpy as np from typing import Any, Dict from ndsl import Namelist, StencilFactory +from ndsl.dsl.typing import Float from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import temperature_adjust from pyFV3.stencils.dyn_core import get_nk_heat_dissipation @@ -42,7 +44,9 @@ def __init__( self.stencil_factory = stencil_factory def compute_from_storage(self, inputs): - inputs["delt_time_factor"] = abs(inputs["bdt"] * self.namelist.delt_max) + inputs["delt_time_factor"] = np.abs( + inputs["bdt"] * self.namelist.delt_max, dtype=Float + ) del inputs["bdt"] self.compute_func(**inputs) return inputs From f1c255944244872de2aac1c78b9f68d58f478459 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 9 May 2025 12:36:22 -0400 Subject: [PATCH 220/252] Switch global reduce to `simplified` version --- pyFV3/stencils/remapping_GEOS.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index fcbd9e41..0bd6b3d8 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -547,11 +547,13 @@ def __call__( inputArray=self._te_2d.view[:] * self._area_64.view[:], communicator=self._comm, stencil_factory=self._stencil_factory, + simplified_reduce=True, ) zsum = mpp_global_sum( inputArray=self._zsum1.view[:] * self._area_64.view[:], communicator=self._comm, stencil_factory=self._stencil_factory, + simplified_reduce=True, ) dtmp = tesum / (CV_AIR * zsum) From 83377460a8d211468172dd9dc9dec08c40035531 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 May 2025 15:42:22 -0400 Subject: [PATCH 221/252] Reset `Tracers` for pyFV3 using `FieldBundle` --- pyFV3/tracers.py | 190 +++++++++++++++-------------------------------- 1 file changed, 59 insertions(+), 131 deletions(-) diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py index 4b938c5f..029847a3 100644 --- a/pyFV3/tracers.py +++ b/pyFV3/tracers.py @@ -1,132 +1,60 @@ -from __future__ import annotations - -from typing import Dict, List - -import numpy as np - -from ndsl import Quantity, QuantityFactory +from ndsl import QuantityFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM - - -# FOR REFERENCE - previous descriptive of the tracers, lining up with Pace work -# tracer_variables = [ -# "qvapor", -# "qliquid", -# "qrain", -# "qice", -# "qsnow", -# "qgraupel", -# "qo3mr", -# "qsgs_tke", -# "qcld", -# ] - - -class Tracers: - unit = "g/kg" - dims = [X_DIM, Y_DIM, Z_DIM] - - def __init__(self, factory: QuantityFactory) -> None: - self._quantities: Dict[str, Quantity] = {} - self._quantity_factory = factory - - def copy_tracer_data( - self, - name: str, - data: np.ndarray, - unit="unknown", - ): - qty = self._quantity_factory.empty(dims=self.dims, units=unit) - if data.shape > qty.data.shape: - raise ValueError( - f"[pyFV3] Tracer {name} size ({data.shape}" - f" is bigger than grid {qty.data.shape})" - ) - qty.data[: data.shape[0], : data.shape[1], : data.shape[2]] = data - self._quantities[name] = qty - - @property - def count(self) -> int: - return len(self._quantities) - - def values(self): - return self._quantities.values() - - def names(self): - return self._quantities.keys() - - def items(self): - return self._quantities.items() - - def as_4D_array(self) -> np.ndarray: - shape = self._quantity_factory.sizer.get_shape(self.dims) - var4d = np.empty( - ( - shape[0] - 1, - shape[1] - 1, - shape[2] - 1, - self.count, - ) - ) - # Skip the extra data point that is meant to align interface - # and non interface fields - for idx, q in enumerate(self.values()): - var4d[:, :, :, idx] = q.data[:-1, :-1, :-1] - return var4d - - def __getitem__(self, key): - return self._quantities[key] - - def __setitem__(self, key, value): - self._quantities[key] = value - - def __str__(self) -> str: - return self.__repr__() - - def __repr__(self) -> str: - msg = f"[pyFV3] {self.count} Tracers:\n" - for q in self.names(): - msg += f" {q}\n" - return msg - - @classmethod - def make( - cls, - quantity_factory: QuantityFactory, - tracer_mapping: List[str], - ): - tracers = cls(quantity_factory) - for name in tracer_mapping: - qty = quantity_factory.empty(dims=cls.dims, units=cls.unit) - tracers._quantities[name] = qty - return tracers - - @staticmethod - def blind_mapping_from_data(tracer_data: np.ndarray): - if len(tracer_data.shape) != 4: - raise ValueError("Expected 4D field as input") - return [f"Tracer_{idx}" for idx in range(tracer_data.shape[3])] - - @classmethod - def make_from_4D_array( - cls, - quantity_factory: QuantityFactory, - tracer_mapping: List[str], - tracer_data: np.ndarray, - ) -> Tracers: - if len(tracer_data.shape) != 4: - raise ValueError("Expected 4D field as input") - count = len(tracer_mapping) - if count > tracer_data.shape[3]: - raise ValueError( - f"Mapping size {len(tracer_mapping)} is bigger than" - f" data dimensionality {tracer_data.shape[3]}" - ) - tracers = cls(quantity_factory) - for idx in range(0, count): - tracers.copy_tracer_data( - name=tracer_mapping[idx] or f"Tracer_{idx}", - data=tracer_data[:, :, :, idx], - unit=cls.unit, - ) - return tracers +from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType +from pyFV3.version import IS_GEOS + +# Defauult maopping for common models +_default_mapping_GEOS = { + "vapor": 0, + "liquid": 1, + "rain": 2, + "snow": 3, + "ice": 4, + "graupel": 5, + "cloud": 6, +} +_default_mapping_PACE = { + "vapor": 0, + "liquid": 1, + "rain": 2, + "ice": 3, + "snow": 4, + "graupel": 5, + "om3r": 6, + "cloud": 7, +} + + +TracersType = FieldBundleType.T("Tracers") + + +def setup_tracers( + number_of_tracers: int, + quantity_factory: QuantityFactory, + mappings: dict[str, int] | None = None, +) -> FieldBundle: + """Setup a FieldBundle for tracers. Should be called only once.""" + + FieldBundleType.register("Tracers", (number_of_tracers,)) + + _unit = "g/kg" + _dims = [X_DIM, Y_DIM, Z_DIM, "tracers"] + + tracers_qty_factory = FieldBundle.extend_3D_quantity_factory( + quantity_factory, {"tracers": number_of_tracers} + ) + data = tracers_qty_factory.zeros(_dims, units=_unit) + + # Some default mappings for ease of use with commonly + # run models + if mappings is None: + if IS_GEOS: + mappings = _default_mapping_GEOS + else: + mappings = _default_mapping_PACE + + return FieldBundle( + "Tracers", + quantity=data, + mapping=mappings, + ) From 1f1279b38c4a4c14bf7724ef09ddf8ddfd71e215 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 May 2025 15:47:37 -0400 Subject: [PATCH 222/252] Update tracers to `FieldBundle` and orchestration --- .gitignore | 1 + pyFV3/stencils/fillz.py | 11 +- pyFV3/stencils/map_single.py | 168 ++++++++++++------ pyFV3/stencils/mapn_tracer.py | 49 ++--- pyFV3/stencils/remapping_GEOS.py | 136 +++++++------- .../translate/translate_MapN_Tracer_2d.py | 88 ++------- .../translate/translate_Pressures_mapU.py | 3 - .../translate/translate_Pressures_mapV.py | 3 - tests/savepoint/translate/translate_fillz.py | 51 ++---- .../translate/translate_map1_ppm_W.py | 1 - .../translate/translate_map1_ppm_delz.py | 1 - .../translate/translate_map_scalar.py | 2 +- .../translate/translate_remapping_GEOS.py | 45 +++-- 13 files changed, 287 insertions(+), 272 deletions(-) diff --git a/.gitignore b/.gitignore index 74cd9b72..c8918d08 100644 --- a/.gitignore +++ b/.gitignore @@ -150,6 +150,7 @@ dmypy.json # GT4Py **/.gt_cache*/ +**/.gt4py_cache/* # Run outputs plot_output/ diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index 0426956e..2a91d033 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -6,7 +6,8 @@ from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, Int, FloatField, FloatFieldIJ, IntFieldIJ -from pyFV3.tracers import Tracers +from pyFV3.tracers import TracersType +import dace @typing.no_type_check @@ -117,8 +118,6 @@ def __init__( self, stencil_factory: StencilFactory, quantity_factory: QuantityFactory, - nq: int, - tracers: Tracers, ): orchestrate( obj=self, @@ -152,16 +151,16 @@ def __init__( def __call__( self, dp2: FloatField, - tracers: Tracers, + tracers: TracersType, ): """ Args: dp2 (in): pressure thickness of atmospheric layer tracers (inout): tracers to fix negative masses in """ - for tracer_name in self._filtered_tracer_dict.keys(): + for i_tracer in dace.nounroll(range(tracers.shape[3])): self._fix_tracer_stencil( - tracers[tracer_name], + tracers.quantity.data[:, :, :, i_tracer], dp2, self._zfix, self._sum0, diff --git a/pyFV3/stencils/map_single.py b/pyFV3/stencils/map_single.py index fc09e817..be7e3808 100644 --- a/pyFV3/stencils/map_single.py +++ b/pyFV3/stencils/map_single.py @@ -87,6 +87,40 @@ def lagrangian_contributions( lev = lev - 1 +class LagrangianContribution: + """Lagrangian contribution as it appears in FV3GFS/SHiELD""" + + def __init__(self, stencil_factory: StencilFactory, dims: Sequence[str]) -> None: + self._lagrangian_contributions = stencil_factory.from_dims_halo( + lagrangian_contributions, + compute_dims=dims, + ) + + def __call__( + self, + q: FloatField, # type: ignore + pe1: FloatField, # type: ignore + pe2: FloatField, # type: ignore + q4_1: FloatField, # type: ignore + q4_2: FloatField, # type: ignore + q4_3: FloatField, # type: ignore + q4_4: FloatField, # type: ignore + dp1: FloatField, # type: ignore + lev: IntFieldIJ, # type: ignore + ): + self._lagrangian_contributions( + q, + pe1, + pe2, + q4_1, + q4_2, + q4_3, + q4_4, + dp1, + lev, + ) + + def lagrangian_contributions_interp( km: int, not_exit_loop: BoolFieldIJ, @@ -220,6 +254,66 @@ def lagrangian_contributions_interp( q = q_temp +class LagrangianContributionInterpolated: + """Lagrangian contribution as it appears in GEOS, modified from original + FV3GFS version""" + + def __init__( + self, + stencil_factory: StencilFactory, + quantity_factory: QuantityFactory, + dims: Sequence[str], + ) -> None: + self._lagrangian_contributions_interp = stencil_factory.from_dims_halo( + lagrangian_contributions_interp, + compute_dims=dims, + ) + + self._INDEX_LM1 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="", + dtype=Int, + ) + + self._INDEX_LP0 = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="", + dtype=Int, + ) + self._km = stencil_factory.grid_indexing.domain[2] + self._not_exit_loop = quantity_factory.zeros( + [X_DIM, Y_DIM], units="", dtype=bool + ) + + def __call__( + self, + q: FloatField, # type: ignore + pe1: FloatField, # type: ignore + pe2: FloatField, # type: ignore + q4_1: FloatField, # type: ignore + q4_2: FloatField, # type: ignore + q4_3: FloatField, # type: ignore + q4_4: FloatField, # type: ignore + dp1: FloatField, # type: ignore + lev: IntFieldIJ, # type: ignore + ): + self._lagrangian_contributions_interp( + km=self._km, + not_exit_loop=self._not_exit_loop, + INDEX_LM1=self._INDEX_LM1, + INDEX_LP0=self._INDEX_LP0, + q=q, + pe1=pe1, + pe2=pe2, + q4_1=q4_1, + q4_2=q4_2, + q4_3=q4_3, + q4_4=q4_4, + dp1=dp1, + lev=lev, + ) + + class MapSingle: """ Fortran name is map_single, test classes are Map1_PPM_2d, Map_Scalar_2d @@ -232,6 +326,7 @@ def __init__( kord: int, mode: int, dims: Sequence[str], + interpolate_contribution: bool = False, ): orchestrate( obj=self, @@ -278,31 +373,14 @@ def make_quantity(): dims=dims, ) - self._lagrangian_contributions = stencil_factory.from_dims_halo( - lagrangian_contributions, - compute_dims=dims, - ) - - self._lagrangian_contributions_interp = stencil_factory.from_dims_halo( - lagrangian_contributions_interp, - compute_dims=dims, - ) - - self._INDEX_LM1 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="", - dtype=Int, - ) - - self._INDEX_LP0 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], - units="", - dtype=Int, - ) - self._km = grid_indexing.domain[2] - self._not_exit_loop = quantity_factory.zeros( - [X_DIM, Y_DIM], units="", dtype=bool - ) + if interpolate_contribution: + self._lagrangian_contributions = LagrangianContributionInterpolated( + stencil_factory, quantity_factory, dims + ) + else: + self._lagrangian_contributions = LagrangianContribution( + stencil_factory, dims + ) @property def i_extent(self): @@ -319,7 +397,6 @@ def __call__( pe2: FloatField, qs: Optional["FloatFieldIJ"] = None, qmin: Float = 0.0, - interp: bool = False, ): """ Compute x-flux using the PPM method. @@ -356,31 +433,14 @@ def __call__( Float(qmin), ) - if interp is False: - self._lagrangian_contributions( - q1, - pe1, - pe2, - self._q4_1, - self._q4_2, - self._q4_3, - self._q4_4, - self._dp1, - self._lev, - ) - else: - self._lagrangian_contributions_interp( - km=self._km, - not_exit_loop=self._not_exit_loop, - INDEX_LM1=self._INDEX_LM1, - INDEX_LP0=self._INDEX_LP0, - q=q1, - pe1=pe1, - pe2=pe2, - q4_1=self._q4_1, - q4_2=self._q4_2, - q4_3=self._q4_3, - q4_4=self._q4_4, - dp1=self._dp1, - lev=self._lev, - ) + self._lagrangian_contributions( + q=q1, + pe1=pe1, + pe2=pe2, + q4_1=self._q4_1, + q4_2=self._q4_2, + q4_3=self._q4_3, + q4_4=self._q4_4, + dp1=self._dp1, + lev=self._lev, + ) diff --git a/pyFV3/stencils/mapn_tracer.py b/pyFV3/stencils/mapn_tracer.py index 79e598c9..3612e412 100644 --- a/pyFV3/stencils/mapn_tracer.py +++ b/pyFV3/stencils/mapn_tracer.py @@ -1,10 +1,10 @@ -import ndsl.dsl.gt4py_utils as utils from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.stencils.map_single import MapSingle -from pyFV3.tracers import Tracers +from pyFV3.tracers import TracersType +from dace import nounroll class MapNTracer: @@ -19,52 +19,51 @@ def __init__( kord: int, nq: int, fill: bool, - tracers: Tracers, + tracers: TracersType, ): orchestrate( obj=self, config=stencil_factory.config.dace_config, dace_compiletime_args=["tracers"], ) - self._nq = int(nq) self._qs = quantity_factory.zeros( [X_DIM, Y_DIM, Z_DIM], units="unknown", dtype=Float, ) - kord_tracer = [kord] * self._nq - kord_tracer[5] = 9 # qcld - - self._list_of_remap_objects = [ - MapSingle( - stencil_factory, - quantity_factory, - kord_tracer[i], - 0, - dims=[X_DIM, Y_DIM, Z_DIM], - ) - for i in range(len(kord_tracer)) - ] + self._map_single = MapSingle( + stencil_factory, + quantity_factory, + kord, + 0, + dims=[X_DIM, Y_DIM, Z_DIM], + ) + self._map_single_kord9 = MapSingle( + stencil_factory, + quantity_factory, + 9, + 0, + dims=[X_DIM, Y_DIM, Z_DIM], + ) if fill: self._fill_negative_tracers = True self._fillz = FillNegativeTracerValues( stencil_factory, quantity_factory, - self._nq, - tracers, ) else: self._fill_negative_tracers = False + self._index_cloud = tracers.index("cloud") + def __call__( self, pe1: FloatField, pe2: FloatField, dp2: FloatField, - # tracers: Tracers, - tracers, + tracers: TracersType, ): """ Remaps the tracer species onto the Eulerian grid @@ -77,8 +76,12 @@ def __call__( dp2 (in): Difference in pressure between Eulerian levels tracers (inout): tracers to be remapped """ - for i, q in enumerate(utils.tracer_variables[0 : self._nq]): - self._list_of_remap_objects[i](tracers[q], pe1, pe2, self._qs) + for i_tracer in nounroll(range(tracers.shape[3])): + if i_tracer != self._index_cloud: + self._map_single( + tracers.quantity.data[:, :, :, i_tracer], pe1, pe2, self._qs + ) + self._map_single_kord9(tracers.cloud, pe1, pe2, self._qs) if self._fill_negative_tracers is True: self._fillz(dp2, tracers) diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 0bd6b3d8..e0bbddcc 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -10,7 +10,7 @@ Z_DIM, Z_INTERFACE_DIM, ) -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldIJ64, FloatFieldK from ndsl.grid import GridData from ndsl.stencils.basic_operations import adjust_divide_stencil from pyFV3._config import RemappingConfig @@ -18,7 +18,7 @@ from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.moist_cv import moist_pt_last_step -from pyFV3.stencils.mpp_global_sum import mpp_global_sum +from pyFV3.stencils.mpp_global_sum import GlobalSum from pyFV3.stencils.remapping import ( CONSV_MIN, init_pe, @@ -32,7 +32,15 @@ from pyFV3.stencils.saturation_adjustment import SatAdjust3d from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment -from pyFV3.tracers import Tracers +from pyFV3.tracers import TracersType + + +def _normalize_to_grid_stencil( + te_2d: FloatFieldIJ, zsum_2d: FloatFieldIJ, area: FloatFieldIJ64 +): + with computation(FORWARD), interval(0, 1): + te_2d = te_2d * area + zsum_2d = zsum_2d * area class LagrangianToEulerian_GEOS: @@ -51,7 +59,7 @@ def __init__( grid_data: GridData, nq, pfull, - tracers: Tracers, + tracers: TracersType, adiabatic: bool, ): orchestrate( @@ -177,6 +185,11 @@ def __init__( # Stencils + self._global_sum = GlobalSum( + communicator=comm, + backend=stencil_factory.backend, + ) + self._init_pe = stencil_factory.from_origin_domain( init_pe, origin=grid_indexing.origin_compute(), @@ -202,6 +215,7 @@ def __init__( self._kord_tm, mode=1, dims=[X_DIM, Y_DIM, Z_DIM], + interpolate_contribution=True, ) self._mapn_tracer = MapNTracer( @@ -210,7 +224,6 @@ def __init__( kord=abs(config.kord_tr), fill=config.fill, tracers=tracers, - exclude_tracers=[], ) self._map_single_w = MapSingle( @@ -333,9 +346,15 @@ def __init__( domain=grid_indexing.domain_compute(), ) + self._normalize_to_grid = stencil_factory.from_origin_domain( + _normalize_to_grid_stencil, + origin=grid_indexing.origin_compute(), + domain=grid_indexing.domain_compute(), + ) + def __call__( self, - tracers: Tracers, + tracers: TracersType, pt: FloatField, # type: ignore delp: FloatField, # type: ignore delz: FloatField, # type: ignore @@ -411,12 +430,12 @@ def __call__( # Build remapping profiles self._init_pe(pe, self._pe1, self._pe2, ptop) self._moist_cv_pt_pressure( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers.vapor, + qliquid=tracers.liquid, + qrain=tracers.rain, + qsnow=tracers.snow, + qice=tracers.ice, + qgraupel=tracers.graupel, q_con=q_con, pt=pt, cappa=cappa, @@ -449,14 +468,13 @@ def __call__( self._pn1, self._pn2, qmin=self._t_min, - interp=True, ) # Map all tracers self._mapn_tracer(self._pe1, self._pe2, self._dp2, tracers) # Map vertical wind - self._map_single_w(w, self._pe1, self._pe2, qs=wsd, interp=False) + self._map_single_w(w, self._pe1, self._pe2, qs=wsd) self._rescale_delz_1(delz, delp) self._map_single_delz(delz, self._pe1, self._pe2) self._rescale_delz_2(delz, self._dp2) @@ -473,14 +491,14 @@ def __call__( # Map horizontal winds, fluxes and courant number self._pressures_mapu(pe, ak, bk, self._pe0, self._pe3, ptop) self._pe0_ptop_xmax(self._pe0, ptop) - self._map_single_u(u, self._pe0, self._pe3, interp=False) - self._map_single_u(mfy, self._pe0, self._pe3, interp=False) - self._map_single_u(cy, self._pe0, self._pe3, interp=False) + self._map_single_u(u, self._pe0, self._pe3) + self._map_single_u(mfy, self._pe0, self._pe3) + self._map_single_u(cy, self._pe0, self._pe3) self._pressures_mapv(pe, ak, bk, self._pe0, self._pe3) - self._map_single_v(v, self._pe0, self._pe3, interp=False) - self._map_single_v(mfx, self._pe0, self._pe3, interp=False) - self._map_single_v(cx, self._pe0, self._pe3, interp=False) + self._map_single_v(v, self._pe0, self._pe3) + self._map_single_v(mfx, self._pe0, self._pe3) + self._map_single_v(cx, self._pe0, self._pe3) self._pe_pk_delp_peln( pe=pe, @@ -497,12 +515,12 @@ def __call__( ) self._moist_cv_pkz( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers.vapor, + qliquid=tracers.liquid, + qrain=tracers.rain, + qsnow=tracers.snow, + qice=tracers.ice, + qgraupel=tracers.graupel, pkz=pkz, pt=pt, cappa=cappa, @@ -515,12 +533,12 @@ def __call__( if last_step: if consv_te > CONSV_MIN: self._moist_cv_te( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers.vapor, + qliquid=tracers.liquid, + qrain=tracers.rain, + qsnow=tracers.snow, + qice=tracers.ice, + qgraupel=tracers.graupel, u=u, v=v, w=w, @@ -543,17 +561,15 @@ def __call__( zsum1=self._zsum1, ) - tesum = mpp_global_sum( - inputArray=self._te_2d.view[:] * self._area_64.view[:], - communicator=self._comm, - stencil_factory=self._stencil_factory, - simplified_reduce=True, + # We can normalize to the same array because + # they are properly reset in the above stencils + self._normalize_to_grid(self._te_2d, self._zsum1, self._area_64) + + tesum: Float = self._global_sum.all_reduce( + array_to_sum=self._te_2d, ) - zsum = mpp_global_sum( - inputArray=self._zsum1.view[:] * self._area_64.view[:], - communicator=self._comm, - stencil_factory=self._stencil_factory, - simplified_reduce=True, + zsum: Float = self._global_sum.all_reduce( + array_to_sum=self._zsum1, ) dtmp = tesum / (CV_AIR * zsum) @@ -570,13 +586,13 @@ def __call__( fast_mp_consv = consv_te > CONSV_MIN self._saturation_adjustment( dp1, - tracers["vapor"], - tracers["liquid"], - tracers["ice"], - tracers["rain"], - tracers["snow"], - tracers["graupel"], - tracers["cloud"], + tracers.vapor, + tracers.liquid, + tracers.ice, + tracers.rain, + tracers.snow, + tracers.graupel, + tracers.cloud, hs, peln, delp, @@ -598,12 +614,12 @@ def __call__( # to the physics, but if we're staying in dynamics we need # to keep it as the virtual potential temperature self._moist_cv_last_step_stencil( - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers.vapor, + qliquid=tracers.liquid, + qrain=tracers.rain, + qsnow=tracers.snow, + qice=tracers.ice, + qgraupel=tracers.graupel, pt=pt, pkz=pkz, dtmp=Float(dtmp), @@ -611,11 +627,11 @@ def __call__( ) self._fill_cond( q_con=q_con, - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qliquid=tracers.liquid, + qrain=tracers.rain, + qsnow=tracers.snow, + qice=tracers.ice, + qgraupel=tracers.graupel, ) else: # converts virtual temperature back to virtual potential temperature diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index b73570be..558da2d3 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -2,6 +2,7 @@ from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.mapn_tracer import MapNTracer +from pyFV3.tracers import setup_tracers class TranslateMapN_Tracer_2d(TranslateFortranData2Py): @@ -12,48 +13,22 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { - "qvapor": { - "kend": grid.npz - 1, - }, - "qliquid": { - "kend": grid.npz - 1, - }, - "qice": { - "kend": grid.npz - 1, - }, - "qrain": { - "kend": grid.npz - 1, - }, - "qsnow": { - "kend": grid.npz - 1, - }, - "qgraupel": { - "kend": grid.npz - 1, - }, - "qcld": { - "kend": grid.npz - 1, - }, - "qo3mr": { - "kend": grid.npz - 1, - }, - "qsgs_tke": { - "kend": grid.npz - 1, - }, - "pe1_": { + "qtracers": {}, + "pe1": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, "kend": grid.npz, }, - "pe2_": { + "pe2": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, "jend": grid.je, "kend": grid.npz, }, - "dp2_": { + "dp2": { "istart": grid.is_, "iend": grid.ie, "jstart": grid.js, @@ -63,33 +38,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto } self.out_vars = { - "qvapor": { - "kend": grid.npz - 1, - }, - "qliquid": { - "kend": grid.npz - 1, - }, - "qice": { - "kend": grid.npz - 1, - }, - "qrain": { - "kend": grid.npz - 1, - }, - "qsnow": { - "kend": grid.npz - 1, - }, - "qgraupel": { - "kend": grid.npz - 1, - }, - "qcld": { - "kend": grid.npz - 1, - }, - "qo3mr": { - "kend": grid.npz - 1, - }, - "qsgs_tke": { - "kend": grid.npz - 1, - }, + "qtracers": {}, } # Value from GEOS @@ -103,32 +52,25 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.fill = True def compute_from_storage(self, inputs): - - tracers = { - "qvapor": inputs["qvapor"], - "qliquid": inputs["qliquid"], - "qice": inputs["qice"], - "qrain": inputs["qrain"], - "qsnow": inputs["qsnow"], - "qgraupel": inputs["qgraupel"], - "qcld": inputs["qcld"], - "qo3mr": inputs["qo3mr"], - "qsgs_tke": inputs["qsgs_tke"], - } + tracers = setup_tracers( + number_of_tracers=inputs["qtracers"].shape[3], + quantity_factory=self.quantity_factory, + mappings={"cloud": 6}, + ) + tracers.quantity.data[:-1, :-1, :-1, :] = inputs["qtracers"] self._compute_func = MapNTracer( self.stencil_factory, self.quantity_factory, abs(self.kord), - self.nq, fill=self.fill, tracers=tracers, ) self._compute_func( - inputs["pe1_"], - inputs["pe2_"], - inputs["dp2_"], + inputs["pe1"], + inputs["pe2"], + inputs["dp2"], tracers, ) diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index f370a858..b2790161 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -145,19 +145,16 @@ def compute_from_storage(self, inputs): inputs["u_"], inputs["pe0_"], inputs["pe3_"], - interp=False, ) self._map1_ppm_u( inputs["mfy_"], inputs["pe0_"], inputs["pe3_"], - interp=False, ) self._map1_ppm_u( inputs["cy_"], inputs["pe0_"], inputs["pe3_"], - interp=False, ) return inputs diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index 147339f6..687243ef 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -136,20 +136,17 @@ def compute_from_storage(self, inputs): inputs["v_"], inputs["pe0_"], inputs["pe3_"], - interp=False, ) self._map1_ppm_v( inputs["mfx_"], inputs["pe0_"], inputs["pe3_"], - interp=False, ) self._map1_ppm_v( inputs["cx_"], inputs["pe0_"], inputs["pe3_"], - interp=False, ) return inputs diff --git a/tests/savepoint/translate/translate_fillz.py b/tests/savepoint/translate/translate_fillz.py index c08b5323..e8ee9482 100644 --- a/tests/savepoint/translate/translate_fillz.py +++ b/tests/savepoint/translate/translate_fillz.py @@ -1,11 +1,11 @@ import numpy as np -import ndsl.dsl.gt4py_utils as utils from ndsl import Namelist, StencilFactory from ndsl.stencils.testing import pad_field_in_j -from ndsl.utils import safe_assign_array -from pyFV3.stencils import fillz +from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.testing import TranslateDycoreFortranData2Py +from pyFV3.tracers import setup_tracers +from ndsl.quantity.field_bundle import FieldBundle class TranslateFillz(TranslateDycoreFortranData2Py): @@ -33,26 +33,28 @@ def __init__( self.max_error = 1e-13 self.ignore_near_zero_errors = {"q2tracers": True} self.stencil_factory = stencil_factory + self._quantity_factory = grid.quantity_factory - def make_storage_data_input_vars(self, inputs, storage_vars=None): - if storage_vars is None: - storage_vars = self.storage_vars() + def make_storage_data_input_vars(self, inputs, tracers: FieldBundle): + storage_vars = self.storage_vars() info = storage_vars["dp2"] inputs["dp2"] = self.make_storage_data( np.squeeze(inputs["dp2"]), istart=info["istart"], axis=info["axis"] ) inputs["tracers"] = {} info = storage_vars["q2tracers"] - for i in range(int(inputs["nq"])): - inputs["tracers"][utils.tracer_variables[i]] = self.make_storage_data( - np.squeeze(inputs["q2tracers"][:, :, i]), - istart=info["istart"], - axis=info["axis"], - ) + tracers.quantity.field[:, :, :, :] = inputs["q2tracers"][:, np.newaxis, :, :] del inputs["q2tracers"] def compute(self, inputs): - self.make_storage_data_input_vars(inputs) + tracers = setup_tracers( + number_of_tracers=inputs["q2tracers"].shape[2], + quantity_factory=self._quantity_factory, + ) + + self.make_storage_data_input_vars(inputs, tracers) + inputs["tracers"] = tracers + for name, value in tuple(inputs.items()): if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: inputs[name] = self.make_storage_data( @@ -60,28 +62,13 @@ def compute(self, inputs): value, self.grid.njd, backend=self.stencil_factory.backend ) ) - for name, value in tuple(inputs["tracers"].items()): - if hasattr(value, "shape") and len(value.shape) > 1 and value.shape[1] == 1: - inputs["tracers"][name] = self.make_storage_data( - pad_field_in_j( - value, self.grid.njd, backend=self.stencil_factory.backend - ) - ) - run_fillz = fillz.FillNegativeTracerValues( + inputs.pop("nq") + fillz = FillNegativeTracerValues( self.stencil_factory, self.grid.quantity_factory, - inputs.pop("nq"), - inputs["tracers"], ) - run_fillz(**inputs) + fillz(**inputs) ds = self.grid.default_domain_dict() ds.update(self.out_vars["q2tracers"]) - tracers = np.zeros((self.grid.nic, self.grid.npz, len(inputs["tracers"]))) - for varname, data in inputs["tracers"].items(): - index = utils.tracer_variables.index(varname) - data[self.grid.slice_dict(ds)] - safe_assign_array( - tracers[:, :, index], np.squeeze(data[self.grid.slice_dict(ds)]) - ) - out = {"q2tracers": tracers} + out = {"q2tracers": tracers.quantity.field[:, 0, :, :]} return out diff --git a/tests/savepoint/translate/translate_map1_ppm_W.py b/tests/savepoint/translate/translate_map1_ppm_W.py index 3948262c..adbb0fc9 100644 --- a/tests/savepoint/translate/translate_map1_ppm_W.py +++ b/tests/savepoint/translate/translate_map1_ppm_W.py @@ -66,6 +66,5 @@ def compute_from_storage(self, inputs): inputs["pe1_"], inputs["pe2_"], qs=inputs["ws_"], - interp=False, ) return inputs diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py index 0da0167e..837ad616 100644 --- a/tests/savepoint/translate/translate_map1_ppm_delz.py +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -110,7 +110,6 @@ def compute_from_storage(self, inputs): inputs["pe1_"], inputs["pe2_"], qs=inputs["gz_"], - interp=False, ) self._rescale_delz_2( diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py index bba68118..e0683376 100644 --- a/tests/savepoint/translate/translate_map_scalar.py +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -55,6 +55,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self._kord_tm, self.mode, dims=[X_DIM, Y_DIM, Z_DIM], + interpolate_contribution=True, ) def compute_from_storage(self, inputs): @@ -63,6 +64,5 @@ def compute_from_storage(self, inputs): inputs["pe1_"], inputs["pe2_"], qmin=inputs["q_min"], - interp=True, ) return inputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index ed2a8947..d1a9e378 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -13,7 +13,7 @@ from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS -from pyFV3.tracers import Tracers +from pyFV3.tracers import TracersType, setup_tracers class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): @@ -372,7 +372,7 @@ def __init__( def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") - def state_from_inputs(self, inputs: dict, tracers: Tracers) -> SimpleNamespace: + def state_from_inputs(self, inputs: dict, tracers: TracersType) -> SimpleNamespace: input_storages = super().state_from_inputs(inputs) # Rename fluxes and courant numbers input_storages["mfx"] = input_storages.pop("mfx_R4") @@ -398,24 +398,39 @@ def outputs_from_state(self, state: dict): else: outputs[name] = state[name] # scalar # Put tracers - storages["tracers"] = state["tracers"].as_4D_array() + storages["tracers"] = state["tracers"].quantity.data[:-1, :-1, :-1, :] outputs.update(self._base.slice_output(storages)) return outputs def compute_parallel(self, inputs, communicator): - tracers_mapping = Tracers.blind_mapping_from_data(inputs["tracers"]) - tracers_mapping[0] = "vapor" - tracers_mapping[1] = "liquid" - tracers_mapping[2] = "rain" - tracers_mapping[3] = "snow" - tracers_mapping[4] = "ice" - tracers_mapping[5] = "graupel" - tracers_mapping[6] = "cloud" - tracers = Tracers.make_from_4D_array( - self.quantity_factory, - tracers_mapping, - inputs["tracers"], + # tracers_mapping = Tracers.blind_mapping_from_data(inputs["tracers"]) + # tracers_mapping[0] = "vapor" + # tracers_mapping[1] = "liquid" + # tracers_mapping[2] = "rain" + # tracers_mapping[3] = "snow" + # tracers_mapping[4] = "ice" + # tracers_mapping[5] = "graupel" + # tracers_mapping[6] = "cloud" + # tracers = Tracers.make_from_4D_array( + # self.quantity_factory, + # tracers_mapping[0:7], + # inputs["tracers"], + # ) + + tracers = setup_tracers( + number_of_tracers=inputs["tracers"].shape[3], + quantity_factory=self.quantity_factory, + mappings={ + "vapor": 0, + "liquid": 1, + "rain": 2, + "snow": 3, + "ice": 4, + "graupel": 5, + "cloud": 6, + }, ) + tracers.quantity.data[:-1, :-1, :-1, :] = inputs["tracers"] inputs["te_2d"] = inputs["te_2d"].astype(Float) state = self.state_from_inputs(inputs, tracers) From 124652570fe18901c3e01828a824a612326b1807 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 May 2025 15:50:11 -0400 Subject: [PATCH 223/252] Update README.md --- README.md | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 728dd95c..a961837b 100644 --- a/README.md +++ b/README.md @@ -37,14 +37,15 @@ Branches: - ⚙️ `fix/GEOSv11_4_2/RiemanSolver`@Florian: Fix for 32-bit RiemanSolver - ⚙️ `fix/GEOSv11_4_2/C_SW`@Florian: Fix for C_SW for 32-bit - ⚙️ `fix/GEOSv11_4_2/Dyncore`@Florian: Fix for Acoustics and DycoreState for 32-bit and `dpx` calculation - - MERGE ORDER: after `fix/GEOS/D_SW` - - MERGE ORDER: after `fix/GEOSv11_4_2/HyperDiffusionDamping` + - MERGE ORDER: after `fix/GEOS/D_SW` + - MERGE ORDER: after `fix/GEOSv11_4_2/HyperDiffusionDamping` - ⚙️ `feature/tracer_rework_part1` @Florian: Allow for update of N Tracers - ⚙️ `fix/GEOS/TracerAdvection` @Florian: Allow for non-update of mass fluxes and courant number, f32 fixes, correct computation of `cmax` and `nsplit`, overcomputation into the algorithm - - BASED ON `tracer_rework_part1` - - REQUIRES: `ndsl` with tracer rework + - BASED ON `tracer_rework_part1` + - REQUIRES: `ndsl` with tracer rework - ⚙️ `feature/fv_mapz/GEOS` @ Chris K: Remapping for GEOS + - REQUIRES: `ndsl` with tracer rework - ⚙️ `fix/GEOSv11_4_2/Dynamics`@Florian: Fix for the f32 & GEOS version of dynamics - - REQUIRES: `ndsl` with tracer rework - - REQUIRES: `tracer_rework_part1`, `fix/GEOSv11_4_2/Dyncore`, `fix/GEOS/TracerAdvection` - - MERGE ORDER: after `fix/GEOSv11_4_2/HyperDiffusionDamping` + - REQUIRES: `ndsl` with tracer rework + - REQUIRES: `tracer_rework_part1`, `fix/GEOSv11_4_2/Dyncore`, `fix/GEOS/TracerAdvection` + - MERGE ORDER: after `fix/GEOSv11_4_2/HyperDiffusionDamping` From 30f79595695331a97c67348ae99b5b7ac10c604e Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 May 2025 16:32:32 -0400 Subject: [PATCH 224/252] Move MPP global sum to sub-modules --- .../mpp_global_sum.py => mpi/mpp_sum.py} | 142 ++++++++---------- pyFV3/mpi/sum.py | 16 ++ pyFV3/stencils/__init__.py | 1 - pyFV3/stencils/remapping.py | 11 +- pyFV3/stencils/remapping_GEOS.py | 12 +- 5 files changed, 90 insertions(+), 92 deletions(-) rename pyFV3/{stencils/mpp_global_sum.py => mpi/mpp_sum.py} (54%) create mode 100644 pyFV3/mpi/sum.py diff --git a/pyFV3/stencils/mpp_global_sum.py b/pyFV3/mpi/mpp_sum.py similarity index 54% rename from pyFV3/stencils/mpp_global_sum.py rename to pyFV3/mpi/mpp_sum.py index e5ad9c5a..0a0ea019 100644 --- a/pyFV3/stencils/mpp_global_sum.py +++ b/pyFV3/mpi/mpp_sum.py @@ -1,82 +1,10 @@ import numpy as np - -from ndsl.comm.comm_abc import ReductionOperator +from ndsl import Quantity, StencilFactory from ndsl.dsl.typing import Float -from ndsl.quantity import Quantity - - -def mpp_global_sum( - inputArray, communicator, stencil_factory=None, simplified_reduce=False -): - if not simplified_reduce: - NUMINT = 6 - NUMBIT = 46 - r_prec = 2.0 ** NUMBIT - prec = 2 ** NUMBIT - I_prec = 1.0 / (2.0 ** NUMBIT) - pr = [ - r_prec ** 2, - r_prec, - 1.0, - 1.0 / r_prec, - 1.0 / r_prec ** 2, - 1.0 / r_prec ** 3, - ] - I_pr = [1.0 / r_prec ** 2, 1.0 / r_prec, 1.0, r_prec, r_prec ** 2, r_prec ** 3] - prec_error = (2 ** 62 + (2 ** 62 - 1)) / 6 - mag_max_term = 0.0 - - ints_sum = Quantity( - data=np.zeros((NUMINT), dtype=Float), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - - ints_sum_reduce = Quantity( - data=np.zeros((NUMINT), dtype=Float), - dims=["K"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - - # Note: This loop range in i and j are for the TBC test case. - for j in range(inputArray.shape[1]): - for i in range(inputArray.shape[0]): - increment_ints_faster( - ints_sum.data, pr, I_pr, inputArray[i, j], mag_max_term - ) - - carry_overflow(ints_sum.data, prec, I_prec, prec_error) +from ndsl.comm.communicator import Communicator, ReductionOperator - communicator.all_reduce(ints_sum, ReductionOperator.SUM, ints_sum_reduce) - regularize_ints(ints_sum_reduce.data, prec, I_prec) - - sum_ = ints_to_real(ints_sum_reduce.data, pr) - else: - inputArray_ = Quantity( - data=inputArray.astype(Float), - dims=["I", "J"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - - ints_sum_reduce = Quantity( - data=np.zeros(inputArray.data.shape, dtype=Float), - dims=["I", "J"], - units="dunno", - gt4py_backend=stencil_factory.backend, - ) - - communicator.all_reduce(inputArray_, ReductionOperator.SUM, ints_sum_reduce) - - sum_ = sum(sum(ints_sum_reduce.data)) - - return sum_ - - -def increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): +def _increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): if (r >= 1e30) == r < 1e30: print("NaN_error") return @@ -94,7 +22,7 @@ def increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): int_sum[i] = int_sum[i] + sgn * ival -def carry_overflow(int_sum, prec, I_prec, prec_error): +def _carry_overflow(int_sum, prec, I_prec, prec_error): for i in range(len(int_sum) - 1, 0, -1): if abs(int_sum[i]) > prec: num_carry = int(int_sum[i] * I_prec) @@ -104,7 +32,7 @@ def carry_overflow(int_sum, prec, I_prec, prec_error): overflow_error = True -def regularize_ints(int_sum, prec, I_prec): +def _regularize_ints(int_sum, prec, I_prec): for i in range(len(int_sum) - 1, 0, -1): if abs(int_sum[i]) > prec: num_carry = int(int_sum[i] * I_prec) @@ -132,10 +60,68 @@ def regularize_ints(int_sum, prec, I_prec): int_sum[i - 1] = int_sum[i - 1] + 1 -def ints_to_real(ints, pr): +def _ints_to_real(ints, pr): r = 0.0 for i in range(len(ints)): r = r + pr[i] * ints[i] return r + + +class MPPGlobalSum: + def __init__( + self, stencil_factory: StencilFactory, communicator: Communicator + ) -> None: + NUMINT = 6 + self._comm = communicator + self._ints_sum = Quantity( + data=np.zeros((NUMINT), dtype=Float), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + self._ints_sum_reduce = Quantity( + data=np.zeros((NUMINT), dtype=Float), + dims=["K"], + units="dunno", + gt4py_backend=stencil_factory.backend, + ) + + def __call__(self, qty_to_sum: Quantity) -> Float: + NUMBIT = 46 + r_prec = 2.0**NUMBIT + prec = 2**NUMBIT + I_prec = 1.0 / (2.0**NUMBIT) + pr = [ + r_prec**2, + r_prec, + 1.0, + 1.0 / r_prec, + 1.0 / r_prec**2, + 1.0 / r_prec**3, + ] + I_pr = [1.0 / r_prec**2, 1.0 / r_prec, 1.0, r_prec, r_prec**2, r_prec**3] + prec_error = (2**62 + (2**62 - 1)) / 6 + mag_max_term = 0.0 + + # Note: This loop range in i and j are for the TBC test case. + self._ints_sum.data[:] = 0 + for j in range(qty_to_sum.field.shape[1]): + for i in range(qty_to_sum.field.shape[0]): + _increment_ints_faster( + self._ints_sum.data, pr, I_pr, qty_to_sum.field[i, j], mag_max_term + ) + + _carry_overflow(self._ints_sum.data, prec, I_prec, prec_error) + + self._comm.all_reduce( + self._ints_sum, + ReductionOperator.SUM, + self._ints_sum_reduce, + ) + + _regularize_ints(self._ints_sum_reduce.data, prec, I_prec) + + return _ints_to_real(self._ints_sum_reduce.data, pr) diff --git a/pyFV3/mpi/sum.py b/pyFV3/mpi/sum.py new file mode 100644 index 00000000..d020bcf5 --- /dev/null +++ b/pyFV3/mpi/sum.py @@ -0,0 +1,16 @@ +from ndsl import Quantity, QuantityFactory +from ndsl.dsl.typing import Float +from ndsl.comm.communicator import Communicator, ReductionOperator +from ndsl.constants import X_DIM, Y_DIM + + +class GlobalSum: + def __init__( + self, quantity_factory: QuantityFactory, communicator: Communicator + ) -> None: + self._comm = communicator + self._tmp_reduce = quantity_factory.empty(dims=[X_DIM, Y_DIM], units="n/a") + + def __call__(self, qty_to_sum: Quantity) -> Float: + self._comm.all_reduce(qty_to_sum, ReductionOperator.SUM, self._tmp_reduce) + return qty_to_sum.field.sum(axis=0).sum(axis=1) diff --git a/pyFV3/stencils/__init__.py b/pyFV3/stencils/__init__.py index d5f5178d..dbf4f1e3 100644 --- a/pyFV3/stencils/__init__.py +++ b/pyFV3/stencils/__init__.py @@ -13,7 +13,6 @@ from .fxadv import FiniteVolumeFluxPrep from .map_single import MapSingle from .mapn_tracer import MapNTracer -from .mpp_global_sum import mpp_global_sum from .neg_adj3 import AdjustNegativeTracerMixingRatio from .nh_p_grad import NonHydrostaticPressureGradient from .pk3_halo import PK3Halo diff --git a/pyFV3/stencils/remapping.py b/pyFV3/stencils/remapping.py index 4174493b..c396c58a 100644 --- a/pyFV3/stencils/remapping.py +++ b/pyFV3/stencils/remapping.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional +from typing import List, Optional, no_type_check from gt4py.cartesian.gtscript import ( BACKWARD, @@ -30,6 +30,7 @@ from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.moist_cv import moist_pt_func, moist_pt_last_step from pyFV3.stencils.saturation_adjustment import SatAdjust3d +from pyFV3.tracers import TracersType # from pyFV3.tracers import Tracers @@ -318,7 +319,6 @@ def pe_pk_delp_peln( akap: Float, ptop: Float, ): - with computation(BACKWARD): with interval(-1, None): pe_bottom = pe @@ -364,8 +364,8 @@ def __init__( area_64, nq, pfull, - # tracers: Tracers, - tracers, + tracers: TracersType, + exclude_tracers: List[str], checkpointer: Optional[Checkpointer] = None, ): orchestrate( @@ -595,9 +595,10 @@ def __init__( domain=grid_indexing.domain_compute(), ) + @no_type_check def __call__( self, - tracers: Dict[str, Quantity], + tracers: TracersType, pt: FloatField, delp: FloatField, delz: FloatField, diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index e0bbddcc..fd551767 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -18,7 +18,7 @@ from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.moist_cv import moist_pt_last_step -from pyFV3.stencils.mpp_global_sum import GlobalSum +from pyFV3.mpi.sum import GlobalSum from pyFV3.stencils.remapping import ( CONSV_MIN, init_pe, @@ -187,7 +187,7 @@ def __init__( self._global_sum = GlobalSum( communicator=comm, - backend=stencil_factory.backend, + quantity_factory=quantity_factory, ) self._init_pe = stencil_factory.from_origin_domain( @@ -565,12 +565,8 @@ def __call__( # they are properly reset in the above stencils self._normalize_to_grid(self._te_2d, self._zsum1, self._area_64) - tesum: Float = self._global_sum.all_reduce( - array_to_sum=self._te_2d, - ) - zsum: Float = self._global_sum.all_reduce( - array_to_sum=self._zsum1, - ) + tesum: Float = self._global_sum(self._te_2d) + zsum: Float = self._global_sum(self._zsum1) dtmp = tesum / (CV_AIR * zsum) elif consv_te < -CONSV_MIN: From 2daf7bfb03f66e2f0e337457188f2dc018828db4 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 May 2025 16:34:09 -0400 Subject: [PATCH 225/252] Orchestrate MPP via callback --- pyFV3/mpi/sum.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyFV3/mpi/sum.py b/pyFV3/mpi/sum.py index d020bcf5..27c62a61 100644 --- a/pyFV3/mpi/sum.py +++ b/pyFV3/mpi/sum.py @@ -2,6 +2,7 @@ from ndsl.dsl.typing import Float from ndsl.comm.communicator import Communicator, ReductionOperator from ndsl.constants import X_DIM, Y_DIM +from ndsl.dsl.dace.orchestration import dace_inhibitor class GlobalSum: @@ -11,6 +12,7 @@ def __init__( self._comm = communicator self._tmp_reduce = quantity_factory.empty(dims=[X_DIM, Y_DIM], units="n/a") + @dace_inhibitor def __call__(self, qty_to_sum: Quantity) -> Float: self._comm.all_reduce(qty_to_sum, ReductionOperator.SUM, self._tmp_reduce) return qty_to_sum.field.sum(axis=0).sum(axis=1) From 1267b2a60c516a6ce0e5409c5f5497ca69394281 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 May 2025 17:12:31 -0400 Subject: [PATCH 226/252] Fix imports on translate Fix sum on 2 axis --- pyFV3/mpi/sum.py | 3 ++- tests/savepoint/translate/translate_mpp_global_sum.py | 10 +++------- tests/savepoint/translate/translate_remapping.py | 1 + 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/pyFV3/mpi/sum.py b/pyFV3/mpi/sum.py index 27c62a61..8bbbd976 100644 --- a/pyFV3/mpi/sum.py +++ b/pyFV3/mpi/sum.py @@ -14,5 +14,6 @@ def __init__( @dace_inhibitor def __call__(self, qty_to_sum: Quantity) -> Float: + assert qty_to_sum.field.shape == 2 # Code handle only 2D quantity self._comm.all_reduce(qty_to_sum, ReductionOperator.SUM, self._tmp_reduce) - return qty_to_sum.field.sum(axis=0).sum(axis=1) + return qty_to_sum.field.sum(axis=0).sum(axis=0) diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index 0980b7f1..54c57d82 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -2,7 +2,7 @@ from ndsl.stencils.testing import ParallelTranslate from ndsl.stencils.testing.grid import Grid from ndsl.typing import Communicator -from pyFV3.stencils.mpp_global_sum import mpp_global_sum +from pyFV3.mpi.mpp_sum import MPPGlobalSum class TranslateMpp_global_sum(ParallelTranslate): @@ -29,11 +29,7 @@ def __init__( self._base.out_vars = {"tesum": {}} def compute_parallel(self, inputs, communicator: Communicator): - - inputs["tesum"] = mpp_global_sum( - inputs["inputArray"], - communicator, - self.stencil_factory, - ) + mpp_sum = MPPGlobalSum(self.stencil_factory, communicator) + inputs["tesum"] = mpp_sum(inputs["inputArray"]) return inputs diff --git a/tests/savepoint/translate/translate_remapping.py b/tests/savepoint/translate/translate_remapping.py index 43ddb27d..66adf6f4 100644 --- a/tests/savepoint/translate/translate_remapping.py +++ b/tests/savepoint/translate/translate_remapping.py @@ -4,6 +4,7 @@ from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import LagrangianToEulerian from pyFV3.testing import TranslateDycoreFortranData2Py +from pyFV3.tracers import TracersType class TranslateRemapping(TranslateDycoreFortranData2Py): From 36b053ebd365dc72b2684eb4651711dfe6d64147 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 19 May 2025 17:13:21 -0400 Subject: [PATCH 227/252] Bad merge --- pyFV3/stencils/fillz.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/pyFV3/stencils/fillz.py b/pyFV3/stencils/fillz.py index 2a91d033..42b157b3 100644 --- a/pyFV3/stencils/fillz.py +++ b/pyFV3/stencils/fillz.py @@ -2,7 +2,6 @@ from gt4py.cartesian.gtscript import BACKWARD, FORWARD, PARALLEL, computation, interval -import ndsl.dsl.gt4py_utils as utils from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, Int, FloatField, FloatFieldIJ, IntFieldIJ @@ -124,7 +123,6 @@ def __init__( config=stencil_factory.config.dace_config, dace_compiletime_args=["tracers"], ) - self._nq = int(nq) self._fix_tracer_stencil = stencil_factory.from_dims_halo( fix_tracer, compute_dims=[X_DIM, Y_DIM, Z_DIM], @@ -144,10 +142,6 @@ def __init__( dtype=Float, ) - self._filtered_tracer_dict = { - name: tracers[name] for name in utils.tracer_variables[0 : self._nq] - } - def __call__( self, dp2: FloatField, From 528270764fc9b16ec9d0622e74323cfdc250f2cd Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 20 May 2025 12:22:21 -0400 Subject: [PATCH 228/252] Move Tracers2D1L to FieldBundle tracers Orchestrate --- pyFV3/stencils/tracer_2d_1l.py | 107 +++++++++--------- .../translate/translate_tracer2d1l.py | 20 ++-- 2 files changed, 60 insertions(+), 67 deletions(-) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index 61dbe61c..c3313d52 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -1,4 +1,5 @@ -from typing import List +import dace +from typing import no_type_check import numpy as np import gt4py.cartesian.gtscript as gtscript @@ -27,11 +28,11 @@ Y_INTERFACE_DIM, Z_DIM, ) +from ndsl.dsl.dace.orchestration import dace_inhibitor from ndsl.dsl.typing import FloatField, FloatFieldIJ, FloatFieldK from ndsl.comm.communicator import Communicator, ReductionOperator from pyFV3.stencils.fvtp2d import FiniteVolumeTransport -from pyFV3.tracers import Tracers -from ndsl.dsl.gt4py_utils import asarray +from pyFV3.tracers import TracersType from ndsl.utils import safe_assign_array @@ -57,6 +58,7 @@ def flux_y(cy, dya, dx, sin_sg4, sin_sg2, yfx): return yfx +@no_type_check def flux_compute( cx: FloatField, cy: FloatField, @@ -91,6 +93,7 @@ def flux_compute( yfx = flux_y(cy, dya, dx, sin_sg4, sin_sg2, yfx) +@no_type_check def divide_fluxes_by_n_substeps( cxd: FloatField, xfx: FloatField, @@ -124,6 +127,7 @@ def divide_fluxes_by_n_substeps( mfyd = mfyd * frac +@no_type_check def apply_mass_flux( dp1: FloatField, x_mass_flux: FloatField, @@ -150,6 +154,7 @@ def apply_mass_flux( ) +@no_type_check def apply_tracer_flux( q: FloatField, dp1: FloatField, @@ -179,6 +184,7 @@ def apply_tracer_flux( # dp1[:] = dp2 # dp2[:] = self._tmp_dp2 # Because dpX can be a quantity or an array +@no_type_check def swap_dp(dp1: FloatField, dp2: FloatField): with computation(PARALLEL), interval(...): tmp = dp1 @@ -210,8 +216,7 @@ def __init__( transport: FiniteVolumeTransport, grid_data: GridData, comm: Communicator, - tracers: Tracers, - exclude_tracers: List[str], + tracers: TracersType, update_mass_courant: bool = True, ): orchestrate( @@ -222,7 +227,6 @@ def __init__( grid_indexing = stencil_factory.grid_indexing self.grid_indexing = grid_indexing # needed for selective validation self.grid_data = grid_data - self._exclude_tracers = exclude_tracers self._update_mass_courant = update_mass_courant if not self._update_mass_courant: @@ -313,23 +317,11 @@ def __init__( ) self.finite_volume_transport: FiniteVolumeTransport = transport - # Setup halo updater for tracers - tracer_halo_spec = quantity_factory.get_quantity_halo_spec( - dims=[X_DIM, Y_DIM, Z_DIM], - n_halo=N_HALO_DEFAULT, - ) - - # We can exclude tracers from advecting and therefore also - # halo exchanging - advected_tracers = {} - for name, tracer in tracers.items(): - if name in exclude_tracers: - continue - advected_tracers[name] = tracer + # Halo exchange of all tracers self._tracers_halo_updater = WrappedHaloUpdater( - comm.get_scalar_halo_updater([tracer_halo_spec] * len(advected_tracers)), - advected_tracers, - [t for t in advected_tracers.keys()], + comm.get_scalar_halo_updater([tracers.quantity.halo_spec(N_HALO_DEFAULT)]), + {"tracers": tracers.quantity}, + ["tracers"], ) # Setup tracer courant max reduction calculation @@ -342,7 +334,7 @@ def __init__( def __call__( self, - tracers: Tracers, + tracers: TracersType, dp1, x_mass_flux, y_mass_flux, @@ -425,10 +417,8 @@ def __call__( # a loop on the highest number of nsplit, but restraining # actual update in `apply_tracer_flux` to only the valid # K level for each tracers - cmax_on_host = asarray(self._cmax.view[:], to_type=np.ndarray) - max_n_split = i32(1.0 + cmax_on_host.max()) - - for current_nsplit in range(int(max_n_split)): + max_n_split = int(1.0 + self._compute_cmax.max_over_column) + for current_nsplit in range(max_n_split): last_call = current_nsplit == max_n_split - 1 # tracer substep self._apply_mass_flux( @@ -438,31 +428,29 @@ def __call__( self.grid_data.rarea, dp2, ) - for name, q in tracers.items(): - if name in self._exclude_tracers: - pass - else: - self.finite_volume_transport( - q, - working_x_courant, - working_y_courant, - self._x_area_flux, - self._y_area_flux, - self._x_flux, - self._y_flux, - x_mass_flux=x_mass_flux, - y_mass_flux=y_mass_flux, - ) - self._apply_tracer_flux( - q, - dp1, - self._x_flux, - self._y_flux, - self.grid_data.rarea, - dp2, - cmax=self._cmax, - current_nsplit=current_nsplit, - ) + for i_tracer in dace.nounroll(range(tracers.shape[3])): + q = tracers.quantity.data[:, :, :, i_tracer] + self.finite_volume_transport( + q, + working_x_courant, + working_y_courant, + self._x_area_flux, + self._y_area_flux, + self._x_flux, + self._y_flux, + x_mass_flux=x_mass_flux, + y_mass_flux=y_mass_flux, + ) + self._apply_tracer_flux( + q, + dp1, + self._x_flux, + self._y_flux, + self.grid_data.rarea, + dp2, + cmax=self._cmax, + current_nsplit=current_nsplit, + ) if not last_call: self._tracers_halo_updater.update() # we can't use variable assignment to avoid a data copy @@ -470,6 +458,7 @@ def __call__( self._swap_dp(dp1, dp2) +@no_type_check def cmax_stencil_low_k( cx: FloatField, cy: FloatField, @@ -479,6 +468,7 @@ def cmax_stencil_low_k( cmax = max(abs(cx), abs(cy)) +@no_type_check def cmax_stencil_high_k( cx: FloatField, cy: FloatField, @@ -501,6 +491,7 @@ def __init__( The maximum courant number for every atmospheric level on the entire grid. """ + orchestrate(obj=self, config=stencil_factory.config.dace_config) self._grid_data = grid_data self._comm = comm grid_indexing = stencil_factory.grid_indexing @@ -535,8 +526,15 @@ def __init__( [Z_DIM], units="unknown", ) + self.max_over_column = 0 - def __call__(self, cx: Quantity, cy: Quantity, cmax: Quantity): + @dace_inhibitor + def _reduce(self, cmax: Quantity): + cmax.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] + self._comm.all_reduce_per_element_in_place(cmax, ReductionOperator.MAX) + self.max_over_column = cmax.field.max() + + def __call__(self, cx, cy, cmax: Quantity): if __debug__: if not isinstance(cmax, Quantity): raise TypeError( @@ -553,5 +551,4 @@ def __call__(self, cx: Quantity, cy: Quantity, cmax: Quantity): sin_sg5=self._grid_data.sin_sg5, cmax=self._tmp_cmax, ) - cmax.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] - self._comm.all_reduce_per_element_in_place(cmax, ReductionOperator.MAX) + self._reduce(cmax) diff --git a/tests/savepoint/translate/translate_tracer2d1l.py b/tests/savepoint/translate/translate_tracer2d1l.py index 90cfd600..ecc658d3 100644 --- a/tests/savepoint/translate/translate_tracer2d1l.py +++ b/tests/savepoint/translate/translate_tracer2d1l.py @@ -4,7 +4,7 @@ from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import ParallelTranslate from pyFV3.stencils import FiniteVolumeTransport, TracerAdvection -from pyFV3.tracers import Tracers +from pyFV3.tracers import TracersType, setup_tracers from pyFV3.utils.functional_validation import get_subset_func @@ -45,17 +45,13 @@ def __init__( n_halo=((0, 0), (0, 0)), ) - def collect_input_data(self, serializer, savepoint): - input_data = self._base.collect_input_data(serializer, savepoint) - return input_data - def compute_parallel(self, inputs, communicator): - tracers = Tracers.make_from_4D_array( + self._base.make_storage_data_input_vars(inputs, dict_4d=False) + tracers = setup_tracers( + number_of_tracers=inputs["tracers"].shape[3], quantity_factory=self._quantity_factory, - tracer_mapping=Tracers.blind_mapping_from_data(inputs["tracers"]), - tracer_data=inputs["tracers"], ) - self._base.make_storage_data_input_vars(inputs, dict_4d=False) + tracers.quantity.data[:] = inputs["tracers"][:] inputs.pop("tracers") inputs.pop("nq") # Fortran NQ is intrinsic to Tracers (e.g Tracers.count) transport = FiniteVolumeTransport( @@ -74,7 +70,6 @@ def compute_parallel(self, inputs, communicator): self.grid.grid_data, communicator, tracers, - exclude_tracers=["cloud"], update_mass_courant=False, ) inputs["x_mass_flux"] = inputs.pop("mfxd_R4") @@ -86,9 +81,10 @@ def compute_parallel(self, inputs, communicator): inputs["mfyd_R4"] = inputs.pop("y_mass_flux") inputs["cxd_R4"] = inputs.pop("x_courant") inputs["cyd_R4"] = inputs.pop("y_courant") - inputs["tracers"] = tracers.as_4D_array() + inputs["tracers"] = tracers.quantity.field outputs = self._base.slice_output(inputs) - outputs["tracers"] = self.subset_output("tracers", outputs["tracers"]) + # outputs["tracers"] = self.subset_output("tracers", outputs["tracers"]) + # outputs["tracers"] = tracers.quantity.field[:] return outputs def compute_sequential(self, a, b): From 7ea01d175a776dbbb86ee00f609075e034839dc5 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Tue, 20 May 2025 12:23:45 -0400 Subject: [PATCH 229/252] Help `mypy` by going to `TypeAlias` --- pyFV3/tracers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py index 029847a3..30a17bc1 100644 --- a/pyFV3/tracers.py +++ b/pyFV3/tracers.py @@ -1,3 +1,4 @@ +from typing import TypeAlias from ndsl import QuantityFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType @@ -25,7 +26,7 @@ } -TracersType = FieldBundleType.T("Tracers") +TracersType: TypeAlias = FieldBundleType.T("Tracers") # type: ignore def setup_tracers( From 8d8b912a3b0561ff2089292557c71909172e0ce1 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 26 Jun 2025 08:28:33 -0400 Subject: [PATCH 230/252] Fix liquid tracers order for GEOS --- pyFV3/tracers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pyFV3/tracers.py b/pyFV3/tracers.py index 30a17bc1..0407c5e6 100644 --- a/pyFV3/tracers.py +++ b/pyFV3/tracers.py @@ -8,9 +8,9 @@ _default_mapping_GEOS = { "vapor": 0, "liquid": 1, - "rain": 2, - "snow": 3, - "ice": 4, + "ice": 2, + "rain": 3, + "snow": 4, "graupel": 5, "cloud": 6, } From ac0c83f2298064a53004bc2ce9629c6d28c50fb8 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 26 Jun 2025 08:28:54 -0400 Subject: [PATCH 231/252] Orchestrate Tracers --- pyFV3/stencils/tracer_2d_1l.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index c3313d52..b5c86363 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -363,10 +363,10 @@ def __call__( working_x_courant = x_courant working_y_courant = y_courant else: - safe_assign_array(self._tmp_mfx.data, x_mass_flux) - safe_assign_array(self._tmp_mfy.data, y_mass_flux) - safe_assign_array(self._tmp_cx.data, x_courant) - safe_assign_array(self._tmp_cy.data, y_courant) + self._tmp_mfx.data = x_mass_flux + self._tmp_mfy.data = y_mass_flux + self._tmp_cx.data = x_courant + self._tmp_cy.data = y_courant working_x_mass_flux = self._tmp_mfx working_y_mass_flux = self._tmp_mfy working_x_courant = self._tmp_cx @@ -530,16 +530,16 @@ def __init__( @dace_inhibitor def _reduce(self, cmax: Quantity): - cmax.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] - self._comm.all_reduce_per_element_in_place(cmax, ReductionOperator.MAX) - self.max_over_column = cmax.field.max() - - def __call__(self, cx, cy, cmax: Quantity): if __debug__: if not isinstance(cmax, Quantity): raise TypeError( f"[pyFV3][Tracer]: cmax must be a quantity, got {type(cmax)}" ) + cmax.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] + self._comm.all_reduce_per_element_in_place(cmax, ReductionOperator.MAX) + self.max_over_column = cmax.field.max() + + def __call__(self, cx, cy, cmax): self._cmax_low_k( cx=cx, cy=cy, From 1b01b2f40b460130dc5f5fe568289e2053f95340 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 26 Jun 2025 08:33:26 -0400 Subject: [PATCH 232/252] Fix DycoreState for tracers, translate test Remove workaround for Float(x) in orchestration --- pyFV3/dycore_state.py | 127 ++++++++++++++++++++++----- pyFV3/stencils/divergence_damping.py | 4 +- pyFV3/stencils/dyn_core.py | 22 +---- pyFV3/testing/translate_dyncore.py | 1 + 4 files changed, 113 insertions(+), 41 deletions(-) diff --git a/pyFV3/dycore_state.py b/pyFV3/dycore_state.py index 9fdeaf12..aee2b206 100644 --- a/pyFV3/dycore_state.py +++ b/pyFV3/dycore_state.py @@ -16,6 +16,66 @@ from ndsl.dsl.typing import Float from ndsl.restart._legacy_restart import open_restart from ndsl.typing import Communicator +from pyFV3.tracers import TracersType, setup_tracers +from ndsl.quantity.field_bundle import FieldBundle + + +DEFAULT_TRACER_PROPERTIES = { + "specific_humidity": { + "pyFV3_key": "vapor", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "sphum", + "units": "g/kg", + }, + "cloud_liquid_water_mixing_ratio": { + "pyFV3_key": "liquid", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "liq_wat", + "units": "g/kg", + }, + "cloud_ice_mixing_ratio": { + "pyFV3_key": "ice", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "ice_wat", + "units": "g/kg", + }, + "rain_mixing_ratio": { + "pyFV3_key": "rain", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "rainwat", + "units": "g/kg", + }, + "snow_mixing_ratio": { + "pyFV3_key": "snow", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "snowwat", + "units": "g/kg", + }, + "graupel_mixing_ratio": { + "pyFV3_key": "graupel", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "graupel", + "units": "g/kg", + }, + "ozone_mixing_ratio": { + "pyFV3_key": "o3mr", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "o3mr", + "units": "g/kg", + }, + "turbulent_kinetic_energy": { + "pyFV3_key": "sgs_tke", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "sgs_tke", + "units": "g/kg", + }, + "cloud_fraction": { + "pyFV3_key": "cloud", + "dims": [Z_DIM, Y_DIM, X_DIM], + "restart_name": "cld_amt", + "units": "g/kg", + }, +} @dataclass() @@ -148,7 +208,7 @@ class DycoreState: "intent": "inout", } ) - qvapor: Quantity = field( + tracers: TracersType = field( metadata={ "name": "specific_humidity", "dims": [X_DIM, Y_DIM, Z_DIM], @@ -313,6 +373,7 @@ def __post_init__(self): def init_zeros( cls, quantity_factory: QuantityFactory, + tracer_count: int, dtype_dict: Optional[Dict[str, type]] = None, allow_mismatch_float_precision=False, ): @@ -327,9 +388,18 @@ def init_zeros( else Float, allow_mismatch_float_precision=allow_mismatch_float_precision, ).data + elif _field.name == "tracers": + qty_factory_tracers = FieldBundle.extend_3D_quantity_factory( + quantity_factory, {"tracers": tracer_count} + ) + initial_storages[_field.name] = qty_factory_tracers.zeros( + [X_DIM, Y_DIM, Z_DIM, "tracers"], + _field.metadata["units"], + dtype=Float, # type: ignore + ).data return cls.init_from_storages( storages=initial_storages, - sizer=quantity_factory.sizer, + quantity_factory=quantity_factory, allow_mismatch_float_precision=allow_mismatch_float_precision, ) @@ -362,7 +432,7 @@ def init_from_numpy_arrays( def init_from_storages( cls, storages: Mapping[str, Any], - sizer: GridSizer, + quantity_factory: QuantityFactory, bdt: float = 0.0, mdt: float = 0.0, allow_mismatch_float_precision=False, @@ -380,6 +450,11 @@ def init_from_storages( allow_mismatch_float_precision=allow_mismatch_float_precision, ) inputs[_field.name] = quantity + elif "tracers" == _field.name: + tracers = setup_tracers(storages["tracers"].shape[3], quantity_factory) + tracers.quantity.data[:] = storages["tracers"][:] + inputs[_field.name] = tracers + return cls(**inputs, bdt=bdt, mdt=mdt) @classmethod @@ -393,7 +468,11 @@ def from_fortran_restart( state_dict: Mapping[str, Quantity] = open_restart( dirname=path, communicator=communicator, - tracer_properties=TRACER_PROPERTIES, + tracer_properties=DEFAULT_TRACER_PROPERTIES, + ) + new = cls.init_zeros( + quantity_factory=quantity_factory, + tracer_count=len(DEFAULT_TRACER_PROPERTIES), ) new = cls.init_zeros(quantity_factory=quantity_factory) @@ -417,31 +496,33 @@ def from_fortran_restart( new.v.view[:] = new.v.np.asarray( state_dict["y_wind"].transpose(new.v.dims).view[:] ) - new.qvapor.view[:] = new.qvapor.np.asarray( - state_dict["specific_humidity"].transpose(new.qvapor.dims).view[:] + new.tracers.vapor.view[:] = new.tracers.vapor.np.asarray( + state_dict["specific_humidity"].transpose(new.tracers.vapor.dims).view[:] ) - new.qliquid.view[:] = new.qliquid.np.asarray( + new.tracers.liquid.view[:] = new.tracers.liquid.np.asarray( state_dict["cloud_liquid_water_mixing_ratio"] - .transpose(new.qliquid.dims) + .transpose(new.tracers.liquid.dims) .view[:] ) - new.qice.view[:] = new.qice.np.asarray( - state_dict["cloud_ice_mixing_ratio"].transpose(new.qice.dims).view[:] + new.tracers.ice.view[:] = new.tracers.ice.np.asarray( + state_dict["cloud_ice_mixing_ratio"].transpose(new.tracers.ice.dims).view[:] ) - new.qrain.view[:] = new.qrain.np.asarray( - state_dict["rain_mixing_ratio"].transpose(new.qrain.dims).view[:] + new.tracers.rain.view[:] = new.tracers.rain.np.asarray( + state_dict["rain_mixing_ratio"].transpose(new.tracers.rain.dims).view[:] ) - new.qsnow.view[:] = new.qsnow.np.asarray( - state_dict["snow_mixing_ratio"].transpose(new.qsnow.dims).view[:] + new.tracers.snow.view[:] = new.tracers.snow.np.asarray( + state_dict["snow_mixing_ratio"].transpose(new.tracers.snow.dims).view[:] ) - new.qgraupel.view[:] = new.qgraupel.np.asarray( - state_dict["graupel_mixing_ratio"].transpose(new.qgraupel.dims).view[:] + new.tracers.graupel.view[:] = new.tracers.graupel.np.asarray( + state_dict["graupel_mixing_ratio"] + .transpose(new.tracers.graupel.dims) + .view[:] ) - new.qo3mr.view[:] = new.qo3mr.np.asarray( - state_dict["ozone_mixing_ratio"].transpose(new.qo3mr.dims).view[:] + new.tracers.o3mr.view[:] = new.tracers.o3mr.np.asarray( + state_dict["ozone_mixing_ratio"].transpose(new.tracers.o3mr.dims).view[:] ) - new.qcld.view[:] = new.qcld.np.asarray( - state_dict["cloud_fraction"].transpose(new.qcld.dims).view[:] + new.tracers.cloud.view[:] = new.tracers.cld.np.asarray( + state_dict["cloud_fraction"].transpose(new.tracers.cld.dims).view[:] ) new.delz.view[:] = new.delz.np.asarray( state_dict["vertical_thickness_of_atmospheric_layer"] @@ -467,6 +548,12 @@ def xr_dataset(self): "units": field_info.metadata.get("units", "unknown"), }, ) + if isinstance(field_info.type, FieldBundle): + data_vars[name] = self._xr_dataarray_from_quantity( + name=name, + metadata=field_info.metadata, + data=getattr(self, name).quantity.data, + ) return xr.Dataset(data_vars=data_vars) def __getitem__(self, item): diff --git a/pyFV3/stencils/divergence_damping.py b/pyFV3/stencils/divergence_damping.py index 46a949ce..39db3a1e 100644 --- a/pyFV3/stencils/divergence_damping.py +++ b/pyFV3/stencils/divergence_damping.py @@ -632,7 +632,7 @@ def __call__( self.v_contra_dxc, ) - da_min_c: Float = self._get_da_min_c() + da_min_c = self._damping_coefficients.da_min_c self._damping( delpc, damped_rel_vort_bgrid, @@ -694,7 +694,7 @@ def __call__( abs(dt), ) - da_min: Float = self._get_da_min() + da_min = self._damping_coefficients.da_min if self._stretched_grid: # reference https://github.com/NOAA-GFDL/GFDL_atmos_cubed_sphere/blob/main/model/sw_core.F90#L1422 # noqa: E501 dd8 = da_min * self._d4_bg ** (self._nonzero_nord + 1) diff --git a/pyFV3/stencils/dyn_core.py b/pyFV3/stencils/dyn_core.py index 2df88a30..92b57aa1 100644 --- a/pyFV3/stencils/dyn_core.py +++ b/pyFV3/stencils/dyn_core.py @@ -642,22 +642,6 @@ def __init__( pkc=self._pkc, ) - # See divergence_damping.py, _get_da_min for explanation of this function - @dace_inhibitor - def _get_da_min(self) -> np.float64: - return self._da_min - - # TODO: fix me - we shouldn't need a function here, Dace is fudging the types - # See https://github.com/GEOS-ESM/pace/issues/9 - @dace_inhibitor - def dt_acoustic_substep(self, timestep: Float) -> Float: - return Float(timestep / self.config.n_split) - - # TODO: Same as above - @dace_inhibitor - def dt2(self, dt_acoustic_substep: Float) -> Float: - return Float(0.5) * dt_acoustic_substep - def __call__( self, state: DycoreState, @@ -674,8 +658,8 @@ def __call__( # akap, ptop, n_map, comm): end_step = n_map == self.config.k_split # dt = state.mdt / self.config.n_split - dt_acoustic_substep = self.dt_acoustic_substep(timestep) - dt2 = self.dt2(dt_acoustic_substep) + dt_acoustic_substep = Float(timestep / self.config.n_split) + dt2 = Float(0.5) * dt_acoustic_substep n_split = self.config.n_split # NOTE: In Fortran model the halo update starts happens in fv_dynamics, not here self._halo_updaters.q_con__cappa.start() @@ -951,7 +935,7 @@ def __call__( if self._do_del2cubed: self._halo_updaters.heat_source.update() # TODO: move dependence on da_min into init of hyperdiffusion class - cd = constants.CNST_0P20 * self._get_da_min() + cd = constants.CNST_0P20 * self._da_min # we want to diffuse the heat source from damping before we apply it, # so that we don't reinforce the same grid-scale patterns we're trying # to damp diff --git a/pyFV3/testing/translate_dyncore.py b/pyFV3/testing/translate_dyncore.py index 47415671..0e1703f6 100644 --- a/pyFV3/testing/translate_dyncore.py +++ b/pyFV3/testing/translate_dyncore.py @@ -149,6 +149,7 @@ def compute_parallel(self, inputs, communicator): state = DycoreState.init_zeros( quantity_factory=self.grid.quantity_factory, dtype_dict=inputs_dtypes, + tracer_count=1, # No tracers used in acoustics allow_mismatch_float_precision=True, ) wsd = self.grid.quantity_factory.zeros( From 1476dcde1aafacf14d452fa31515fe6e995c028c Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 26 Jun 2025 08:36:12 -0400 Subject: [PATCH 233/252] Fix ComputeTotalEnergy w/ new Tracers --- pyFV3/stencils/compute_total_energy.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/compute_total_energy.py b/pyFV3/stencils/compute_total_energy.py index 919193d8..1d1dd8b5 100644 --- a/pyFV3/stencils/compute_total_energy.py +++ b/pyFV3/stencils/compute_total_energy.py @@ -2,7 +2,7 @@ from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, GRAV from pyFV3._config import DynamicalCoreConfig -from pyFV3.tracers import Tracers +from pyFV3.tracers import TracersType from pyFV3.stencils.moist_cv import moist_cv_nwat6_fn from gt4py.cartesian.gtscript import ( BACKWARD, @@ -135,7 +135,7 @@ def __call__( u: FloatField, # type: ignore v: FloatField, # type: ignore w: FloatField, # type: ignore - tracers: Tracers, + tracers: TracersType, te_2d: FloatFieldIJ, # type: ignore ) -> None: self._compute_total_energy( From b893fabb6968e7b9a4165ab52e37f4d113989eae Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 27 Jun 2025 12:08:04 -0400 Subject: [PATCH 234/252] Orchestrater compute total energy --- pyFV3/stencils/compute_total_energy.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/pyFV3/stencils/compute_total_energy.py b/pyFV3/stencils/compute_total_energy.py index 1d1dd8b5..bb4f0283 100644 --- a/pyFV3/stencils/compute_total_energy.py +++ b/pyFV3/stencils/compute_total_energy.py @@ -1,4 +1,4 @@ -from ndsl import StencilFactory, QuantityFactory +from ndsl import StencilFactory, QuantityFactory, orchestrate from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM, GRAV from pyFV3._config import DynamicalCoreConfig @@ -100,6 +100,11 @@ def __init__( quantity_factory: QuantityFactory, grid_data: GridData, ) -> None: + orchestrate( + obj=self, + config=stencil_factory.config.dace_config, + dace_compiletime_args=["tracers"], + ) if config.hydrostatic: raise NotImplementedError( "Dynamics (Compute Total Energy): " @@ -147,12 +152,12 @@ def __call__( u=u, v=v, w=w, - qvapor=tracers["vapor"], - qliquid=tracers["liquid"], - qrain=tracers["rain"], - qsnow=tracers["snow"], - qice=tracers["ice"], - qgraupel=tracers["graupel"], + qvapor=tracers.vapor, + qliquid=tracers.liquid, + qrain=tracers.rain, + qsnow=tracers.snow, + qice=tracers.ice, + qgraupel=tracers.graupel, rsin2=self._rsin2, cosa_s=self._cosa_s, phyz=self._phyz, From c7ce6b23ec0904fee8af5ae31c829e505a1f9037 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 27 Jun 2025 12:08:22 -0400 Subject: [PATCH 235/252] Orchestrate dynamics + new tracers structure --- pyFV3/stencils/fv_dynamics.py | 39 ++++++++++++++--------------------- 1 file changed, 16 insertions(+), 23 deletions(-) diff --git a/pyFV3/stencils/fv_dynamics.py b/pyFV3/stencils/fv_dynamics.py index afe8eca5..2aedcca5 100644 --- a/pyFV3/stencils/fv_dynamics.py +++ b/pyFV3/stencils/fv_dynamics.py @@ -312,13 +312,13 @@ def __init__( "graupel", "cloud", ] - if not all(n in state.tracers.names() for n in required_tracers): + if not all(n in state.tracers._indexer.keys() for n in required_tracers): raise NotImplementedError( "Dynamical core (fv_dynamics):" " missing required tracers. Dynamics requires:\n" f" {required_tracers}\n" "but only the following where given:\n" - f" {state.tracers.names()}" + f" {state.tracers._indexer.keys()}" ) self._comm = comm @@ -359,7 +359,6 @@ def __init__( self.grid_data, comm, state.tracers, - exclude_tracers=exclude_tracers, ) self._ak = grid_data.ak self._bk = grid_data.bk @@ -541,11 +540,6 @@ def __init__( domain=grid_indexing.domain_compute(add=(1, 1, 0)), ) - # See divergence_damping.py, _get_da_min for explanation of this function - @dace_inhibitor - def _get_da_min(self) -> NDSL_64BIT_FLOAT_TYPE: # type: ignore - return self._da_min - def step_dynamics( self, state: DycoreState, @@ -573,12 +567,12 @@ def compute_preamble(self, state: DycoreState, is_root_rank: bool): self._set_value(state.cyd, Float(0.0)) self._fv_setup_stencil( - state.tracers["vapor"], - state.tracers["liquid"], - state.tracers["rain"], - state.tracers["snow"], - state.tracers["ice"], - state.tracers["graupel"], + state.tracers.vapor, + state.tracers.liquid, + state.tracers.rain, + state.tracers.snow, + state.tracers.ice, + state.tracers.graupel, state.q_con, self._cvm, state.pkz, @@ -782,7 +776,6 @@ def _compute(self, state: DycoreState, timer: Timer): self._increment(state.cyd, self._cy_local) if last_step: - da_min = self._get_da_min() if not self.config.hydrostatic: if __debug__: log_on_rank_0("Omega") @@ -798,18 +791,18 @@ def _compute(self, state: DycoreState, timer: Timer): if __debug__: log_on_rank_0("Del2Cubed") self._omega_halo_updater.update() - self._hyperdiffusion(state.omga, Float(0.18) * da_min) + self._hyperdiffusion(state.omga, Float(0.18) * self._da_min) if __debug__: log_on_rank_0("Neg Adj 3") self._adjust_tracer_mixing_ratio( - state.tracers["vapor"], - state.tracers["liquid"], - state.tracers["rain"], - state.tracers["snow"], - state.tracers["ice"], - state.tracers["graupel"], - state.tracers["cloud"], + state.tracers.vapor, + state.tracers.liquid, + state.tracers.rain, + state.tracers.snow, + state.tracers.ice, + state.tracers.graupel, + state.tracers.cloud, state.pt, state.delp, ) From fa339ff4bf7b774e589401ebb07f9fd1e6dec2bb Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 27 Jun 2025 12:08:46 -0400 Subject: [PATCH 236/252] Fix FVDynamics translate tests --- pyFV3/testing/translate_fvdynamics.py | 106 +++++--------------------- 1 file changed, 18 insertions(+), 88 deletions(-) diff --git a/pyFV3/testing/translate_fvdynamics.py b/pyFV3/testing/translate_fvdynamics.py index 147520e9..4951ecbf 100644 --- a/pyFV3/testing/translate_fvdynamics.py +++ b/pyFV3/testing/translate_fvdynamics.py @@ -5,7 +5,7 @@ import pytest import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, Quantity, QuantityFactory, StencilFactory +from ndsl import Namelist, Quantity, QuantityFactory, StencilFactory, FieldBundle from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -33,27 +33,6 @@ def __init__( self.namelist = DynamicalCoreConfig.from_namelist(namelist) -TRACERS_IN_PYFV3 = [ - "vapor", - "liquid", - "ice", - "rain", - "snow", - "graupel", - "cloud", -] - -TRACERS_IN_FORTRAN = [ - "qvapor", - "qliquid", - "qice", - "qrain", - "qsnow", - "qgraupel", - "qcld", -] - - class TranslateFVDynamics(ParallelTranslateBaseSlicing): compute_grid_option = True inputs: Dict[str, Any] = { @@ -185,41 +164,6 @@ class TranslateFVDynamics(ParallelTranslateBaseSlicing): "units": "m^2 s^-2", "dims": [X_DIM, Y_DIM], }, - "qvapor": { - "name": "specific_humidity", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - }, - "qliquid": { - "name": "cloud_water_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - }, - "qice": { - "name": "cloud_ice_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - }, - "qrain": { - "name": "rain_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - }, - "qsnow": { - "name": "snow_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - }, - "qgraupel": { - "name": "graupel_mixing_ratio", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "kg/kg", - }, - "qcld": { - "name": "cloud_fraction", - "dims": [X_DIM, Y_DIM, Z_DIM], - "units": "", - }, "omga": { "name": "vertical_pressure_velocity", "dims": [X_DIM, Y_DIM, Z_DIM], @@ -230,6 +174,7 @@ class TranslateFVDynamics(ParallelTranslateBaseSlicing): } outputs = inputs.copy() + outputs["tracers"] = {} for name in ("bdt", "ak", "bk", "ptop", "ua"): outputs.pop(name) @@ -248,13 +193,6 @@ def __init__( "v": grid.x3d_domain_dict(), "w": {}, "delz": {}, - "qvapor": grid.compute_dict(), - "qliquid": grid.compute_dict(), - "qice": grid.compute_dict(), - "qrain": grid.compute_dict(), - "qsnow": grid.compute_dict(), - "qgraupel": grid.compute_dict(), - "qcld": {}, "ps": {}, "pe": { "istart": grid.is_ - 1, @@ -294,19 +232,12 @@ def __init__( self._base.out_vars.update(fv_dynamics_vars) self._base.out_vars["ps"] = {"kstart": grid.npz - 1, "kend": grid.npz - 1} self._base.out_vars["phis"] = {"kstart": grid.npz - 1, "kend": grid.npz - 1} + self._base.out_vars["tracers"] = {} self._base.out_vars.pop("ua") self.max_error = 1e-5 self.ignore_near_zero_errors = {} - self.ignore_near_zero_errors["qvapor"] = True - self.ignore_near_zero_errors["qliquid"] = True - self.ignore_near_zero_errors["qice"] = True - self.ignore_near_zero_errors["qrain"] = True - self.ignore_near_zero_errors["qsnow"] = True - self.ignore_near_zero_errors["qgraupel"] = True - self.ignore_near_zero_errors["qcld"] = True - self.ignore_near_zero_errors["q_con"] = True self.dycore: Optional[fv_dynamics.DynamicalCore] = None self.stencil_factory = stencil_factory self._quantity_factory = QuantityFactory.from_backend( @@ -316,15 +247,17 @@ def __init__( self.namelist: DynamicalCoreConfig = DynamicalCoreConfig.from_namelist(namelist) def state_from_inputs(self, inputs): + tracers = self._quantity_factory._numpy.empty( + ( + inputs["tracers"].shape[0] + 1, + inputs["tracers"].shape[1] + 1, + inputs["tracers"].shape[2] + 1, + inputs["tracers"].shape[3], + ) + ) + tracers[:-1, :-1, :-1, :] = inputs.pop("tracers") input_storages = super().state_from_inputs(inputs) - # extract tracers - input_storages["vapor"] = input_storages.pop("qvapor") - input_storages["liquid"] = input_storages.pop("qliquid") - input_storages["ice"] = input_storages.pop("qice") - input_storages["rain"] = input_storages.pop("qrain") - input_storages["snow"] = input_storages.pop("qsnow") - input_storages["graupel"] = input_storages.pop("qgraupel") - input_storages["cloud"] = input_storages.pop("qcld") + input_storages["tracers"] = tracers # Move fluxes and courant numbers input_storages["mfxd"] = input_storages.pop("mfxd_FV") input_storages["mfyd"] = input_storages.pop("mfyd_FV") @@ -332,7 +265,6 @@ def state_from_inputs(self, inputs): input_storages["cyd"] = input_storages.pop("cyd_FV") # making sure we init DycoreState with the exact set of variables accepted_keys = [_field.name for _field in fields(DycoreState)] - accepted_keys += TRACERS_IN_PYFV3 todelete = [] for name in input_storages.keys(): if name not in accepted_keys: @@ -342,7 +274,6 @@ def state_from_inputs(self, inputs): state = DycoreState.init_from_storages( storages=input_storages, quantity_factory=self._quantity_factory, - tracer_list=TRACERS_IN_PYFV3, ) return state @@ -382,17 +313,16 @@ def compute_parallel(self, inputs, communicator): outputs = self.outputs_from_state(state) return outputs - def outputs_from_state(self, state: dict): + def outputs_from_state(self, state: DycoreState): if len(self.outputs) == 0: return {} outputs = {} storages = {} - for name, properties in self.outputs.items(): - if name in TRACERS_IN_FORTRAN: - idx = TRACERS_IN_FORTRAN.index(name) - storages[name] = state["tracers"][TRACERS_IN_PYFV3[idx]].data - elif name in ["mfxd_FV", "mfyd_FV", "cxd_FV", "cyd_FV"]: + for name, _properties in self.outputs.items(): + if name in ["mfxd_FV", "mfyd_FV", "cxd_FV", "cyd_FV"]: storages[name] = state[name[:-3]].data + elif isinstance(state[name], FieldBundle): + storages[name] = state[name].quantity.data elif isinstance(state[name], Quantity): storages[name] = state[name].data elif len(self.outputs[name]["dims"]) > 0: From c4c9ffeeeaeba834bb2359f65da511134a9f3023 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 27 Jun 2025 13:42:30 -0400 Subject: [PATCH 237/252] Orchestrate `ray_fast` w/ workaround for lazy call --- pyFV3/stencils/ray_fast.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index e1dea877..4af0a7ed 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -25,6 +25,7 @@ Y_DIM, ) from ndsl.dsl.typing import Float, FloatField, FloatFieldK +import numpy as np # NOTE: The fortran version of this computes rf in the first timestep only. Then @@ -211,14 +212,14 @@ def __init__( # We compute the damping increment once using a trick to write a # FloatFieldK as a (1, 1, K) 3D writable Field - K_stencil_factory, K_quantity_factory = get_factories_single_tile( + _K_stencil_factory, K_quantity_factory = get_factories_single_tile( 1, 1, domain[2], 0, stencil_factory.backend, ) - self._ray_fast_damping_increment = K_stencil_factory.from_origin_domain( + self._ray_fast_damping_increment = stencil_factory.from_origin_domain( ray_fast_damping_increment, origin=(0, 0, origin[2]), domain=(1, 1, domain[2]), @@ -230,7 +231,7 @@ def __init__( self._damping_increment = K_quantity_factory.ones( [X_DIM, Y_DIM, Z_DIM], units="n/a" ) - self._initialize_damping_increment = False + self._initialize_damping_increment = np.ones((1,), dtype=int) def __call__( self, @@ -254,17 +255,19 @@ def __call__( """ rf_cutoff_nudge = self._rf_cutoff + min(Float(100.0), Float(10.0) * ptop) - if not self._initialize_damping_increment: + # TODO: this is a bad fix to go around an orchestration issue + # on compile-time values. Do better. + if self._initialize_damping_increment[0] == 1: self._ray_fast_damping_increment( pfull=pfull, dt=dt, ptop=ptop, rf=self._damping_increment ) - self._initialize_damping_increment = True + self._initialize_damping_increment[0] = 0 self._ray_fast_wind_compute( u=u, v=v, w=w, delta_p_ref=dp, pfull=pfull, - rf=self._damping_increment.view[0, 0, :], + rf=self._damping_increment.field[0, 0, :], rf_cutoff_nudge=rf_cutoff_nudge, ) From e25f246e00ae80c3dab24b0c9d4e718f7166ff0d Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 27 Jun 2025 13:44:21 -0400 Subject: [PATCH 238/252] Fix assert in `sum` --- pyFV3/mpi/sum.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyFV3/mpi/sum.py b/pyFV3/mpi/sum.py index 8bbbd976..3ec4a8f2 100644 --- a/pyFV3/mpi/sum.py +++ b/pyFV3/mpi/sum.py @@ -14,6 +14,6 @@ def __init__( @dace_inhibitor def __call__(self, qty_to_sum: Quantity) -> Float: - assert qty_to_sum.field.shape == 2 # Code handle only 2D quantity + assert len(qty_to_sum.field.shape) == 2 # Code handle only 2D quantity self._comm.all_reduce(qty_to_sum, ReductionOperator.SUM, self._tmp_reduce) return qty_to_sum.field.sum(axis=0).sum(axis=0) From c209450a43ae41a73ed661c717ec6b1c2f81ab77 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Mon, 21 Jul 2025 14:26:23 -0400 Subject: [PATCH 239/252] Remove redundant cast --- pyFV3/stencils/remapping_GEOS.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index fd551767..83932e3c 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -567,7 +567,7 @@ def __call__( tesum: Float = self._global_sum(self._te_2d) zsum: Float = self._global_sum(self._zsum1) - dtmp = tesum / (CV_AIR * zsum) + dtmp: Float = tesum / (CV_AIR * zsum) elif consv_te < -CONSV_MIN: raise NotImplementedError( @@ -618,7 +618,7 @@ def __call__( qgraupel=tracers.graupel, pt=pt, pkz=pkz, - dtmp=Float(dtmp), + dtmp=dtmp, r_vir=zvir, ) self._fill_cond( From 60422dc1420e7989c69eef5cf7ad540fe3bf02bf Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 3 Sep 2025 13:46:35 -0400 Subject: [PATCH 240/252] Update code to M2 standards --- pyFV3/stencils/a2b_ord4.py | 1 - pyFV3/stencils/divergence_damping.py | 6 +++--- pyFV3/stencils/ray_fast.py | 4 ++-- pyFV3/stencils/tracer_2d_1l.py | 6 +++--- 4 files changed, 8 insertions(+), 9 deletions(-) diff --git a/pyFV3/stencils/a2b_ord4.py b/pyFV3/stencils/a2b_ord4.py index 0744796f..5f9a9fe6 100644 --- a/pyFV3/stencils/a2b_ord4.py +++ b/pyFV3/stencils/a2b_ord4.py @@ -10,7 +10,6 @@ region, sin, sqrt, - f64, ) from ndsl import GridIndexing, QuantityFactory, StencilFactory, orchestrate diff --git a/pyFV3/stencils/divergence_damping.py b/pyFV3/stencils/divergence_damping.py index 35cc1f20..480ff85f 100644 --- a/pyFV3/stencils/divergence_damping.py +++ b/pyFV3/stencils/divergence_damping.py @@ -4,7 +4,7 @@ __INLINED, PARALLEL, computation, - f32, + float32, horizontal, interval, region, @@ -25,13 +25,13 @@ @gtscript.function def damp_tmp(q, da_min_c, d2_bg, dddmp): - damp: f32 = da_min_c * max(d2_bg, min(0.2, dddmp * abs(q))) + damp: float32 = da_min_c * max(d2_bg, min(0.2, dddmp * abs(q))) return damp @gtscript.function def damp_tmp2(q, da_min_c, d2_bg, dddmp): - damp: f32 = da_min_c * max(d2_bg, min(0.2, dddmp * q)) + damp: float32 = da_min_c * max(d2_bg, min(0.2, dddmp * q)) return damp diff --git a/pyFV3/stencils/ray_fast.py b/pyFV3/stencils/ray_fast.py index 4af0a7ed..96523da9 100644 --- a/pyFV3/stencils/ray_fast.py +++ b/pyFV3/stencils/ray_fast.py @@ -5,7 +5,7 @@ FORWARD, PARALLEL, computation, - f64, + float64, horizontal, interval, log, @@ -43,7 +43,7 @@ def compute_rf_vals(pfull, bdt, rf_cutoff, tau0, ptop): @gtscript.function def compute_rff_vals(pfull, dt, rf_cutoff, tau0, ptop): rffvals = compute_rf_vals(pfull, dt, rf_cutoff, tau0, ptop) - rffvals = f64(1.0) / (f64(1.0) + rffvals) + rffvals = float64(1.0) / (float64(1.0) + rffvals) return rffvals diff --git a/pyFV3/stencils/tracer_2d_1l.py b/pyFV3/stencils/tracer_2d_1l.py index b5c86363..6a4c5552 100644 --- a/pyFV3/stencils/tracer_2d_1l.py +++ b/pyFV3/stencils/tracer_2d_1l.py @@ -9,7 +9,7 @@ horizontal, interval, region, - i32, + int32, ) from ndsl import ( @@ -116,7 +116,7 @@ def divide_fluxes_by_n_substeps( mfyd (inout): """ with computation(PARALLEL), interval(...): - n_split = i32(1.0 + cmax) + n_split = int32(1.0 + cmax) if n_split > 1: frac = 1.0 / n_split cxd = cxd * frac @@ -175,7 +175,7 @@ def apply_tracer_flux( dp2 (in): """ with computation(PARALLEL), interval(...): - if current_nsplit < i32(1.0 + cmax): + if current_nsplit < int32(1.0 + cmax): q = (q * dp1 + ((fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) * rarea) / dp2 From 98601fbece9e25b318fd0fa0fb28a41471f13bc1 Mon Sep 17 00:00:00 2001 From: Charles Kropiewnicki Date: Mon, 22 Sep 2025 13:37:48 -0400 Subject: [PATCH 241/252] Fixed UpdateDzC - computation now occurs at all needed points, no longer misses outermost x/y --- pyFV3/stencils/updatedzc.py | 204 +++++++++++++++++++----------------- 1 file changed, 107 insertions(+), 97 deletions(-) diff --git a/pyFV3/stencils/updatedzc.py b/pyFV3/stencils/updatedzc.py index 0e346e9f..46c46108 100644 --- a/pyFV3/stencils/updatedzc.py +++ b/pyFV3/stencils/updatedzc.py @@ -4,7 +4,9 @@ FORWARD, PARALLEL, computation, + horizontal, interval, + region, ) from ndsl import Quantity, QuantityFactory, StencilFactory @@ -13,44 +15,6 @@ from ndsl.stencils import corners -@gtscript.function -def p_weighted_average_top(vel, top_dp_ratio): - # TODO: ratio is a constant, where should this be placed? - return vel + (vel - vel[0, 0, 1]) * top_dp_ratio - - -@gtscript.function -def p_weighted_average_bottom(vel, bottom_dp_ratio): - return vel[0, 0, -1] + (vel[0, 0, -1] - vel[0, 0, -2]) * bottom_dp_ratio - - -@gtscript.function -def p_weighted_average_domain(vel, dp0): - int_ratio = 1.0 / (dp0[-1] + dp0) - return (dp0 * vel[0, 0, -1] + dp0[-1] * vel) * int_ratio - - -@gtscript.function -def xy_flux(gz_x, gz_y, xfx, yfx): - """ - Compute first-order upwind fluxes of gz in x and y directions. - - Args: - gz_x: gz with corners copied to perform derivatives in x-direction - gz_y: gz with corners copied to perform derivatives in y-direction - xfx (out): contravariant c-grid u-wind interpolated to layer interfaces, - including metric terms to make it a "volume flux" - yfx (out): contravariant c-grid v-wind interpolated to layer interfaces - - Returns: - fx: first-order upwind x-flux of gz - fy: first-order upwind y-flux of gz - """ - fx = xfx * (gz_x[-1, 0, 0] if xfx > 0.0 else gz_x) - fy = yfx * (gz_y[0, -1, 0] if yfx > 0.0 else gz_y) - return fx, fy - - def double_copy(q_in: FloatField, copy_1: FloatField, copy_2: FloatField): with computation(PARALLEL), interval(...): copy_1 = q_in @@ -62,61 +26,67 @@ def copy(q_in: FloatField, q_copy: FloatField): q_copy = q_in -def update_dz_c( +def compute_weighted_average( dp_ref: FloatFieldK, - zs: FloatFieldIJ, - area: FloatFieldIJ, - ut: FloatField, - vt: FloatField, - gz: FloatField, - gz_x: FloatField, - gz_y: FloatField, - gz_filled: FloatField, - ws: FloatFieldIJ, - *, - dt: Float, - dz_min: Float, + vel: FloatField, + avg: FloatField, ): - """ - Step dz forward on c-grid - Eusures gz is monotonically increasing in z at the end - Args: - dp_ref(in): layer thickness in Pa - zs(in): surface height in m - area(in): - ut(in): horizontal wind (TODO: covariant or contravariant?) - vt(in): horizontal wind (TODO: covariant or contravariant?) - gz(inout): geopotential height on model interfaces - gz_x(in): gz with corners copied to perform derivatives in x-direction - gz_y(in): gz with corners copied to perform derivatives in y-direction - ws(out): lagrangian (parcel-following) surface vertical wind implied by - lowest-level gz change note that a parcel moving horizontally - across terrain will be moving in the vertical (eqn 5.5 in documentation) - dt(in): timestep over which to evolve the geopotential height, in seconds - dz_min(in): Controls minimum thickness in NH solver - """ # there's some complexity due to gz being defined on interfaces # have to interpolate winds to layer interfaces first, using higher-order # cubic spline interpolation with computation(PARALLEL): with interval(0, 1): - # TODO: inline some or all of these functions - xfx = p_weighted_average_top(ut, dp_ref) - yfx = p_weighted_average_top(vt, dp_ref) + top_ratio = dp_ref / (dp_ref + dp_ref[1]) + avg = vel + (vel - vel[0, 0, 1]) * top_ratio with interval(1, -1): - xfx = p_weighted_average_domain(ut, dp_ref) - yfx = p_weighted_average_domain(vt, dp_ref) + int_ratio = 1.0 / (dp_ref[-1] + dp_ref) + avg = (dp_ref * vel[0, 0, -1] + dp_ref[-1] * vel) * int_ratio with interval(-1, None): - xfx = p_weighted_average_bottom(ut, dp_ref) - yfx = p_weighted_average_bottom(vt, dp_ref) - # xfx/yfx are now ut/vt interpolated to layer interfaces + bot_ratio = dp_ref[-1] / (dp_ref[-2] + dp_ref[-1]) + avg = vel[0, 0, -1] + (vel[0, 0, -1] - vel[0, 0, -2]) * bot_ratio + + +def compute_fx_fy( + gz_x: FloatField, + gz_y: FloatField, + xfx: FloatField, + yfx: FloatField, + fx: FloatField, + fy: FloatField, +): + with computation(PARALLEL), interval(...): + if xfx > 0.0: + fx = gz_x[-1, 0, 0] + else: + fx = gz_x + fx = xfx * fx + + if yfx > 0.0: + fy = gz_y[0, -1, 0] + else: + fy = gz_y + fy = yfx * fy + + +def compute_gz_ws( + gz_y: FloatField, + area: FloatFieldIJ, + fx: FloatField, + fy: FloatField, + xfx: FloatField, + yfx: FloatField, + dz_min: Float, + dt: Float, + zs: FloatFieldIJ, + ws: FloatFieldIJ, + gz: FloatField, +): with computation(PARALLEL), interval(...): - fx, fy = xy_flux(gz_x, gz_y, xfx, yfx) - gz = (gz_filled * area + (fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) / ( + gz = (gz_y * area + (fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) / ( area + (xfx - xfx[1, 0, 0]) + (yfx - yfx[0, 1, 0]) ) - with computation(FORWARD), interval(-1, None): + with computation(FORWARD), interval(...): rdt = 1.0 / dt ws = (zs - gz) * rdt with computation(BACKWARD), interval(0, -1): @@ -170,6 +140,26 @@ def __init__( units="m**2/s**2", dtype=Float, ) + self._xfx = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._yfx = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._fx = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + self._fy = quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) full_origin = grid_indexing.origin_full() full_domain = grid_indexing.domain_full(add=(0, 0, 1)) self._double_copy_stencil = stencil_factory.from_origin_domain( @@ -199,12 +189,26 @@ def __init__( domain=full_domain, ) - self._update_dz_c = stencil_factory.from_origin_domain( - update_dz_c, + self._compute_weighted_average = stencil_factory.from_origin_domain( + compute_weighted_average, + origin=grid_indexing.origin_compute(add=(-1, -1, 0)), + domain=grid_indexing.domain_compute(add=(3, 3, 1)), + ) + + self._compute_flux = stencil_factory.from_origin_domain( + compute_fx_fy, + origin=grid_indexing.origin_compute(add=(-1, -1, 0)), + domain=grid_indexing.domain_compute(add=(3, 3, 1)), + ) + + self._compute_gz_ws = stencil_factory.from_origin_domain( + compute_gz_ws, origin=grid_indexing.origin_compute(add=(-1, -1, 0)), domain=grid_indexing.domain_compute(add=(2, 2, 1)), ) + self.DEBUG_VAR_1 = quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM], "n/a") + def __call__( self, zs: FloatFieldIJ, @@ -228,28 +232,34 @@ def __call__( # TODO: is this advecting gz, and if so can we name it that? # Can we reduce duplication of advection logic with other stencils? - self._copy_stencil(gz, self._gz_filled) - self._fill_corners_x_stencil(self._gz_filled, self._gz_filled) - self._fill_corners_y_stencil(self._gz_filled, self._gz_filled) - self._double_copy_stencil(gz, self._gz_x, self._gz_y) - # TODO(eddied): We pass the same fields 2x to avoid GTC validation errors if self._grid_type < 3: self._fill_corners_x_stencil(self._gz_x, self._gz_x) self._fill_corners_y_stencil(self._gz_y, self._gz_y) - self._update_dz_c( - dp_ref=self._dp_ref, - zs=zs, - area=self._area, - ut=ut, - vt=vt, - gz=gz, + self._compute_weighted_average(dp_ref=self._dp_ref, vel=ut, avg=self._xfx) + self._compute_weighted_average(dp_ref=self._dp_ref, vel=vt, avg=self._yfx) + + self._compute_flux( gz_x=self._gz_x, gz_y=self._gz_y, - gz_filled=self._gz_filled, - ws=ws, - dt=dt, + xfx=self._xfx, + yfx=self._yfx, + fx=self._fx, + fy=self._fy, + ) + + self._compute_gz_ws( + gz_y=self._gz_y, + area=self._area, + fx=self._fx, + fy=self._fy, + xfx=self._xfx, + yfx=self._yfx, dz_min=self._dz_min, + dt=dt, + zs=zs, + ws=ws, + gz=gz, ) From c891618d6f179dd02f3c2c4d7a1307fa1a79004a Mon Sep 17 00:00:00 2001 From: Charles Kropiewnicki Date: Mon, 22 Sep 2025 14:05:59 -0400 Subject: [PATCH 242/252] Brought back UpdateDzC documentation --- pyFV3/stencils/updatedzc.py | 43 ++++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/pyFV3/stencils/updatedzc.py b/pyFV3/stencils/updatedzc.py index 46c46108..2c546475 100644 --- a/pyFV3/stencils/updatedzc.py +++ b/pyFV3/stencils/updatedzc.py @@ -31,10 +31,16 @@ def compute_weighted_average( vel: FloatField, avg: FloatField, ): + """ + Perform a cubic spline interpolation of wind velocity from grid center to grid edge + Args: + dp_ref(in): layer thickness in Pa + vel(in): grid center wind speed + avg(out: interpolated (grid edge) wind speed + """ # there's some complexity due to gz being defined on interfaces # have to interpolate winds to layer interfaces first, using higher-order - # cubic spline interpolation with computation(PARALLEL): with interval(0, 1): top_ratio = dp_ref / (dp_ref + dp_ref[1]) @@ -55,6 +61,19 @@ def compute_fx_fy( fx: FloatField, fy: FloatField, ): + """ + Compute first-order upwind fluxes of gz in x and y directions. + + Args: + gz_x(in): gz with corners copied to perform derivatives in x-direction + gz_y(in): gz with corners copied to perform derivatives in y-direction + xfx(in): contravariant c-grid u-wind interpolated to layer interfaces, + including metric terms to make it a "volume flux" + yfx(in): contravariant c-grid v-wind interpolated to layer interfaces + fx(out): first-order upwind x-flux of gz + fy(out): first-order upwind y-flux of gz + """ + with computation(PARALLEL), interval(...): if xfx > 0.0: fx = gz_x[-1, 0, 0] @@ -82,6 +101,26 @@ def compute_gz_ws( ws: FloatFieldIJ, gz: FloatField, ): + """ + Compute gz and wd, eusures gz is monotonically increasing in z at the end + + Args + gz_y(in): gz with corners copied to perform derivatives in y-direction + area(in): + fx(in): first-order upwind x-flux of gz + fy(in): first-order upwind y-flux of gz + xfx(in): contravariant c-grid u-wind interpolated to layer interfaces, + including metric terms to make it a "volume flux" + yfx(in): contravariant c-grid v-wind interpolated to layer interfaces + dz_min(in): Controls minimum thickness in NH solver + dt(in): timestep over which to evolve the geopotential height, in seconds + zs(in): surface height in m + ws(out): lagrangian (parcel-following) surface vertical wind implied by + lowest-level gz change note that a parcel moving horizontally + across terrain will be moving in the vertical (eqn 5.5 in documentation) + gz(out): geopotential height on model interfaces + """ + with computation(PARALLEL), interval(...): gz = (gz_y * area + (fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) / ( area + (xfx - xfx[1, 0, 0]) + (yfx - yfx[0, 1, 0]) @@ -219,6 +258,8 @@ def __call__( dt: Float, ): """ + Step dz forward on c-grid + Args: dp_ref: layer thickness in Pa zs: surface height in m From 802e35874e95d3cb8c1cf2555a3a5ce4f78c9e1f Mon Sep 17 00:00:00 2001 From: Christopher Kung Date: Wed, 24 Sep 2025 21:15:42 -0400 Subject: [PATCH 243/252] Fixed GlobalSum class and updated translate tests --- pyFV3/mpi/sum.py | 33 +++++++++++++++-- pyFV3/stencils/remapping_GEOS.py | 7 +++- .../translate/translate_MapN_Tracer_2d.py | 23 ++++++++---- .../translate/translate_moistcvpluspt_2d.py | 16 ++++++-- .../translate_moistcvpluspt_2d_last_step.py | 6 +-- .../translate/translate_remapping_GEOS.py | 37 +++++++++++-------- 6 files changed, 86 insertions(+), 36 deletions(-) diff --git a/pyFV3/mpi/sum.py b/pyFV3/mpi/sum.py index 3ec4a8f2..e5697326 100644 --- a/pyFV3/mpi/sum.py +++ b/pyFV3/mpi/sum.py @@ -1,19 +1,44 @@ +import numpy as np + from ndsl import Quantity, QuantityFactory -from ndsl.dsl.typing import Float from ndsl.comm.communicator import Communicator, ReductionOperator from ndsl.constants import X_DIM, Y_DIM from ndsl.dsl.dace.orchestration import dace_inhibitor +from ndsl.dsl.stencil import GridIndexing +from ndsl.dsl.typing import Float +from ndsl.optional_imports import cupy as cp class GlobalSum: def __init__( - self, quantity_factory: QuantityFactory, communicator: Communicator + self, + quantity_factory: QuantityFactory, + communicator: Communicator, + grid_indexing: GridIndexing = None, ) -> None: self._comm = communicator - self._tmp_reduce = quantity_factory.empty(dims=[X_DIM, Y_DIM], units="n/a") + # self._tmp_reduce = quantity_factory.empty(dims=[X_DIM, Y_DIM], units="n/a") + self._tmp_reduce = quantity_factory.zeros(dims=[X_DIM, Y_DIM], units="n/a") + self._isc = grid_indexing.isc + self._iec = grid_indexing.iec + self._jsc = grid_indexing.jsc + self._jec = grid_indexing.jec @dace_inhibitor def __call__(self, qty_to_sum: Quantity) -> Float: assert len(qty_to_sum.field.shape) == 2 # Code handle only 2D quantity self._comm.all_reduce(qty_to_sum, ReductionOperator.SUM, self._tmp_reduce) - return qty_to_sum.field.sum(axis=0).sum(axis=0) + if isinstance(self._tmp_reduce.data, np.ndarray): + return np.sum( + self._tmp_reduce.data[ + self._isc : self._iec + 1, self._jsc : self._jec + 1 + ] + ) + elif isinstance(self._tmp_reduce.data, cp.ndarray) and cp is not None: + return cp.sum( + self._tmp_reduce.data[ + self._isc : self._iec + 1, self._jsc : self._jec + 1 + ] + ) + else: + raise TypeError("Unsupported array type for reduction result.") diff --git a/pyFV3/stencils/remapping_GEOS.py b/pyFV3/stencils/remapping_GEOS.py index 83932e3c..2376f4d9 100644 --- a/pyFV3/stencils/remapping_GEOS.py +++ b/pyFV3/stencils/remapping_GEOS.py @@ -1,3 +1,5 @@ +from gt4py.cartesian.gtscript import FORWARD, computation, interval + from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.comm.communicator import Communicator from ndsl.constants import ( @@ -14,11 +16,11 @@ from ndsl.grid import GridData from ndsl.stencils.basic_operations import adjust_divide_stencil from pyFV3._config import RemappingConfig +from pyFV3.mpi.sum import GlobalSum from pyFV3.stencils import moist_cv from pyFV3.stencils.map_single import MapSingle from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.stencils.moist_cv import moist_pt_last_step -from pyFV3.mpi.sum import GlobalSum from pyFV3.stencils.remapping import ( CONSV_MIN, init_pe, @@ -188,6 +190,7 @@ def __init__( self._global_sum = GlobalSum( communicator=comm, quantity_factory=quantity_factory, + grid_indexing=stencil_factory.grid_indexing, ) self._init_pe = stencil_factory.from_origin_domain( @@ -567,7 +570,7 @@ def __call__( tesum: Float = self._global_sum(self._te_2d) zsum: Float = self._global_sum(self._zsum1) - dtmp: Float = tesum / (CV_AIR * zsum) + dtmp = tesum / (CV_AIR * zsum) elif consv_te < -CONSV_MIN: raise NotImplementedError( diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index 558da2d3..b72e99dc 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -51,27 +51,34 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.fill = True + self._are_tracers_setup = False + + self._tracers = None + def compute_from_storage(self, inputs): - tracers = setup_tracers( - number_of_tracers=inputs["qtracers"].shape[3], - quantity_factory=self.quantity_factory, - mappings={"cloud": 6}, - ) - tracers.quantity.data[:-1, :-1, :-1, :] = inputs["qtracers"] + if not self._are_tracers_setup: + self._are_tracers_setup = True + self._tracers = setup_tracers( + number_of_tracers=inputs["qtracers"].shape[3], + quantity_factory=self.quantity_factory, + # mappings={"cloud": 6}, + ) + # tracers.quantity.data[:-1, :-1, :-1, :] = inputs["qtracers"] + self._tracers.quantity.data = inputs["qtracers"] self._compute_func = MapNTracer( self.stencil_factory, self.quantity_factory, abs(self.kord), fill=self.fill, - tracers=tracers, + tracers=self._tracers, ) self._compute_func( inputs["pe1"], inputs["pe2"], inputs["dp2"], - tracers, + self._tracers, ) return inputs diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d.py b/tests/savepoint/translate/translate_moistcvpluspt_2d.py index 3a0748cf..271430c0 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d.py @@ -1,6 +1,7 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval from ndsl import StencilFactory +from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float, FloatField from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j from pyFV3.stencils import moist_cv @@ -53,6 +54,12 @@ def __init__( domain=(grid.nic, 1, grid.npz), ) + self._q_con = grid.quantity_factory.zeros( + [X_DIM, Y_DIM, Z_DIM], + units="unknown", + dtype=Float, + ) + def __call__( self, qvapor: FloatField, @@ -61,7 +68,7 @@ def __call__( qsnow: FloatField, qice: FloatField, qgraupel: FloatField, - q_con: FloatField, + # q_con: FloatField, pt: FloatField, cappa: FloatField, delp: FloatField, @@ -75,7 +82,8 @@ def __call__( qsnow, qice, qgraupel, - q_con, + # q_con, + self._q_con, pt, cappa, delp, @@ -98,7 +106,7 @@ def __init__(self, grid, namelist, stencil_factory): "qgraupel": {"serialname": "qgraupel_js"}, "delp": {}, "delz": {}, - "q_con": {}, + # "q_con": {}, "pt": {}, "cappa": {}, } @@ -111,7 +119,7 @@ def __init__(self, grid, namelist, stencil_factory): self.out_vars = { "pt": {}, "cappa": {}, - "q_con": {}, + # "q_con": {}, } def compute_from_storage(self, inputs): diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py index 46cf6402..f947444a 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py @@ -27,7 +27,7 @@ def __init__(self, grid, namelist, stencil_factory): "kend": grid.npz - 1, }, "pt": {}, - "pkz": {}, + "pkz": {"istart": grid.is_, "jstart": grid.js}, } self.in_vars["parameters"] = ["r_vir", "dtmp"] @@ -61,10 +61,10 @@ def compute_from_storage(self, inputs): inputs["qsnow"], inputs["qice"], inputs["qgraupel"], - self._gz, + # self._gz, inputs["pt"], inputs["pkz"], - inputs["dtmp"], + Float(inputs["dtmp"]), inputs["r_vir"], ) return inputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index d1a9e378..6aa6916f 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -369,6 +369,10 @@ def __init__( self.namelist = DynamicalCoreConfig.from_namelist(namelist) self.grid = grid + self._are_tracers_setup = False + + self._tracers = None + def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") @@ -417,23 +421,26 @@ def compute_parallel(self, inputs, communicator): # inputs["tracers"], # ) - tracers = setup_tracers( - number_of_tracers=inputs["tracers"].shape[3], - quantity_factory=self.quantity_factory, - mappings={ - "vapor": 0, - "liquid": 1, - "rain": 2, - "snow": 3, - "ice": 4, - "graupel": 5, - "cloud": 6, - }, - ) - tracers.quantity.data[:-1, :-1, :-1, :] = inputs["tracers"] + if not self._are_tracers_setup: + self._are_tracers_setup = True + self._tracers = setup_tracers( + number_of_tracers=inputs["tracers"].shape[3], + quantity_factory=self.quantity_factory, + mappings={ + "vapor": 0, + "liquid": 1, + "rain": 3, + "snow": 4, + "ice": 2, + "graupel": 5, + "cloud": 6, + }, + ) + + self._tracers.quantity.data[:-1, :-1, :-1, :] = inputs["tracers"] inputs["te_2d"] = inputs["te_2d"].astype(Float) - state = self.state_from_inputs(inputs, tracers) + state = self.state_from_inputs(inputs, self._tracers) l_to_e = LagrangianToEulerian_GEOS( self.stencil_factory, From 172d4ea589c536f47adfc40e22c66f24ace680eb Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 26 Nov 2025 12:28:36 -0500 Subject: [PATCH 244/252] Update to NDSL 2025.10.00 API (Namelist, Corners) --- pyFV3/__init__.py | 7 + pyFV3/_config.py | 305 ++++++++---------- pyFV3/stencils/corners.py | 169 ++++++++++ pyFV3/stencils/fvtp2d.py | 9 +- pyFV3/testing/translate_dyncore.py | 3 +- pyFV3/testing/translate_fvdynamics.py | 3 +- .../translate/translate_MapN_Tracer_2d.py | 3 +- .../translate/translate_Pressures_mapU.py | 3 +- .../translate/translate_Pressures_mapV.py | 3 +- .../savepoint/translate/translate_a2b_ord4.py | 3 +- tests/savepoint/translate/translate_c_sw.py | 3 +- .../savepoint/translate/translate_corners.py | 3 +- .../translate/translate_cs_profile.py | 3 +- .../translate/translate_cubedtolatlon.py | 3 +- .../translate/translate_d2a2c_vect.py | 3 +- tests/savepoint/translate/translate_d_sw.py | 3 +- .../translate/translate_del2cubed.py | 3 +- .../translate/translate_del6vtflux.py | 3 +- .../savepoint/translate/translate_delnflux.py | 3 +- .../translate/translate_divergencedamping.py | 3 +- tests/savepoint/translate/translate_fillz.py | 5 +- .../translate/translate_fvsubgridz.py | 3 +- tests/savepoint/translate/translate_fvtp2d.py | 3 +- tests/savepoint/translate/translate_fxadv.py | 3 +- .../translate/translate_getMPIprop.py | 3 +- tests/savepoint/translate/translate_grid.py | 3 +- .../translate/translate_haloupdate.py | 3 +- .../translate/translate_init_case.py | 3 +- ...ranslate_lagrangian_contribution_interp.py | 4 +- .../translate/translate_last_step.py | 3 +- .../translate/translate_map1_ppm_W.py | 3 +- .../translate/translate_map1_ppm_delz.py | 3 +- .../translate/translate_map_scalar.py | 3 +- .../translate/translate_moistcvpluspkz_2d.py | 3 +- .../translate/translate_mpp_global_sum.py | 3 +- .../savepoint/translate/translate_neg_adj3.py | 3 +- .../translate/translate_nh_p_grad.py | 3 +- .../savepoint/translate/translate_pe_halo.py | 3 +- .../translate/translate_pe_pk_delp_peln.py | 4 +- .../savepoint/translate/translate_pk3_halo.py | 3 +- ...ssureadjustedtemperature_nonhydrostatic.py | 3 +- tests/savepoint/translate/translate_qsinit.py | 3 +- .../savepoint/translate/translate_ray_fast.py | 3 +- .../translate/translate_remapping.py | 3 +- .../translate/translate_remapping_GEOS.py | 3 +- .../translate/translate_riem_solver3.py | 3 +- .../translate/translate_riem_solver_c.py | 3 +- .../translate/translate_satadjust3d.py | 3 +- .../translate/translate_scalar_profile.py | 3 +- .../translate/translate_tracer2d1l.py | 10 +- .../translate/translate_tracer2d1l_cmax.py | 3 +- .../translate/translate_updatedzc.py | 3 +- .../translate/translate_updatedzd.py | 3 +- tests/savepoint/translate/translate_xppm.py | 3 +- tests/savepoint/translate/translate_xtp_u.py | 3 +- tests/savepoint/translate/translate_yppm.py | 3 +- tests/savepoint/translate/translate_ytp_v.py | 3 +- 57 files changed, 421 insertions(+), 239 deletions(-) create mode 100644 pyFV3/stencils/corners.py diff --git a/pyFV3/__init__.py b/pyFV3/__init__.py index 4024a8f1..c8081971 100644 --- a/pyFV3/__init__.py +++ b/pyFV3/__init__.py @@ -11,3 +11,10 @@ """ __version__ = "0.2.0" + +__all__ = [ + "DynamicalCoreConfig", + "DycoreState", + "DryConvectiveAdjustment", + "DynamicalCore", +] diff --git a/pyFV3/_config.py b/pyFV3/_config.py index dcb2e8de..0f226979 100644 --- a/pyFV3/_config.py +++ b/pyFV3/_config.py @@ -1,18 +1,25 @@ +from __future__ import annotations + import dataclasses from datetime import timedelta from math import floor -from typing import Optional, Tuple import f90nml import yaml +from dacite import Config, from_dict -from ndsl.namelist import Namelist, NamelistDefaults +from ndsl.utils import f90nml_as_dict DEFAULT_INT = 0 DEFAULT_STR = "" DEFAULT_FLOAT = 0.0 DEFAULT_BOOL = False +DEFAULT_DYCORE_NML_GROUPS = ( + "main_nml", + "coupler_nml", + "fv_core_nml", +) @dataclasses.dataclass(frozen=True) @@ -112,7 +119,6 @@ class AcousticDynamicsConfig: """ mainly for backwards compatibility, not really used anymore """ - dz_min: float riemann: RiemannConfig d_grid_shallow_water: DGridShallowWaterLagrangianDynamicsConfig @@ -196,97 +202,101 @@ class DynamicalCoreConfig: vtdm4: float = DEFAULT_FLOAT z_tracer: bool = DEFAULT_BOOL do_qa: bool = DEFAULT_BOOL - layout: Tuple[int, int] = NamelistDefaults.layout - grid_type: int = NamelistDefaults.grid_type - u_max: float = NamelistDefaults.u_max # max windspeed for dp config - do_f3d: bool = NamelistDefaults.do_f3d - inline_q: bool = NamelistDefaults.inline_q - do_skeb: bool = NamelistDefaults.do_skeb # save dissipation estimate - use_logp: bool = NamelistDefaults.use_logp - moist_phys: bool = NamelistDefaults.moist_phys - check_negative: bool = NamelistDefaults.check_negative + layout: tuple[int, int] = (1, 1) + grid_type: int = 0 + u_max: float = 350.0 + """max windspeed for dp config""" + do_f3d: bool = False + inline_q: bool = False + do_skeb: bool = False + """save dissipation estimate""" + use_logp: bool = False + moist_phys: bool = True + check_negative: bool = False # gfdl_cloud_microphys.F90 - tau_r2g: float = NamelistDefaults.tau_r2g # rain freezing during fast_sat - tau_smlt: float = NamelistDefaults.tau_smlt # snow melting - tau_g2r: float = NamelistDefaults.tau_g2r # graupel melting to rain - tau_imlt: float = NamelistDefaults.tau_imlt # cloud ice melting - tau_i2s: float = NamelistDefaults.tau_i2s # cloud ice to snow auto - conversion - tau_l2r: float = NamelistDefaults.tau_l2r # cloud water to rain auto - conversion - tau_g2v: float = NamelistDefaults.tau_g2v # graupel sublimation - tau_v2g: float = ( - NamelistDefaults.tau_v2g - ) # graupel deposition -- make it a slow process - sat_adj0: float = ( - NamelistDefaults.sat_adj0 - ) # adjustment factor (0: no 1: full) during fast_sat_adj + tau_r2g: float = 900.0 + """rain freezing during fast_sat""" + tau_smlt: float = 900.0 + """snow melting""" + tau_g2r: float = 600.0 + """graupel melting to rain""" + tau_imlt: float = 600.0 + """cloud ice melting""" + tau_i2s: float = 1000.0 + """cloud ice to snow auto - conversion""" + tau_l2r: float = 900.0 + """cloud water to rain auto - conversion""" + tau_g2v: float = 1200.0 + """graupel sublimation""" + tau_v2g: float = 21600.0 + """graupel deposition -- make it a slow process""" + sat_adj0: float = 0.90 + """adjustment factor (0: no 1: full) during fast_sat_adj""" ql_gen: float = ( 1.0e-3 # max new cloud water during remapping step if fast_sat_adj = .t. ) - ql_mlt: float = ( - NamelistDefaults.ql_mlt - ) # max value of cloud water allowed from melted cloud ice - qs_mlt: float = NamelistDefaults.qs_mlt # max cloud water due to snow melt - ql0_max: float = ( - NamelistDefaults.ql0_max - ) # max cloud water value (auto converted to rain) - t_sub: float = NamelistDefaults.t_sub # min temp for sublimation of cloud ice - qi_gen: float = ( - NamelistDefaults.qi_gen - ) # max cloud ice generation during remapping step - qi_lim: float = ( - NamelistDefaults.qi_lim - ) # cloud ice limiter to prevent large ice build up - qi0_max: float = NamelistDefaults.qi0_max # max cloud ice value (by other sources) - rad_snow: bool = ( - NamelistDefaults.rad_snow - ) # consider snow in cloud fraction calculation - rad_rain: bool = ( - NamelistDefaults.rad_rain - ) # consider rain in cloud fraction calculation - rad_graupel: bool = ( - NamelistDefaults.rad_graupel - ) # consider graupel in cloud fraction calculation - tintqs: bool = ( - NamelistDefaults.tintqs - ) # use temperature in the saturation mixing in PDF - dw_ocean: float = NamelistDefaults.dw_ocean # base value for ocean - dw_land: float = ( - NamelistDefaults.dw_land - ) # base value for subgrid deviation / variability over land + ql_mlt: float = 2.0e-3 + """max value of cloud water allowed from melted cloud ice""" + qs_mlt: float = 1.0e-6 + """max cloud water due to snow melt""" + ql0_max: float = 2.0e-3 + """max cloud water value (auto converted to rain)""" + t_sub: float = 184.0 + """min temp for sublimation of cloud ice""" + qi_gen: float = 1.82e-6 + """max cloud ice generation during remapping step""" + qi_lim: float = 1.0 + """cloud ice limiter to prevent large ice build up""" + qi0_max: float = 1.0e-4 + """max cloud ice value (by other sources)""" + rad_snow: bool = True + """consider snow in cloud fraction calculation""" + rad_rain: bool = True + """consider rain in cloud fraction calculation""" + rad_graupel: bool = True + """consider graupel in cloud fraction calculation""" + tintqs: bool = False + """use temperature in the saturation mixing in PDF""" + dw_ocean: float = 0.10 + """base value for ocean""" + dw_land: float = 0.15 + """base value for subgrid deviation / variability over land""" # cloud scheme 0 - ? # 1: old fvgfs gfdl) mp implementation # 2: binary cloud scheme (0 / 1) - icloud_f: int = NamelistDefaults.icloud_f - cld_min: float = NamelistDefaults.cld_min # !< minimum cloud fraction - tau_l2v: float = ( - NamelistDefaults.tau_l2v - ) # cloud water to water vapor (evaporation) - tau_v2l: float = ( - NamelistDefaults.tau_v2l - ) # water vapor to cloud water (condensation) - c2l_ord: int = NamelistDefaults.c2l_ord - regional: bool = NamelistDefaults.regional - m_split: int = NamelistDefaults.m_split - convert_ke: bool = NamelistDefaults.convert_ke - breed_vortex_inline: bool = NamelistDefaults.breed_vortex_inline - use_old_omega: bool = NamelistDefaults.use_old_omega - rf_fast: bool = NamelistDefaults.rf_fast - adiabatic: bool = NamelistDefaults.adiabatic - nf_omega: int = NamelistDefaults.nf_omega - fv_sg_adj: int = NamelistDefaults.fv_sg_adj - n_sponge: int = NamelistDefaults.n_sponge - dz_min: float = 2.0 - """Controls minimum thickness in NH solver""" - namelist_override: Optional[str] = None - - def __post_init__(self): + icloud_f: int = 0 + cld_min: float = 0.05 + """!< minimum cloud fraction""" + tau_l2v: float = 300.0 + """cloud water to water vapor (evaporation)""" + tau_v2l: float = 90.0 + """water vapor to cloud water (condensation)""" + c2l_ord: int = 4 + regional: bool = False + m_split: int = 0 + convert_ke: bool = False + breed_vortex_inline: bool = False + use_old_omega: bool = True + rf_fast: bool = False + adiabatic: bool = False + nf_omega: int = 1 + fv_sg_adj: int = -1 + n_sponge: int = 1 + sw_dynamics: bool = False + """shallow water conditions""" + namelist_override: str | None = None + target_nml_groups: tuple[str, ...] | None = DEFAULT_DYCORE_NML_GROUPS + + def __post_init__(self) -> None: if self.namelist_override is not None: try: f90_nml = f90nml.read(self.namelist_override) except FileNotFoundError: print(f"{self.namelist_override} does not exist") raise - dycore_config = self.from_f90nml(f90_nml) + # TODO: Find a better way to do below. Passing self.* as an argument + # to a class function of the same class is always a bit fishy. + dycore_config = self.from_f90nml(f90_nml, self.target_nml_groups) for var in dycore_config.__dict__.keys(): setattr(self, var, dycore_config.__dict__[var]) # Single tile cartesian grids @@ -294,102 +304,52 @@ def __post_init__(self): self.nf_omega = 0 @classmethod - def from_f90nml(cls, f90_namelist: f90nml.Namelist) -> "DynamicalCoreConfig": - namelist = Namelist.from_f90nml(f90_namelist) - return cls.from_namelist(namelist) + def from_f90nml( + cls, + nml: f90nml.Namelist, + target_groups: tuple[str, ...] | None = DEFAULT_DYCORE_NML_GROUPS, + ) -> DynamicalCoreConfig: + """Uses the nml to create a DynamicalCoreConfig. + + Args: + nml: f90nml.Namelist + target_groups: tuple[str,...] | None + This list will be used to specify which groups in the nml to + use when initializing the DynamicalCoreConfig. If None, all + groups will be used. (Default: DEFAULT_DYCORE_NML_GROUPS) + """ + groups = list(target_groups) if target_groups is not None else None + nml_dict = f90nml_as_dict(nml, flatten=True, target_groups=groups) + nml_dict["target_nml_groups"] = target_groups + return cls.from_dict(nml_dict) @classmethod - def from_namelist(cls, namelist: Namelist) -> "DynamicalCoreConfig": - return cls( - dt_atmos=namelist.dt_atmos, - a_imp=namelist.a_imp, - beta=namelist.beta, - consv_te=namelist.consv_te, - d2_bg=namelist.d2_bg, - d2_bg_k1=namelist.d2_bg_k1, - d2_bg_k2=namelist.d2_bg_k2, - d4_bg=namelist.d4_bg, - d_con=namelist.d_con, - d_ext=namelist.d_ext, - dddmp=namelist.dddmp, - delt_max=namelist.delt_max, - do_sat_adj=namelist.do_sat_adj, - do_vort_damp=namelist.do_vort_damp, - fill=namelist.fill, - hord_dp=namelist.hord_dp, - hord_mt=namelist.hord_mt, - hord_tm=namelist.hord_tm, - hord_tr=namelist.hord_tr, - hord_vt=namelist.hord_vt, - hydrostatic=namelist.hydrostatic, - k_split=namelist.k_split, - ke_bg=namelist.ke_bg, - kord_mt=namelist.kord_mt, - kord_tm=namelist.kord_tm, - kord_tr=namelist.kord_tr, - kord_wz=namelist.kord_wz, - n_split=namelist.n_split, - nord=namelist.nord, - npx=namelist.npx, - npy=namelist.npy, - npz=namelist.npz, - ntiles=namelist.ntiles, - nwat=namelist.nwat, - p_fac=namelist.p_fac, - rf_cutoff=namelist.rf_cutoff, - tau=namelist.tau, - vtdm4=namelist.vtdm4, - z_tracer=namelist.z_tracer, - do_qa=namelist.do_qa, - layout=namelist.layout, - grid_type=namelist.grid_type, - u_max=namelist.u_max, - do_f3d=namelist.do_f3d, - inline_q=namelist.inline_q, - do_skeb=namelist.do_skeb, - check_negative=namelist.check_negative, - tau_r2g=namelist.tau_r2g, - tau_smlt=namelist.tau_smlt, - tau_g2r=namelist.tau_g2r, - tau_imlt=namelist.tau_imlt, - tau_i2s=namelist.tau_i2s, - tau_l2r=namelist.tau_l2r, - tau_g2v=namelist.tau_g2v, - tau_v2g=namelist.tau_v2g, - sat_adj0=namelist.sat_adj0, - ql_gen=namelist.ql_gen, - ql_mlt=namelist.ql_mlt, - qs_mlt=namelist.qs_mlt, - ql0_max=namelist.ql0_max, - t_sub=namelist.t_sub, - qi_gen=namelist.qi_gen, - qi_lim=namelist.qi_lim, - qi0_max=namelist.qi0_max, - rad_snow=namelist.rad_snow, - rad_rain=namelist.rad_rain, - rad_graupel=namelist.rad_graupel, - tintqs=namelist.tintqs, - dw_ocean=namelist.dw_ocean, - dw_land=namelist.dw_land, - icloud_f=namelist.icloud_f, - cld_min=namelist.cld_min, - tau_l2v=namelist.tau_l2v, - tau_v2l=namelist.tau_v2l, - c2l_ord=namelist.c2l_ord, - regional=namelist.regional, - m_split=namelist.m_split, - convert_ke=namelist.convert_ke, - breed_vortex_inline=namelist.breed_vortex_inline, - use_old_omega=namelist.use_old_omega, - rf_fast=namelist.rf_fast, - adiabatic=namelist.adiabatic, - nf_omega=namelist.nf_omega, - fv_sg_adj=namelist.fv_sg_adj, - n_sponge=namelist.n_sponge, + def from_dict( + cls, + data: dict, + ) -> DynamicalCoreConfig: + """Create a DynamicalCoreConfig from the given data. + + Args: + data: "flattened" dictionary where the keys match the class member variables + """ + # NOTE: We're setting strict to False so that extra keys in the data are + # ignored. Eventually, we'd like to turn this to True once we move away from + # expecting dicts that are basically flattened yamls and f90nml files. + dacite_config = Config( + strict=False, + type_hooks={ + tuple[int, int]: lambda x: tuple(x), + tuple[str, ...]: lambda x: tuple(x) if x is not None else None, + }, + ) + dycore_config = from_dict( + data_class=DynamicalCoreConfig, data=data, config=dacite_config ) + return dycore_config @classmethod - def from_yaml(cls, yaml_config: str) -> "DynamicalCoreConfig": + def from_yaml(cls, yaml_config: str) -> DynamicalCoreConfig: config = cls() with open(yaml_config, "r") as f: raw_config = yaml.safe_load(f) @@ -491,7 +451,6 @@ def acoustic_dynamics(self) -> AcousticDynamicsConfig: use_old_omega=self.use_old_omega, riemann=self.riemann, d_grid_shallow_water=self.d_grid_shallow_water, - dz_min=self.dz_min, ) @property diff --git a/pyFV3/stencils/corners.py b/pyFV3/stencils/corners.py new file mode 100644 index 00000000..266627df --- /dev/null +++ b/pyFV3/stencils/corners.py @@ -0,0 +1,169 @@ +from ndsl import StencilFactory, orchestrate +from ndsl.dsl.typing import FloatField + + +def corner_copy_x(field_to_copy): + """Equivalent to the copy_corners_x functions in fortran. + + This is written to operate on plain ndarrarys and not use the GT4Py framework. + This choice was made because we've seen a lot of performance left on the table using + orchestration without explicitly describing the operations but rather have full 3d- + sweeps with conditionals. + Since DaCe can handle (simple) operations on ndarrays directly this gives us a more + explicit entrypoint to the language and more optimization-potential. + + Args: + field_to_copy (ndarray): field to apply the corner copy on. + This is explicitly not type-hinted for orchestration + """ + field_to_copy[0, 0] = field_to_copy[0, 5] + field_to_copy[0, 1] = field_to_copy[1, 5] + field_to_copy[0, 2] = field_to_copy[2, 5] + + field_to_copy[1, 0] = field_to_copy[0, 4] + field_to_copy[1, 1] = field_to_copy[1, 4] + field_to_copy[1, 2] = field_to_copy[2, 4] + + field_to_copy[2, 0] = field_to_copy[0, 3] + field_to_copy[2, 1] = field_to_copy[1, 3] + field_to_copy[2, 2] = field_to_copy[2, 3] + + field_to_copy[0, -4] = field_to_copy[2, -7] + field_to_copy[0, -3] = field_to_copy[1, -7] + field_to_copy[0, -2] = field_to_copy[0, -7] + + field_to_copy[1, -4] = field_to_copy[2, -6] + field_to_copy[1, -3] = field_to_copy[1, -6] + field_to_copy[1, -2] = field_to_copy[0, -6] + + field_to_copy[2, -4] = field_to_copy[2, -5] + field_to_copy[2, -3] = field_to_copy[1, -5] + field_to_copy[2, -2] = field_to_copy[0, -5] + + field_to_copy[-4, 0] = field_to_copy[-2, 3] + field_to_copy[-4, 1] = field_to_copy[-3, 3] + field_to_copy[-4, 2] = field_to_copy[-4, 3] + + field_to_copy[-3, 0] = field_to_copy[-2, 4] + field_to_copy[-3, 1] = field_to_copy[-3, 4] + field_to_copy[-3, 2] = field_to_copy[-4, 4] + + field_to_copy[-2, 0] = field_to_copy[-2, 5] + field_to_copy[-2, 1] = field_to_copy[-3, 5] + field_to_copy[-2, 2] = field_to_copy[-4, 5] + + field_to_copy[-4, -2] = field_to_copy[-2, -5] + field_to_copy[-4, -3] = field_to_copy[-3, -5] + field_to_copy[-4, -4] = field_to_copy[-4, -5] + + field_to_copy[-3, -2] = field_to_copy[-2, -6] + field_to_copy[-3, -3] = field_to_copy[-3, -6] + field_to_copy[-3, -4] = field_to_copy[-4, -6] + + field_to_copy[-2, -2] = field_to_copy[-2, -7] + field_to_copy[-2, -3] = field_to_copy[-3, -7] + field_to_copy[-2, -4] = field_to_copy[-4, -7] + + +def corner_copy_y(field_to_copy): + """Equivalent to the copy_corners_y functions in fortran. + + This is written to operate on plain ndarrarys and not use the GT4Py framework. + This choice was made because we've seen a lot of performance left on the table using + orchestration without explicitly describing the operations but rather have full 3d- + sweeps with conditionals. + Since DaCe can handle (simple) operations on ndarrays directly this gives us a more + explicit entrypoint to the language and more optimization-potential. + + Args: + field_to_copy (ndarray): field to apply the corner copy on. + This is explicitly not type-hinted for orchestration + """ + field_to_copy[0, 0] = field_to_copy[5, 0] + field_to_copy[1, 0] = field_to_copy[5, 1] + field_to_copy[2, 0] = field_to_copy[5, 2] + + field_to_copy[0, 1] = field_to_copy[4, 0] + field_to_copy[1, 1] = field_to_copy[4, 1] + field_to_copy[2, 1] = field_to_copy[4, 2] + + field_to_copy[0, 2] = field_to_copy[3, 0] + field_to_copy[1, 2] = field_to_copy[3, 1] + field_to_copy[2, 2] = field_to_copy[3, 2] + + field_to_copy[-4, 0] = field_to_copy[-7, 2] + field_to_copy[-3, 0] = field_to_copy[-7, 1] + field_to_copy[-2, 0] = field_to_copy[-7, 0] + + field_to_copy[-4, 1] = field_to_copy[-6, 2] + field_to_copy[-3, 1] = field_to_copy[-6, 1] + field_to_copy[-2, 1] = field_to_copy[-6, 0] + + field_to_copy[-4, 2] = field_to_copy[-5, 2] + field_to_copy[-3, 2] = field_to_copy[-5, 1] + field_to_copy[-2, 2] = field_to_copy[-5, 0] + + field_to_copy[0, -2] = field_to_copy[5, -2] + field_to_copy[0, -3] = field_to_copy[4, -2] + field_to_copy[0, -4] = field_to_copy[3, -2] + + field_to_copy[1, -2] = field_to_copy[5, -3] + field_to_copy[1, -3] = field_to_copy[4, -3] + field_to_copy[1, -4] = field_to_copy[3, -3] + + field_to_copy[2, -2] = field_to_copy[5, -4] + field_to_copy[2, -3] = field_to_copy[4, -4] + field_to_copy[2, -4] = field_to_copy[3, -4] + + field_to_copy[-2, -4] = field_to_copy[-5, -2] + field_to_copy[-2, -3] = field_to_copy[-6, -2] + field_to_copy[-2, -2] = field_to_copy[-7, -2] + + field_to_copy[-3, -4] = field_to_copy[-5, -3] + field_to_copy[-3, -3] = field_to_copy[-6, -3] + field_to_copy[-3, -2] = field_to_copy[-7, -3] + + field_to_copy[-4, -4] = field_to_copy[-5, -4] + field_to_copy[-4, -3] = field_to_copy[-6, -4] + field_to_copy[-4, -2] = field_to_copy[-7, -4] + + +class CopyCornersX: + """ + Helper-class to copy corners corresponding to the fortran function copy_corners_x + """ + + def __init__(self, stencil_factory: StencilFactory) -> None: + orchestrate( + obj=self, + config=stencil_factory.config.dace_config, + ) + + if stencil_factory.grid_indexing.n_halo != 3: + raise NotImplementedError( + "Corner-Copy only implemented for exactly 3 Halo-Points" + ) + + def __call__(self, field: FloatField): + corner_copy_x(field) + + +class CopyCornersY: + """ + Helper-class to copy corners corresponding to the fortran function + copy_corners_y + """ + + def __init__(self, stencil_factory: StencilFactory) -> None: + orchestrate( + obj=self, + config=stencil_factory.config.dace_config, + ) + + if stencil_factory.grid_indexing.n_halo != 3: + raise NotImplementedError( + "Corner-Copy only implemented for exactly 3 Halo-Points" + ) + + def __call__(self, field: FloatField): + corner_copy_y(field) diff --git a/pyFV3/stencils/fvtp2d.py b/pyFV3/stencils/fvtp2d.py index 4e5fe3e2..d3e5b356 100644 --- a/pyFV3/stencils/fvtp2d.py +++ b/pyFV3/stencils/fvtp2d.py @@ -11,6 +11,7 @@ from pyFV3.stencils.delnflux import DelnFlux from pyFV3.stencils.xppm import XPiecewiseParabolic from pyFV3.stencils.yppm import YPiecewiseParabolic +from pyFV3.stencils.corners import CopyCornersX, CopyCornersY @gtscript.function @@ -179,9 +180,7 @@ def make_quantity(): # self.delnflux = None self._do_delnflux = False - self._copy_corners_y: corners.CopyCorners = corners.CopyCorners( - "y", stencil_factory - ) + self._copy_corners_y = CopyCornersY(stencil_factory) self.y_piecewise_parabolic_inner = YPiecewiseParabolic( stencil_factory=stencil_factory, dya=grid_data.dya, @@ -204,9 +203,7 @@ def make_quantity(): domain=idx.domain_compute(add=(1, 1, 1)), ) - self._copy_corners_x: corners.CopyCorners = corners.CopyCorners( - "x", stencil_factory - ) + self._copy_corners_x = CopyCornersX(stencil_factory) self.x_piecewise_parabolic_inner = XPiecewiseParabolic( stencil_factory=stencil_factory, dxa=grid_data.dxa, diff --git a/pyFV3/testing/translate_dyncore.py b/pyFV3/testing/translate_dyncore.py index 0e1703f6..1ae8d2cd 100644 --- a/pyFV3/testing/translate_dyncore.py +++ b/pyFV3/testing/translate_dyncore.py @@ -1,5 +1,6 @@ import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, Quantity, StencilFactory +from ndsl import Quantity, StencilFactory +from f90nml.namelist import Namelist from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.stencils.testing import ParallelTranslate2PyState from numpy import dtype diff --git a/pyFV3/testing/translate_fvdynamics.py b/pyFV3/testing/translate_fvdynamics.py index 4951ecbf..ed78141d 100644 --- a/pyFV3/testing/translate_fvdynamics.py +++ b/pyFV3/testing/translate_fvdynamics.py @@ -5,7 +5,8 @@ import pytest import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, Quantity, QuantityFactory, StencilFactory, FieldBundle +from ndsl import Quantity, QuantityFactory, StencilFactory, FieldBundle +from f90nml import Namelist from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index b72e99dc..a13c5d8e 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.mapn_tracer import MapNTracer diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index b2790161..e4f6c85b 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index 687243ef..81620694 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_a2b_ord4.py b/tests/savepoint/translate/translate_a2b_ord4.py index dfba1ed0..92ec2e0e 100644 --- a/tests/savepoint/translate/translate_a2b_ord4.py +++ b/tests/savepoint/translate/translate_a2b_ord4.py @@ -1,7 +1,8 @@ import numpy as np from typing import Any, Dict -from ndsl import Namelist, StencilFactory, orchestrate +from ndsl import StencilFactory, orchestrate +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from pyFV3.stencils import DivergenceDamping from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_c_sw.py b/tests/savepoint/translate/translate_c_sw.py index 13abd93c..b7255525 100644 --- a/tests/savepoint/translate/translate_c_sw.py +++ b/tests/savepoint/translate/translate_c_sw.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, QuantityFactory, StencilFactory +from ndsl import QuantityFactory, StencilFactory +from f90nml.namelist import Namelist from pyFV3.stencils import CGridShallowWaterDynamics from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_corners.py b/tests/savepoint/translate/translate_corners.py index 4aa58497..7bcd8f39 100644 --- a/tests/savepoint/translate/translate_corners.py +++ b/tests/savepoint/translate/translate_corners.py @@ -1,7 +1,8 @@ from typing import Any, Dict import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.stencils import corners from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_cs_profile.py b/tests/savepoint/translate/translate_cs_profile.py index c63617f2..d256c507 100644 --- a/tests/savepoint/translate/translate_cs_profile.py +++ b/tests/savepoint/translate/translate_cs_profile.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_cubedtolatlon.py b/tests/savepoint/translate/translate_cubedtolatlon.py index 6610a46e..32b15779 100644 --- a/tests/savepoint/translate/translate_cubedtolatlon.py +++ b/tests/savepoint/translate/translate_cubedtolatlon.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, Quantity, StencilFactory +from ndsl import Quantity, StencilFactory +from f90nml.namelist import Namelist from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.stencils.c2l_ord import CubedToLatLon from ndsl.stencils.testing import ParallelTranslate2Py diff --git a/tests/savepoint/translate/translate_d2a2c_vect.py b/tests/savepoint/translate/translate_d2a2c_vect.py index d3561e42..c081f110 100644 --- a/tests/savepoint/translate/translate_d2a2c_vect.py +++ b/tests/savepoint/translate/translate_d2a2c_vect.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.stencils import DGrid2AGrid2CGridVectors from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_d_sw.py b/tests/savepoint/translate/translate_d_sw.py index 34f34cd7..d9508cbd 100644 --- a/tests/savepoint/translate/translate_d_sw.py +++ b/tests/savepoint/translate/translate_d_sw.py @@ -2,7 +2,8 @@ import pyFV3 import pyFV3.stencils.d_sw as d_sw -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_del2cubed.py b/tests/savepoint/translate/translate_del2cubed.py index 07bf24c2..cc33515c 100644 --- a/tests/savepoint/translate/translate_del2cubed.py +++ b/tests/savepoint/translate/translate_del2cubed.py @@ -1,6 +1,7 @@ from typing import Any, Dict -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.stencils import HyperdiffusionDamping from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_del6vtflux.py b/tests/savepoint/translate/translate_del6vtflux.py index f163c83b..ae500403 100644 --- a/tests/savepoint/translate/translate_del6vtflux.py +++ b/tests/savepoint/translate/translate_del6vtflux.py @@ -1,5 +1,6 @@ import pyFV3.stencils.delnflux as delnflux -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import Z_DIM from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_delnflux.py b/tests/savepoint/translate/translate_delnflux.py index 1db0b33c..4ca8bbc9 100644 --- a/tests/savepoint/translate/translate_delnflux.py +++ b/tests/savepoint/translate/translate_delnflux.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import Z_DIM from pyFV3.stencils import delnflux from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_divergencedamping.py b/tests/savepoint/translate/translate_divergencedamping.py index 001bec4a..8d5cbedb 100644 --- a/tests/savepoint/translate/translate_divergencedamping.py +++ b/tests/savepoint/translate/translate_divergencedamping.py @@ -1,6 +1,7 @@ from typing import Optional -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import Z_DIM from pyFV3.stencils import DivergenceDamping from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_fillz.py b/tests/savepoint/translate/translate_fillz.py index 605b32d7..eb2377be 100644 --- a/tests/savepoint/translate/translate_fillz.py +++ b/tests/savepoint/translate/translate_fillz.py @@ -1,8 +1,7 @@ -from typing import List - import numpy as np -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.stencils.testing import pad_field_in_j from pyFV3.stencils.fillz import FillNegativeTracerValues from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_fvsubgridz.py b/tests/savepoint/translate/translate_fvsubgridz.py index 726bb7be..c427f984 100644 --- a/tests/savepoint/translate/translate_fvsubgridz.py +++ b/tests/savepoint/translate/translate_fvsubgridz.py @@ -1,7 +1,8 @@ from types import SimpleNamespace import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM from ndsl.stencils.testing import ParallelTranslateBaseSlicing from pyFV3 import DryConvectiveAdjustment diff --git a/tests/savepoint/translate/translate_fvtp2d.py b/tests/savepoint/translate/translate_fvtp2d.py index 81a0cff4..55762dda 100644 --- a/tests/savepoint/translate/translate_fvtp2d.py +++ b/tests/savepoint/translate/translate_fvtp2d.py @@ -1,5 +1,6 @@ import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import Z_DIM from ndsl.dsl.typing import Float from pyFV3.stencils import FiniteVolumeTransport diff --git a/tests/savepoint/translate/translate_fxadv.py b/tests/savepoint/translate/translate_fxadv.py index 76afe3b2..324d9e1d 100644 --- a/tests/savepoint/translate/translate_fxadv.py +++ b/tests/savepoint/translate/translate_fxadv.py @@ -1,6 +1,7 @@ import numpy as np -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from pyFV3.stencils import FiniteVolumeFluxPrep from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_getMPIprop.py b/tests/savepoint/translate/translate_getMPIprop.py index 236ca062..561c8b5c 100644 --- a/tests/savepoint/translate/translate_getMPIprop.py +++ b/tests/savepoint/translate/translate_getMPIprop.py @@ -1,6 +1,7 @@ import numpy as np -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.quantity import Quantity from ndsl.stencils.testing import ParallelTranslate from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_grid.py b/tests/savepoint/translate/translate_grid.py index 25d77b21..3e7f94ff 100644 --- a/tests/savepoint/translate/translate_grid.py +++ b/tests/savepoint/translate/translate_grid.py @@ -4,7 +4,8 @@ import pytest import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, diff --git a/tests/savepoint/translate/translate_haloupdate.py b/tests/savepoint/translate/translate_haloupdate.py index 12b078d7..1591edef 100644 --- a/tests/savepoint/translate/translate_haloupdate.py +++ b/tests/savepoint/translate/translate_haloupdate.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import ( N_HALO_DEFAULT, X_DIM, diff --git a/tests/savepoint/translate/translate_init_case.py b/tests/savepoint/translate/translate_init_case.py index 06d8fcc0..41391883 100644 --- a/tests/savepoint/translate/translate_init_case.py +++ b/tests/savepoint/translate/translate_init_case.py @@ -7,7 +7,8 @@ import pyFV3.initialization.analytic_init as analytic_init import pyFV3.initialization.init_utils as init_utils import pyFV3.initialization.test_cases.initialize_baroclinic as baroclinic_init -from ndsl import Namelist, Quantity, QuantityFactory, StencilFactory, SubtileGridSizer +from ndsl import Quantity, QuantityFactory, StencilFactory, SubtileGridSizer +from f90nml import Namelist from ndsl.constants import ( N_HALO_DEFAULT, X_DIM, diff --git a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py index 33b5f7d6..e65753a1 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Bool, BoolFieldIJ, FloatField, Int, IntField, IntFieldIJ from ndsl.stencils.testing import TranslateFortranData2Py @@ -127,7 +128,6 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto } def compute_from_storage(self, inputs): - self._not_exit_loop = self.quantity_factory.zeros( [X_DIM, Y_DIM], units="", diff --git a/tests/savepoint/translate/translate_last_step.py b/tests/savepoint/translate/translate_last_step.py index 1d2d7525..c6a9a4df 100644 --- a/tests/savepoint/translate/translate_last_step.py +++ b/tests/savepoint/translate/translate_last_step.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.stencils import moist_cv from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_map1_ppm_W.py b/tests/savepoint/translate/translate_map1_ppm_W.py index adbb0fc9..07ca73da 100644 --- a/tests/savepoint/translate/translate_map1_ppm_W.py +++ b/tests/savepoint/translate/translate_map1_ppm_W.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py index 837ad616..9384076e 100644 --- a/tests/savepoint/translate/translate_map1_ppm_delz.py +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -1,6 +1,7 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import FloatField from ndsl.stencils.testing import TranslateFortranData2Py diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py index e0683376..35968971 100644 --- a/tests/savepoint/translate/translate_map_scalar.py +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_moistcvpluspkz_2d.py b/tests/savepoint/translate/translate_moistcvpluspkz_2d.py index 657d27dc..e7331103 100644 --- a/tests/savepoint/translate/translate_moistcvpluspkz_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluspkz_2d.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.dsl.typing import FloatField from ndsl.stencils.testing import pad_field_in_j from pyFV3.stencils import moist_cv diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index 54c57d82..397b6d2c 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.stencils.testing import ParallelTranslate from ndsl.stencils.testing.grid import Grid from ndsl.typing import Communicator diff --git a/tests/savepoint/translate/translate_neg_adj3.py b/tests/savepoint/translate/translate_neg_adj3.py index f7ba63b1..b7f34ac6 100644 --- a/tests/savepoint/translate/translate_neg_adj3.py +++ b/tests/savepoint/translate/translate_neg_adj3.py @@ -1,7 +1,8 @@ from typing import Any, Dict import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.stencils import AdjustNegativeTracerMixingRatio from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_nh_p_grad.py b/tests/savepoint/translate/translate_nh_p_grad.py index 8220f83a..6ec3b424 100644 --- a/tests/savepoint/translate/translate_nh_p_grad.py +++ b/tests/savepoint/translate/translate_nh_p_grad.py @@ -1,5 +1,6 @@ import pyFV3.stencils.nh_p_grad as NH_P_Grad -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_pe_halo.py b/tests/savepoint/translate/translate_pe_halo.py index 0c5c2e09..a5ed4596 100644 --- a/tests/savepoint/translate/translate_pe_halo.py +++ b/tests/savepoint/translate/translate_pe_halo.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.stencils import pe_halo from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_pe_pk_delp_peln.py b/tests/savepoint/translate/translate_pe_pk_delp_peln.py index 89dab68c..bbf80271 100644 --- a/tests/savepoint/translate/translate_pe_pk_delp_peln.py +++ b/tests/savepoint/translate/translate_pe_pk_delp_peln.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.remapping import pe_pk_delp_peln @@ -135,7 +136,6 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto ) def compute_from_storage(self, inputs): - self._pe_pk_delp_peln( inputs["pe_"], inputs["pk"], diff --git a/tests/savepoint/translate/translate_pk3_halo.py b/tests/savepoint/translate/translate_pk3_halo.py index 4f13c02d..35475b1f 100644 --- a/tests/savepoint/translate/translate_pk3_halo.py +++ b/tests/savepoint/translate/translate_pk3_halo.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.stencils import PK3Halo from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py b/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py index 7f9afeac..c8e47a9d 100644 --- a/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py +++ b/tests/savepoint/translate/translate_pressureadjustedtemperature_nonhydrostatic.py @@ -1,7 +1,8 @@ import numpy as np from typing import Any, Dict -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.dsl.typing import Float from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import temperature_adjust diff --git a/tests/savepoint/translate/translate_qsinit.py b/tests/savepoint/translate/translate_qsinit.py index e1228ced..6b58393e 100644 --- a/tests/savepoint/translate/translate_qsinit.py +++ b/tests/savepoint/translate/translate_qsinit.py @@ -2,7 +2,8 @@ import ndsl.dsl.gt4py_utils as utils import pyFV3.stencils.saturation_adjustment as satadjust -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_ray_fast.py b/tests/savepoint/translate/translate_ray_fast.py index 32de3d55..70a1abf9 100644 --- a/tests/savepoint/translate/translate_ray_fast.py +++ b/tests/savepoint/translate/translate_ray_fast.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3.stencils import RayleighDamping from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_remapping.py b/tests/savepoint/translate/translate_remapping.py index 79c66b99..f0dbdfde 100644 --- a/tests/savepoint/translate/translate_remapping.py +++ b/tests/savepoint/translate/translate_remapping.py @@ -1,5 +1,6 @@ import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, QuantityFactory, StencilFactory +from ndsl import QuantityFactory, StencilFactory +from f90nml.namelist import Namelist from ndsl.constants import Z_DIM from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import LagrangianToEulerian diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 6aa6916f..982ba391 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,6 +1,7 @@ from types import SimpleNamespace -from ndsl import Namelist, StencilFactory, Quantity +from ndsl import StencilFactory, Quantity +from f90nml import Namelist from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, diff --git a/tests/savepoint/translate/translate_riem_solver3.py b/tests/savepoint/translate/translate_riem_solver3.py index 76c4a5c3..962ef56f 100644 --- a/tests/savepoint/translate/translate_riem_solver3.py +++ b/tests/savepoint/translate/translate_riem_solver3.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3 import _config as spec from pyFV3.stencils import NonhydrostaticVerticalSolver from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_riem_solver_c.py b/tests/savepoint/translate/translate_riem_solver_c.py index 9b5f0d11..15d29b75 100644 --- a/tests/savepoint/translate/translate_riem_solver_c.py +++ b/tests/savepoint/translate/translate_riem_solver_c.py @@ -1,5 +1,6 @@ import numpy as np -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from pyFV3.stencils import NonhydrostaticVerticalSolverCGrid from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_satadjust3d.py b/tests/savepoint/translate/translate_satadjust3d.py index 99f33b55..f335e519 100644 --- a/tests/savepoint/translate/translate_satadjust3d.py +++ b/tests/savepoint/translate/translate_satadjust3d.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import SatAdjust3d from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_scalar_profile.py b/tests/savepoint/translate/translate_scalar_profile.py index 12c08355..9a9c8150 100644 --- a/tests/savepoint/translate/translate_scalar_profile.py +++ b/tests/savepoint/translate/translate_scalar_profile.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_tracer2d1l.py b/tests/savepoint/translate/translate_tracer2d1l.py index ecc658d3..5be216b9 100644 --- a/tests/savepoint/translate/translate_tracer2d1l.py +++ b/tests/savepoint/translate/translate_tracer2d1l.py @@ -1,11 +1,13 @@ import pytest -from ndsl import Namelist, QuantityFactory, StencilFactory +from ndsl import QuantityFactory, StencilFactory +from f90nml.namelist import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import ParallelTranslate from pyFV3.stencils import FiniteVolumeTransport, TracerAdvection from pyFV3.tracers import TracersType, setup_tracers from pyFV3.utils.functional_validation import get_subset_func +from pyFV3 import DynamicalCoreConfig class TranslateTracer2D1L(ParallelTranslate): @@ -38,12 +40,12 @@ def __init__( sizer=stencil_factory.grid_indexing._sizer, backend=stencil_factory.backend, ) - self.namelist = namelist self._subset = get_subset_func( self.grid.grid_indexing, dims=[X_DIM, Y_DIM, Z_DIM], n_halo=((0, 0), (0, 0)), ) + self.config = DynamicalCoreConfig.from_f90nml(namelist) def compute_parallel(self, inputs, communicator): self._base.make_storage_data_input_vars(inputs, dict_4d=False) @@ -60,7 +62,7 @@ def compute_parallel(self, inputs, communicator): grid_data=self.grid.grid_data, damping_coefficients=self.grid.damping_coefficients, grid_type=self.grid.grid_type, - hord=self.namelist.hord_tr, + hord=self.config.hord_tr, ) self.tracer_advection = TracerAdvection( @@ -87,7 +89,7 @@ def compute_parallel(self, inputs, communicator): # outputs["tracers"] = tracers.quantity.field[:] return outputs - def compute_sequential(self, a, b): + def compute_sequential(self, inputs_list, communicator_list): pytest.skip( f"{self.__class__} only has a mpirun implementation, " "not running in mock-parallel" diff --git a/tests/savepoint/translate/translate_tracer2d1l_cmax.py b/tests/savepoint/translate/translate_tracer2d1l_cmax.py index 3eff9416..0cde88e1 100644 --- a/tests/savepoint/translate/translate_tracer2d1l_cmax.py +++ b/tests/savepoint/translate/translate_tracer2d1l_cmax.py @@ -1,4 +1,5 @@ -from ndsl import Namelist, StencilFactory, QuantityFactory, Quantity +from ndsl import StencilFactory, QuantityFactory, Quantity +from f90nml import Namelist from ndsl.stencils.testing import ParallelTranslate2Py from pyFV3.stencils.tracer_2d_1l import TracerCMax from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM diff --git a/tests/savepoint/translate/translate_updatedzc.py b/tests/savepoint/translate/translate_updatedzc.py index ed8d82bc..ded74351 100644 --- a/tests/savepoint/translate/translate_updatedzc.py +++ b/tests/savepoint/translate/translate_updatedzc.py @@ -1,6 +1,7 @@ import numpy as np -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from pyFV3.stencils import UpdateGeopotentialHeightOnCGrid from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_updatedzd.py b/tests/savepoint/translate/translate_updatedzd.py index 358e4dd0..1921914a 100644 --- a/tests/savepoint/translate/translate_updatedzd.py +++ b/tests/savepoint/translate/translate_updatedzd.py @@ -1,6 +1,7 @@ import numpy as np -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from pyFV3 import DynamicalCoreConfig from pyFV3.stencils import UpdateHeightOnDGrid, d_sw diff --git a/tests/savepoint/translate/translate_xppm.py b/tests/savepoint/translate/translate_xppm.py index 2789c747..b33325f6 100644 --- a/tests/savepoint/translate/translate_xppm.py +++ b/tests/savepoint/translate/translate_xppm.py @@ -1,5 +1,6 @@ import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.stencils.testing import TranslateGrid from pyFV3.stencils import XPiecewiseParabolic from pyFV3.testing import TranslateDycoreFortranData2Py diff --git a/tests/savepoint/translate/translate_xtp_u.py b/tests/savepoint/translate/translate_xtp_u.py index 4a0cbbb6..60f64957 100644 --- a/tests/savepoint/translate/translate_xtp_u.py +++ b/tests/savepoint/translate/translate_xtp_u.py @@ -1,6 +1,7 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.dsl.typing import FloatField, FloatFieldIJ from ndsl.grid import GridData from pyFV3.stencils import xtp_u diff --git a/tests/savepoint/translate/translate_yppm.py b/tests/savepoint/translate/translate_yppm.py index 4b9430d0..6d296da2 100644 --- a/tests/savepoint/translate/translate_yppm.py +++ b/tests/savepoint/translate/translate_yppm.py @@ -1,5 +1,6 @@ import ndsl.dsl.gt4py_utils as utils -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.dsl.typing import Float from ndsl.stencils.testing import TranslateGrid from pyFV3.stencils import YPiecewiseParabolic diff --git a/tests/savepoint/translate/translate_ytp_v.py b/tests/savepoint/translate/translate_ytp_v.py index 038f889f..cfb2cc57 100644 --- a/tests/savepoint/translate/translate_ytp_v.py +++ b/tests/savepoint/translate/translate_ytp_v.py @@ -1,6 +1,7 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval -from ndsl import Namelist, StencilFactory +from ndsl import StencilFactory +from f90nml import Namelist from ndsl.dsl.typing import FloatField, FloatFieldIJ from ndsl.grid import GridData from pyFV3 import DynamicalCoreConfig From ddb37d98f88f2e99034177895a625376d1a91d4e Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 11 Mar 2026 17:28:48 -0400 Subject: [PATCH 245/252] [WIP] Flak8 fixes --- pyfv3/dycore_state.py | 10 +- .../test_cases/initialize_baroclinic.py | 8 +- .../test_cases/initialize_tc.py | 6 +- pyfv3/stencils/a2b_ord4.py | 8 +- pyfv3/stencils/c_sw.py | 2 +- pyfv3/stencils/compute_total_energy.py | 67 +++++++------- pyfv3/stencils/d2a2c_vect.py | 2 +- pyfv3/stencils/d_sw.py | 4 +- pyfv3/stencils/divergence_damping.py | 4 +- pyfv3/stencils/dyn_core.py | 3 +- pyfv3/stencils/fillz.py | 8 +- pyfv3/stencils/fv_dynamics.py | 92 +++++++++++-------- pyfv3/stencils/map_single.py | 45 ++++----- pyfv3/stencils/mapn_tracer.py | 11 ++- pyfv3/stencils/pk3_halo.py | 2 +- pyfv3/stencils/ray_fast.py | 10 +- pyfv3/stencils/remapping.py | 8 +- pyfv3/stencils/remapping_GEOS.py | 56 +++++------ pyfv3/stencils/tracer_2d_1l.py | 17 ++-- pyfv3/testing/translate_fvdynamics.py | 2 +- pyfv3/tracers.py | 5 +- .../translate/translate_MapN_Tracer_2d.py | 7 +- .../translate/translate_Pressures_mapU.py | 7 +- .../translate/translate_Pressures_mapV.py | 7 +- .../savepoint/translate/translate_a2b_ord4.py | 1 + .../translate/translate_cs_profile.py | 5 +- tests/savepoint/translate/translate_fillz.py | 9 +- .../translate/translate_getMPIprop.py | 2 +- ...ranslate_lagrangian_contribution_interp.py | 5 +- .../translate/translate_map1_ppm_W.py | 5 +- .../translate/translate_map1_ppm_delz.py | 4 +- .../translate/translate_map_scalar.py | 5 +- .../translate/translate_mpp_global_sum.py | 5 +- .../translate/translate_pe_pk_delp_peln.py | 5 +- .../translate/translate_pn2_pk_delp.py | 3 +- .../translate/translate_remapping.py | 44 +++++---- .../translate/translate_remapping_GEOS.py | 9 +- .../translate/translate_riem_solver_c.py | 5 +- .../translate/translate_scalar_profile.py | 5 +- .../translate/translate_tracer2d1l.py | 10 +- .../translate/translate_tracer2d1l_cmax.py | 5 +- .../translate_w_fix_consrv_moment.py | 3 +- 42 files changed, 282 insertions(+), 239 deletions(-) diff --git a/pyfv3/dycore_state.py b/pyfv3/dycore_state.py index b24c2ce5..0f6a771a 100644 --- a/pyfv3/dycore_state.py +++ b/pyfv3/dycore_state.py @@ -327,7 +327,7 @@ def init_zeros( dtype_dict[_field.name] if dtype_dict and _field.name in dtype_dict.keys() else Float - ), # type: ignore + ), allow_mismatch_float_precision=allow_mismatch_float_precision, ).data elif _field.name == "tracers": @@ -337,7 +337,7 @@ def init_zeros( initial_storages[_field.name] = qty_factory_tracers.zeros( [I_DIM, J_DIM, K_DIM, "tracers"], _field.metadata["units"], - dtype=Float, # type: ignore + dtype=Float, ).data return cls.init_from_storages( storages=initial_storages, @@ -372,7 +372,7 @@ def init_from_numpy_arrays( extent=quantity_factory.sizer.get_extent(dims), backend=backend, ) - elif issubclass(_field.type, Tracers): + elif issubclass(_field.type, TracersType): if len(dict_of_numpy_arrays[_field.name]) != len(tracer_list): raise ValueError( "[pyFV3] DycoreState init:" @@ -384,7 +384,7 @@ def init_from_numpy_arrays( quantity_factory=quantity_factory, tracer_mapping=tracer_list, ) - state = cls(**dict_state) # type: ignore + state = cls(**dict_state) return state @classmethod @@ -493,7 +493,7 @@ def from_fortran_restart( return new - def _xr_dataarray_from_quantity(self, name: str, metadata: Dict[str, Any], data): + def _xr_dataarray_from_quantity(self, name: str, metadata: dict[str, Any], data): dims = [f"{dim_name}_{name}" for dim_name in metadata["dims"]] return xr.DataArray( gt_utils.asarray(data), diff --git a/pyfv3/initialization/test_cases/initialize_baroclinic.py b/pyfv3/initialization/test_cases/initialize_baroclinic.py index 8dfc2e90..ef3203de 100644 --- a/pyfv3/initialization/test_cases/initialize_baroclinic.py +++ b/pyfv3/initialization/test_cases/initialize_baroclinic.py @@ -35,7 +35,9 @@ def apply_perturbation(u_component, up, lon, lat, is_steady: bool = False): R = constants.RADIUS / Float(10.0) # Perturbation radius for test case 13 r = np.zeros((u_component.shape[0], u_component.shape[1], 1)) # Equation (11), distance from perturbation at 20E, 40N in JRMS2006 - r = great_circle_distance_lon_lat(PCEN[0], lon, PCEN[1], lat, constants.RADIUS, np)[:, :, None] + r = great_circle_distance_lon_lat(PCEN[0], lon, PCEN[1], lat, constants.RADIUS, np)[ + :, :, None + ] r3d = np.repeat(r, u_component.shape[2], axis=2) near_perturbation = (r3d / R) ** 2.0 < 40.0 # Equation(10) in JRMS2006 perturbation applied to u_component @@ -264,7 +266,9 @@ def baroclinic_initialization( qvapor[:nx, :ny, :-1] = init_utils.specific_humidity( delp[slice_3d], peln[slice_3d], lat_agrid[slice_2d] ) - pt[slice_3d] = init_utils.moisture_adjusted_temperature(pt[slice_3d], qvapor[slice_3d]) + pt[slice_3d] = init_utils.moisture_adjusted_temperature( + pt[slice_3d], qvapor[slice_3d] + ) def init_baroclinic_state( diff --git a/pyfv3/initialization/test_cases/initialize_tc.py b/pyfv3/initialization/test_cases/initialize_tc.py index 71ef7ea4..9576ce1b 100644 --- a/pyfv3/initialization/test_cases/initialize_tc.py +++ b/pyfv3/initialization/test_cases/initialize_tc.py @@ -568,14 +568,10 @@ def init_tc_state( numpy_state.w[:] = w state = DycoreState.init_from_numpy_arrays( numpy_state.__dict__, -<<<<<<< HEAD:pyFV3/initialization/test_cases/initialize_tc.py quantity_factory=quantity_factory, - backend=sample_quantity.metadata.gt4py_backend, + backend=sample_quantity.metadata.backend, tracer_list=["vapor", "liquid", "rain", "snow", "ice", "graupel", "cloud"], -======= sizer=quantity_factory.sizer, - backend=sample_quantity.metadata.backend, ->>>>>>> NOAA/develop:pyfv3/initialization/test_cases/initialize_tc.py ) state.tracers["vapor"].view[:] = qvapor diff --git a/pyfv3/stencils/a2b_ord4.py b/pyfv3/stencils/a2b_ord4.py index 37eaf7f7..71c7c80b 100644 --- a/pyfv3/stencils/a2b_ord4.py +++ b/pyfv3/stencils/a2b_ord4.py @@ -3,7 +3,13 @@ from ndsl.dsl.gt4py import PARALLEL, asin, computation, cos from ndsl.dsl.gt4py import function as gtfunction from ndsl.dsl.gt4py import horizontal, interval, region, sin, sqrt -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.typing import ( + Float, + FloatField, + FloatFieldI64, + FloatFieldIJ, + FloatFieldIJ64, +) from ndsl.grid import GridData from ndsl.stencils.basic_operations import copy diff --git a/pyfv3/stencils/c_sw.py b/pyfv3/stencils/c_sw.py index 9d0f8fd6..954e8406 100644 --- a/pyfv3/stencils/c_sw.py +++ b/pyfv3/stencils/c_sw.py @@ -1,7 +1,7 @@ from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate from ndsl.constants import I_DIM, I_INTERFACE_DIM, J_DIM, J_INTERFACE_DIM, K_DIM from ndsl.dsl.gt4py import PARALLEL, computation, horizontal, interval, region # noqa -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, I, J from ndsl.grid import GridData from ndsl.stencils import corners from pyfv3.stencils.d2a2c_vect import DGrid2AGrid2CGridVectors diff --git a/pyfv3/stencils/compute_total_energy.py b/pyfv3/stencils/compute_total_energy.py index 95455e10..4ca3fc13 100644 --- a/pyfv3/stencils/compute_total_energy.py +++ b/pyfv3/stencils/compute_total_energy.py @@ -1,33 +1,32 @@ from gt4py.cartesian.gtscript import BACKWARD, FORWARD, K, computation, interval -from pyFV3._config import DynamicalCoreConfig -from pyFV3.stencils.moist_cv import moist_cv_nwat6_fn -from pyFV3.tracers import TracersType from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import GRAV, X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from ndsl.grid import GridData +from pyfv3._config import DynamicalCoreConfig +from pyfv3.stencils.moist_cv import moist_cv_nwat6_fn def _compute_total_energy__stencil( - hs: FloatFieldIJ, # type: ignore - delp: FloatField, # type: ignore - delz: FloatField, # type: ignore - qc: FloatField, # type:ignore - pt: FloatField, # type: ignore - u: FloatField, # type: ignore - v: FloatField, # type: ignore - w: FloatField, # type: ignore - qvapor: FloatField, # type: ignore - qliquid: FloatField, # type: ignore - qrain: FloatField, # type: ignore - qsnow: FloatField, # type: ignore - qice: FloatField, # type: ignore - qgraupel: FloatField, # type: ignore - rsin2: FloatFieldIJ, # type: ignore - cosa_s: FloatFieldIJ, # type: ignore - phyz: FloatField, # type: ignore - te_2d: FloatFieldIJ, # type: ignore + hs: FloatFieldIJ, + delp: FloatField, + delz: FloatField, + qc: FloatField, + pt: FloatField, + u: FloatField, + v: FloatField, + w: FloatField, + qvapor: FloatField, + qliquid: FloatField, + qrain: FloatField, + qsnow: FloatField, + qice: FloatField, + qgraupel: FloatField, + rsin2: FloatFieldIJ, + cosa_s: FloatFieldIJ, + phyz: FloatField, + te_2d: FloatFieldIJ, ): """ Dev Note: this is _very_ close to moist_cv.moist_te. The only numerical differences @@ -102,14 +101,12 @@ def __init__( ) if config.hydrostatic: raise NotImplementedError( - "Dynamics (Compute Total Energy): " - " hydrostatic option is not implemented." + "Dynamics (Compute Total Energy): hydrostatic option is not implemented." ) if not config.moist_phys: raise NotImplementedError( - "Dynamics (Compute Total Energy): " - " moist_phys=False option is not implemented." + "Dynamics (Compute Total Energy): moist_phys=False option is not implemented." ) self._phyz = quantity_factory.zeros( @@ -127,16 +124,16 @@ def __init__( def __call__( self, - hs: FloatFieldIJ, # type: ignore - delp: FloatField, # type: ignore - delz: FloatField, # type: ignore - qc: FloatField, # type:ignore - pt: FloatField, # type: ignore - u: FloatField, # type: ignore - v: FloatField, # type: ignore - w: FloatField, # type: ignore - tracers: TracersType, - te_2d: FloatFieldIJ, # type: ignore + hs: FloatFieldIJ, + delp: FloatField, + delz: FloatField, + qc: FloatField, + pt: FloatField, + u: FloatField, + v: FloatField, + w: FloatField, + tracers, + te_2d: FloatFieldIJ, ) -> None: self._compute_total_energy( hs=hs, diff --git a/pyfv3/stencils/d2a2c_vect.py b/pyfv3/stencils/d2a2c_vect.py index 0282b74c..70711eb3 100644 --- a/pyfv3/stencils/d2a2c_vect.py +++ b/pyfv3/stencils/d2a2c_vect.py @@ -3,7 +3,7 @@ from ndsl.dsl.gt4py import PARALLEL, computation from ndsl.dsl.gt4py import function as gtfunction from ndsl.dsl.gt4py import horizontal, interval, region -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, get_precision from ndsl.grid import GridData from ndsl.stencils import corners from pyfv3.stencils.a2b_ord4 import a1, a2, lagrange_x_func, lagrange_y_func diff --git a/pyfv3/stencils/d_sw.py b/pyfv3/stencils/d_sw.py index 6d6e053e..dd185b27 100644 --- a/pyfv3/stencils/d_sw.py +++ b/pyfv3/stencils/d_sw.py @@ -5,7 +5,7 @@ from ndsl.dsl.gt4py import PARALLEL, I, J, computation from ndsl.dsl.gt4py import function as gtfunction from ndsl.dsl.gt4py import horizontal, interval, region -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK +from ndsl.dsl.typing import Float, FloatField, FloatField64, FloatFieldIJ, FloatFieldK from ndsl.grid import DampingCoefficients, GridData from pyfv3._config import DGridShallowWaterLagrangianDynamicsConfig from pyfv3.stencils import delnflux @@ -1025,7 +1025,7 @@ def make_quantity(): ) self._accumulate_delp = stencil_factory.from_dims_halo( func=delp_increment_accumulation, - compute_dims=[X_DIM, Y_DIM, Z_DIM], + compute_dims=[I_DIM, J_DIM, K_DIM], ) def __call__( diff --git a/pyfv3/stencils/divergence_damping.py b/pyfv3/stencils/divergence_damping.py index eb356165..b1b1a0b4 100644 --- a/pyfv3/stencils/divergence_damping.py +++ b/pyfv3/stencils/divergence_damping.py @@ -5,7 +5,7 @@ from ndsl import Quantity, QuantityFactory, StencilFactory from ndsl.constants import I_DIM, I_INTERFACE_DIM, J_DIM, J_INTERFACE_DIM, K_DIM from ndsl.dsl.dace.orchestration import dace_inhibitor, orchestrate -from ndsl.dsl.gt4py import PARALLEL, computation +from ndsl.dsl.gt4py import PARALLEL, computation, float32 from ndsl.dsl.gt4py import function as gtfunction from ndsl.dsl.gt4py import horizontal, interval, region, sqrt from ndsl.dsl.stencil import get_stencils_with_varied_bounds @@ -24,7 +24,7 @@ def damp_tmp(q, da_min_c, d2_bg, dddmp): return damp -@gtscript.function +@gtfunction def damp_tmp2(q, da_min_c, d2_bg, dddmp): damp: float32 = da_min_c * max(d2_bg, min(0.2, dddmp * q)) return damp diff --git a/pyfv3/stencils/dyn_core.py b/pyfv3/stencils/dyn_core.py index 012fe8b5..e0dd9ec6 100644 --- a/pyfv3/stencils/dyn_core.py +++ b/pyfv3/stencils/dyn_core.py @@ -28,7 +28,6 @@ K_DIM, K_INTERFACE_DIM, ) -from ndsl.dsl.dace.orchestration import dace_inhibitor from ndsl.dsl.gt4py import ( BACKWARD, FORWARD, @@ -38,7 +37,7 @@ interval, region, ) -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.typing import Float, FloatField, FloatField64, FloatFieldIJ from ndsl.grid import DampingCoefficients, GridData from ndsl.typing import Communicator from pyfv3._config import AcousticDynamicsConfig diff --git a/pyfv3/stencils/fillz.py b/pyfv3/stencils/fillz.py index f42e1885..0057f509 100644 --- a/pyfv3/stencils/fillz.py +++ b/pyfv3/stencils/fillz.py @@ -1,7 +1,7 @@ -import typing +import dace +from typing_extensions import no_type_check -import ndsl.dsl.gt4py_utils as utils -from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, computation, interval, max, min from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, Int, IntFieldIJ @@ -136,7 +136,7 @@ def __init__( def __call__( self, dp2: FloatField, - tracers: TracersType, + tracers, ): """ Args: diff --git a/pyfv3/stencils/fv_dynamics.py b/pyfv3/stencils/fv_dynamics.py index abd93a7f..ab45d72d 100644 --- a/pyfv3/stencils/fv_dynamics.py +++ b/pyfv3/stencils/fv_dynamics.py @@ -3,13 +3,22 @@ from dace.frontend.python.interface import nounroll as dace_no_unroll -import ndsl.dsl.gt4py_utils as utils import pyfv3.stencils.moist_cv as moist_cv from ndsl import Quantity, QuantityFactory, StencilFactory, WrappedHaloUpdater from ndsl.comm.mpi import MPI -from ndsl.constants import I_DIM, J_DIM, K_DIM, K_INTERFACE_DIM, KAPPA, NQ, ZVIR +from ndsl.constants import ( + I_DIM, + I_INTERFACE_DIM, + J_DIM, + J_INTERFACE_DIM, + K_DIM, + K_INTERFACE_DIM, + KAPPA, + NQ, + ZVIR, +) from ndsl.dsl.dace.orchestration import dace_inhibitor, orchestrate -from ndsl.dsl.gt4py import PARALLEL, computation, interval +from ndsl.dsl.gt4py import FORWARD, PARALLEL, computation, interval from ndsl.dsl.typing import ( NDSL_64BIT_FLOAT_TYPE, Float, @@ -21,16 +30,19 @@ from ndsl.grid import DampingCoefficients, GridData from ndsl.logging import ndsl_log from ndsl.performance import Timer -from ndsl.stencils.basic_operations import copy +from ndsl.stencils.basic_operations import copy, set_value from ndsl.stencils.c2l_ord import CubedToLatLon -from ndsl.typing import Checkpointer, Communicator +from ndsl.typing import Communicator from pyfv3._config import DynamicalCoreConfig from pyfv3.dycore_state import DycoreState from pyfv3.stencils import fvtp2d, tracer_2d_1l +from pyfv3.stencils.compute_total_energy import ComputeTotalEnergy from pyfv3.stencils.del2cubed import HyperdiffusionDamping from pyfv3.stencils.dyn_core import AcousticDynamics from pyfv3.stencils.neg_adj3 import AdjustNegativeTracerMixingRatio from pyfv3.stencils.remapping import LagrangianToEulerian +from pyfv3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyfv3.version import IS_GEOS class DryMassRoundOff: @@ -43,19 +55,19 @@ def __init__( hydrostatic: bool, ) -> None: self.psx_2d = quantity_factory.zeros( - dims=[X_DIM, Y_DIM], + dims=[I_DIM, J_DIM], units="unknown", dtype=NDSL_64BIT_FLOAT_TYPE, allow_mismatch_float_precision=True, ) self.dpx = quantity_factory.zeros( - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], units="unknown", dtype=NDSL_64BIT_FLOAT_TYPE, allow_mismatch_float_precision=True, ) self.dpx0_2d = quantity_factory.zeros( - dims=[X_DIM, Y_DIM], + dims=[I_DIM, J_DIM], units="unknown", dtype=NDSL_64BIT_FLOAT_TYPE, allow_mismatch_float_precision=True, @@ -78,7 +90,7 @@ def __init__( ) halo_spec = quantity_factory.get_quantity_halo_spec( - dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], + dims=[I_DIM, J_DIM, K_INTERFACE_DIM], n_halo=stencil_factory.grid_indexing.n_halo, dtype=Float, ) @@ -92,9 +104,9 @@ def __init__( @staticmethod def _reset_stencil( - dpx: FloatField64, # type:ignore - psx_2d: FloatFieldIJ64, # type:ignore - pe: FloatField, # type:ignore + dpx: FloatField64, + psx_2d: FloatFieldIJ64, + pe: FloatField, ): with computation(PARALLEL), interval(...): dpx = 0.0 @@ -103,9 +115,9 @@ def _reset_stencil( @staticmethod def _apply_dpx_to_psx_stencil( - dpx: FloatField64, # type:ignore - dpx0_2d: FloatFieldIJ64, # type:ignore - psx_2d: FloatFieldIJ64, # type:ignore + dpx: FloatField64, + dpx0_2d: FloatFieldIJ64, + psx_2d: FloatFieldIJ64, ): with computation(FORWARD), interval(0, 1): dpx0_2d = dpx @@ -118,42 +130,42 @@ def _apply_dpx_to_psx_stencil( @staticmethod def _apply_psx_to_pe_stencil( - psx_2d: FloatFieldIJ64, # type:ignore - pe: FloatField, # type:ignore + psx_2d: FloatFieldIJ64, + pe: FloatField, ): with computation(FORWARD), interval(-1, None): pe[0, 0, 1] = psx_2d - def reset(self, pe: FloatField): # type:ignore + def reset(self, pe: FloatField): self._reset(dpx=self.dpx, psx_2d=self.psx_2d, pe=pe) - def apply(self, pe: FloatField): # type:ignore + def apply(self, pe: FloatField): self._apply_dpx_to_psx(self.dpx, self.dpx0_2d, self.psx_2d) self._pe_halo_updater.update() self._apply_psx_to_pe(self.psx_2d, pe) def _increment_stencil( - value: FloatField, # type:ignore - increment: FloatField, # type:ignore + value: FloatField, + increment: FloatField, ): with computation(PARALLEL), interval(...): value += increment def _copy_cast_defn( - q_in_64: FloatField64, # type:ignore - q_out: FloatField, # type:ignore + q_in_64: FloatField64, + q_out: FloatField, ): with computation(PARALLEL), interval(...): q_out = q_in_64 def pt_to_potential_density_pt( - pkz: FloatField, # type: ignore - dp_initial: FloatField, # type: ignore - q_con: FloatField, # type: ignore - pt: FloatField, # type: ignore + pkz: FloatField, + dp_initial: FloatField, + q_con: FloatField, + pt: FloatField, ): """ Args: @@ -169,10 +181,10 @@ def pt_to_potential_density_pt( def omega_from_w( - delp: FloatField, # type: ignore - delz: FloatField, # type: ignore - w: FloatField, # type: ignore - omega: FloatField, # type: ignore + delp: FloatField, + delz: FloatField, + w: FloatField, + omega: FloatField, ): """ Args: @@ -464,51 +476,51 @@ def __init__( self._f32_correction = get_precision() == 32 if self._f32_correction: self._mfx_f64 = quantity_factory.zeros( - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + dims=[I_INTERFACE_DIM, J_DIM, K_DIM], units="unknown", dtype=NDSL_64BIT_FLOAT_TYPE, allow_mismatch_float_precision=True, ) self._mfy_f64 = quantity_factory.zeros( - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + dims=[I_DIM, J_INTERFACE_DIM, K_DIM], units="unknown", dtype=NDSL_64BIT_FLOAT_TYPE, allow_mismatch_float_precision=True, ) self._cx_f64 = quantity_factory.zeros( - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + dims=[I_INTERFACE_DIM, J_DIM, K_DIM], units="unknown", dtype=NDSL_64BIT_FLOAT_TYPE, allow_mismatch_float_precision=True, ) self._cy_f64 = quantity_factory.zeros( - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + dims=[I_DIM, J_INTERFACE_DIM, K_DIM], units="unknown", dtype=NDSL_64BIT_FLOAT_TYPE, allow_mismatch_float_precision=True, ) self._mfx_local = quantity_factory.zeros( - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + dims=[I_INTERFACE_DIM, J_DIM, K_DIM], units="unknown", dtype=Float, ) self._mfy_local = quantity_factory.zeros( - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + dims=[I_DIM, J_INTERFACE_DIM, K_DIM], units="unknown", dtype=Float, ) self._cx_local = quantity_factory.zeros( - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + dims=[I_INTERFACE_DIM, J_DIM, K_DIM], units="unknown", dtype=Float, ) self._cy_local = quantity_factory.zeros( - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + dims=[I_DIM, J_INTERFACE_DIM, K_DIM], units="unknown", dtype=Float, ) self._set_value = stencil_factory.from_origin_domain( - func=set_value_defn, + func=set_value, origin=grid_indexing.origin_compute(), domain=grid_indexing.domain_compute(add=(1, 1, 0)), ) diff --git a/pyfv3/stencils/map_single.py b/pyfv3/stencils/map_single.py index bd7bf4f7..e7420a54 100644 --- a/pyfv3/stencils/map_single.py +++ b/pyfv3/stencils/map_single.py @@ -6,10 +6,13 @@ from ndsl.dsl.gt4py import FORWARD, PARALLEL, computation, interval from ndsl.dsl.typing import ( # noqa: F401 Bool, + BoolField, + BoolFieldIJ, Float, FloatField, FloatFieldIJ, Int, + IntField, IntFieldIJ, ) from ndsl.stencils.basic_operations import copy @@ -97,15 +100,15 @@ def __init__(self, stencil_factory: StencilFactory, dims: Sequence[str]) -> None def __call__( self, - q: FloatField, # type: ignore - pe1: FloatField, # type: ignore - pe2: FloatField, # type: ignore - q4_1: FloatField, # type: ignore - q4_2: FloatField, # type: ignore - q4_3: FloatField, # type: ignore - q4_4: FloatField, # type: ignore - dp1: FloatField, # type: ignore - lev: IntFieldIJ, # type: ignore + q: FloatField, + pe1: FloatField, + pe2: FloatField, + q4_1: FloatField, + q4_2: FloatField, + q4_3: FloatField, + q4_4: FloatField, + dp1: FloatField, + lev: IntFieldIJ, ): self._lagrangian_contributions( q, @@ -269,32 +272,32 @@ def __init__( ) self._INDEX_LM1 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="", dtype=Int, ) self._INDEX_LP0 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="", dtype=Int, ) self._km = stencil_factory.grid_indexing.domain[2] self._not_exit_loop = quantity_factory.zeros( - [X_DIM, Y_DIM], units="", dtype=bool + [I_DIM, J_DIM], units="", dtype=bool ) def __call__( self, - q: FloatField, # type: ignore - pe1: FloatField, # type: ignore - pe2: FloatField, # type: ignore - q4_1: FloatField, # type: ignore - q4_2: FloatField, # type: ignore - q4_3: FloatField, # type: ignore - q4_4: FloatField, # type: ignore - dp1: FloatField, # type: ignore - lev: IntFieldIJ, # type: ignore + q: FloatField, + pe1: FloatField, + pe2: FloatField, + q4_1: FloatField, + q4_2: FloatField, + q4_3: FloatField, + q4_4: FloatField, + dp1: FloatField, + lev: IntFieldIJ, ): self._lagrangian_contributions_interp( km=self._km, diff --git a/pyfv3/stencils/mapn_tracer.py b/pyfv3/stencils/mapn_tracer.py index 04b577fb..b40a9d45 100644 --- a/pyfv3/stencils/mapn_tracer.py +++ b/pyfv3/stencils/mapn_tracer.py @@ -1,5 +1,6 @@ -import ndsl.dsl.gt4py_utils as utils -from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate +import dace + +from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.dsl.typing import Float, FloatField from pyfv3.stencils.fillz import FillNegativeTracerValues @@ -17,7 +18,7 @@ def __init__( quantity_factory: QuantityFactory, kord: int, fill: bool, - tracers: TracersType, + tracers, ): orchestrate( obj=self, @@ -61,7 +62,7 @@ def __call__( pe1: FloatField, pe2: FloatField, dp2: FloatField, - tracers: TracersType, + tracers, ): """ Remaps the tracer species onto the Eulerian grid @@ -74,7 +75,7 @@ def __call__( dp2 (in): Difference in pressure between Eulerian levels tracers (inout): tracers to be remapped """ - for i_tracer in nounroll(range(tracers.shape[3])): + for i_tracer in dace.nounroll(range(tracers.shape[3])): if i_tracer != self._index_cloud: self._map_single( tracers.quantity.data[:, :, :, i_tracer], pe1, pe2, self._qs diff --git a/pyfv3/stencils/pk3_halo.py b/pyfv3/stencils/pk3_halo.py index a1915623..f078bb12 100644 --- a/pyfv3/stencils/pk3_halo.py +++ b/pyfv3/stencils/pk3_halo.py @@ -1,6 +1,6 @@ from ndsl import QuantityFactory, StencilFactory from ndsl.constants import I_DIM, J_DIM -from ndsl.dsl.gt4py import FORWARD, computation, horizontal, interval, region +from ndsl.dsl.gt4py import FORWARD, computation, exp, horizontal, interval, log, region from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ diff --git a/pyfv3/stencils/ray_fast.py b/pyfv3/stencils/ray_fast.py index 872f7dfd..a6232156 100644 --- a/pyfv3/stencils/ray_fast.py +++ b/pyfv3/stencils/ray_fast.py @@ -47,10 +47,10 @@ def dm_layer(rf, dp, wind): def ray_fast_damping_increment( - pfull: FloatFieldK, # type:ignore - dt: Float, # type:ignore - ptop: Float, # type:ignore - rf: FloatField, # type:ignore + pfull: FloatFieldK, + dt: Float, + ptop: Float, + rf: FloatField, ): """rf is rayleigh damping increment, fraction of vertical velocity left after doing rayleigh damping (w -> w * rf) @@ -85,7 +85,7 @@ def ray_fast_wind_compute( rf_cutoff_nudge (in): ks (in): """ - from __externals__ import hydrostatic, local_ie, local_je, rf_cutoff, tau + from __externals__ import hydrostatic, local_ie, local_je, rf_cutoff # dm_stencil with computation(FORWARD): diff --git a/pyfv3/stencils/remapping.py b/pyfv3/stencils/remapping.py index 548c6c05..8f6032d5 100644 --- a/pyfv3/stencils/remapping.py +++ b/pyfv3/stencils/remapping.py @@ -1,3 +1,5 @@ +from typing import no_type_check + from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.constants import ( I_DIM, @@ -29,8 +31,6 @@ from pyfv3.stencils.saturation_adjustment import SatAdjust3d -from gt4py.cartesian.gtscript import __INLINED # isort:skip - # TODO: Should this be set here or in global_constants? CONSV_MIN = 0.001 @@ -358,7 +358,7 @@ def __init__( config: RemappingConfig, area_64, pfull, - tracers: TracersType, + tracers, exclude_tracers: list[str], checkpointer: Checkpointer | None = None, ): @@ -591,7 +591,7 @@ def __init__( @no_type_check def __call__( self, - tracers: TracersType, + tracers, pt: FloatField, delp: FloatField, delz: FloatField, diff --git a/pyfv3/stencils/remapping_GEOS.py b/pyfv3/stencils/remapping_GEOS.py index 85621dcd..920da858 100644 --- a/pyfv3/stencils/remapping_GEOS.py +++ b/pyfv3/stencils/remapping_GEOS.py @@ -358,35 +358,35 @@ def __init__( def __call__( self, tracers: TracersType, - pt: FloatField, # type: ignore - delp: FloatField, # type: ignore - delz: FloatField, # type: ignore - peln: FloatField, # type: ignore - u: FloatField, # type: ignore - v: FloatField, # type: ignore - w: FloatField, # type: ignore - mfx: FloatField, # type: ignore - mfy: FloatField, # type: ignore - cx: FloatField, # type: ignore - cy: FloatField, # type: ignore - cappa: FloatField, # type: ignore - q_con: FloatField, # type: ignore - pkz: FloatField, # type: ignore - pk: FloatField, # type: ignore - pe: FloatField, # type: ignore - hs: FloatFieldIJ, # type: ignore - te0_2d: FloatFieldIJ, # type: ignore - ps: FloatFieldIJ, # type: ignore - wsd: FloatFieldIJ, # type: ignore - ak: FloatFieldK, # type: ignore - bk: FloatFieldK, # type: ignore - dp1: FloatField, # type: ignore - ptop: Float, # type: ignore - akap: Float, # type: ignore - zvir: Float, # type: ignore + pt: FloatField, + delp: FloatField, + delz: FloatField, + peln: FloatField, + u: FloatField, + v: FloatField, + w: FloatField, + mfx: FloatField, + mfy: FloatField, + cx: FloatField, + cy: FloatField, + cappa: FloatField, + q_con: FloatField, + pkz: FloatField, + pk: FloatField, + pe: FloatField, + hs: FloatFieldIJ, + te0_2d: FloatFieldIJ, + ps: FloatFieldIJ, + wsd: FloatFieldIJ, + ak: FloatFieldK, + bk: FloatFieldK, + dp1: FloatField, + ptop: Float, + akap: Float, + zvir: Float, last_step: bool, - consv_te: Float, # type: ignore - mdt: Float, # type: ignore + consv_te: Float, + mdt: Float, ): """ Remap the deformed Lagrangian surfaces onto the reference, or "Eulerian", diff --git a/pyfv3/stencils/tracer_2d_1l.py b/pyfv3/stencils/tracer_2d_1l.py index 55723d7b..045ea162 100644 --- a/pyfv3/stencils/tracer_2d_1l.py +++ b/pyfv3/stencils/tracer_2d_1l.py @@ -1,4 +1,6 @@ -import math +from typing import no_type_check + +import dace from ndsl import ( Quantity, @@ -7,6 +9,7 @@ WrappedHaloUpdater, orchestrate, ) +from ndsl.comm.mpi import ReductionOperator from ndsl.constants import ( I_DIM, I_INTERFACE_DIM, @@ -15,13 +18,15 @@ K_DIM, N_HALO_DEFAULT, ) +from ndsl.dsl.dace.orchestration import dace_inhibitor from ndsl.dsl.gt4py import PARALLEL, computation from ndsl.dsl.gt4py import function as gtfunction -from ndsl.dsl.gt4py import horizontal, interval, region -from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ +from ndsl.dsl.gt4py import horizontal, int32, interval, region +from ndsl.dsl.typing import FloatField, FloatFieldIJ, FloatFieldK from ndsl.grid import GridData from ndsl.typing import Communicator from pyfv3.stencils.fvtp2d import FiniteVolumeTransport +from pyfv3.tracers import TracersType @gtfunction @@ -322,7 +327,7 @@ def __init__( def __call__( self, - tracers: TracersType, + tracers: TracersType, # type: ignore dp1, x_mass_flux, y_mass_flux, @@ -507,11 +512,11 @@ def __init__( ), ) self._tmp_cmax = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="unknown", ) self._tmp_cmax_in_K = quantity_factory.zeros( - [Z_DIM], + [K_DIM], units="unknown", ) self.max_over_column = 0 diff --git a/pyfv3/testing/translate_fvdynamics.py b/pyfv3/testing/translate_fvdynamics.py index f410c9e2..27da22b1 100644 --- a/pyfv3/testing/translate_fvdynamics.py +++ b/pyfv3/testing/translate_fvdynamics.py @@ -6,7 +6,7 @@ from f90nml import Namelist import ndsl.dsl.gt4py_utils as utils -from ndsl import Quantity, QuantityFactory, StencilFactory +from ndsl import FieldBundle, Quantity, QuantityFactory, StencilFactory from ndsl.constants import ( I_DIM, I_INTERFACE_DIM, diff --git a/pyfv3/tracers.py b/pyfv3/tracers.py index 09d22adf..6faefe74 100644 --- a/pyfv3/tracers.py +++ b/pyfv3/tracers.py @@ -1,10 +1,9 @@ from typing import TypeAlias -from pyFV3.version import IS_GEOS - from ndsl import QuantityFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType +from pyfv3.version import IS_GEOS # Defauult maopping for common models @@ -29,7 +28,7 @@ } -TracersType: TypeAlias = FieldBundleType.T("Tracers") # type: ignore +TracersType: TypeAlias = FieldBundleType.T("Tracers") # type: ignore # noqa def setup_tracers( diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index a13c5d8e..48826b98 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -1,10 +1,11 @@ -from ndsl import StencilFactory from f90nml import Namelist -from ndsl.stencils.testing import TranslateFortranData2Py -from ndsl.stencils.testing.grid import Grid from pyFV3.stencils.mapn_tracer import MapNTracer from pyFV3.tracers import setup_tracers +from ndsl import StencilFactory +from ndsl.stencils.testing import TranslateFortranData2Py +from ndsl.stencils.testing.grid import Grid + class TranslateMapN_Tracer_2d(TranslateFortranData2Py): def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFactory): diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index e4f6c85b..07350db5 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -1,10 +1,11 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.map_single import MapSingle +from pyFV3.stencils.remapping import pe0_ptop_xmax, pressures_mapu + +from ndsl import StencilFactory from ndsl.constants import X_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.map_single import MapSingle -from pyFV3.stencils.remapping import pe0_ptop_xmax, pressures_mapu class TranslatePressures_mapU(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index 81620694..6427b268 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -1,10 +1,11 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.map_single import MapSingle +from pyFV3.stencils.remapping import pressures_mapv + +from ndsl import StencilFactory from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.map_single import MapSingle -from pyFV3.stencils.remapping import pressures_mapv class TranslatePressures_mapV(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_a2b_ord4.py b/tests/savepoint/translate/translate_a2b_ord4.py index c06b89b0..e4b03bd3 100644 --- a/tests/savepoint/translate/translate_a2b_ord4.py +++ b/tests/savepoint/translate/translate_a2b_ord4.py @@ -7,6 +7,7 @@ from ndsl.constants import I_DIM, J_DIM, K_DIM from pyfv3.stencils import DivergenceDamping from pyfv3.testing import TranslateDycoreFortranData2Py +from pyfv3.utils.functional_validation import get_subset_func class A2B_Ord4Compute: diff --git a/tests/savepoint/translate/translate_cs_profile.py b/tests/savepoint/translate/translate_cs_profile.py index d256c507..12e90ded 100644 --- a/tests/savepoint/translate/translate_cs_profile.py +++ b/tests/savepoint/translate/translate_cs_profile.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.remap_profile import RemapProfile + +from ndsl import StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.remap_profile import RemapProfile class TranslateCS_Profile(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_fillz.py b/tests/savepoint/translate/translate_fillz.py index b3fd4a3e..51a42d34 100644 --- a/tests/savepoint/translate/translate_fillz.py +++ b/tests/savepoint/translate/translate_fillz.py @@ -1,12 +1,11 @@ import numpy as np from f90nml import Namelist -import ndsl.dsl.gt4py_utils as utils from ndsl import StencilFactory from ndsl.stencils.testing import pad_field_in_j -from ndsl.utils import safe_assign_array from pyfv3.stencils import fillz from pyfv3.testing import TranslateDycoreFortranData2Py +from pyfv3.tracers import setup_tracers class TranslateFillz(TranslateDycoreFortranData2Py): @@ -36,7 +35,7 @@ def __init__( self.stencil_factory = stencil_factory self._quantity_factory = grid.quantity_factory - def make_storage_data_input_vars(self, inputs, tracers: FieldBundle): + def make_storage_data_input_vars(self, inputs, tracers): storage_vars = self.storage_vars() info = storage_vars["dp2"] inputs["dp2"] = self.make_storage_data( @@ -63,11 +62,11 @@ def compute(self, inputs): ) ) inputs.pop("nq") - fillz = FillNegativeTracerValues( + fillz_ = fillz.FillNegativeTracerValues( self.stencil_factory, self.grid.quantity_factory, ) - fillz(**inputs) + fillz_(**inputs) ds = self.grid.default_domain_dict() ds.update(self.out_vars["q2tracers"]) out = {"q2tracers": tracers.quantity.field[:, 0, :, :]} diff --git a/tests/savepoint/translate/translate_getMPIprop.py b/tests/savepoint/translate/translate_getMPIprop.py index 561c8b5c..d0b4848d 100644 --- a/tests/savepoint/translate/translate_getMPIprop.py +++ b/tests/savepoint/translate/translate_getMPIprop.py @@ -1,7 +1,7 @@ import numpy as np +from f90nml import Namelist from ndsl import StencilFactory -from f90nml import Namelist from ndsl.quantity import Quantity from ndsl.stencils.testing import ParallelTranslate from ndsl.stencils.testing.grid import Grid diff --git a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py index e65753a1..daf29630 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py @@ -1,10 +1,11 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.map_single import lagrangian_contributions_interp + +from ndsl import StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Bool, BoolFieldIJ, FloatField, Int, IntField, IntFieldIJ from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.map_single import lagrangian_contributions_interp class test_Lagragian_Contribution_Interp: diff --git a/tests/savepoint/translate/translate_map1_ppm_W.py b/tests/savepoint/translate/translate_map1_ppm_W.py index 07ca73da..6c399474 100644 --- a/tests/savepoint/translate/translate_map1_ppm_W.py +++ b/tests/savepoint/translate/translate_map1_ppm_W.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.map_single import MapSingle + +from ndsl import StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.map_single import MapSingle class TranslateMap1_PPM_W(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py index 9384076e..e85d41bc 100644 --- a/tests/savepoint/translate/translate_map1_ppm_delz.py +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -1,12 +1,12 @@ +from f90nml import Namelist from gt4py.cartesian.gtscript import PARALLEL, computation, interval +from pyFV3.stencils.map_single import MapSingle from ndsl import StencilFactory -from f90nml import Namelist from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import FloatField from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.map_single import MapSingle def rescale_delz_1( diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py index 35968971..decfb0af 100644 --- a/tests/savepoint/translate/translate_map_scalar.py +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.map_single import MapSingle + +from ndsl import StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.map_single import MapSingle class TranslateMap_Scalar(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index 397b6d2c..5471d033 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.mpi.mpp_sum import MPPGlobalSum + +from ndsl import StencilFactory from ndsl.stencils.testing import ParallelTranslate from ndsl.stencils.testing.grid import Grid from ndsl.typing import Communicator -from pyFV3.mpi.mpp_sum import MPPGlobalSum class TranslateMpp_global_sum(ParallelTranslate): diff --git a/tests/savepoint/translate/translate_pe_pk_delp_peln.py b/tests/savepoint/translate/translate_pe_pk_delp_peln.py index bbf80271..c7526edb 100644 --- a/tests/savepoint/translate/translate_pe_pk_delp_peln.py +++ b/tests/savepoint/translate/translate_pe_pk_delp_peln.py @@ -1,8 +1,9 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.remapping import pe_pk_delp_peln + +from ndsl import StencilFactory from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.remapping import pe_pk_delp_peln class TranslatePE_pk_delp_peln(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_pn2_pk_delp.py b/tests/savepoint/translate/translate_pn2_pk_delp.py index ba29adb4..4016e91c 100644 --- a/tests/savepoint/translate/translate_pn2_pk_delp.py +++ b/tests/savepoint/translate/translate_pn2_pk_delp.py @@ -1,7 +1,8 @@ +from pyFV3.stencils.remapping import pn2_pk_delp + from ndsl import StencilFactory from ndsl.dsl.typing import Float, FloatField from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils.remapping import pn2_pk_delp class testClass: diff --git a/tests/savepoint/translate/translate_remapping.py b/tests/savepoint/translate/translate_remapping.py index 7f8f9043..63c1b707 100644 --- a/tests/savepoint/translate/translate_remapping.py +++ b/tests/savepoint/translate/translate_remapping.py @@ -1,16 +1,18 @@ from f90nml import Namelist import ndsl.dsl.gt4py_utils as utils -from ndsl import StencilFactory +from ndsl import QuantityFactory, StencilFactory from ndsl.constants import K_DIM +from ndsl.stencils.testing import Grid from pyfv3.stencils import LagrangianToEulerian from pyfv3.testing import TranslateDycoreFortranData2Py +from pyfv3.tracers import setup_tracers class TranslateRemapping(TranslateDycoreFortranData2Py): def __init__( self, - grid, + grid: Grid, namelist: Namelist, stencil_factory: StencilFactory, ): @@ -97,10 +99,11 @@ def __init__( self.near_zero = 3e-18 self.ignore_near_zero_errors = {"q_con": True, "tracers": True} self.stencil_factory = stencil_factory - self._quantity_factory = QuantityFactory( - sizer=stencil_factory.grid_indexing._sizer, + self.quantity_factory = QuantityFactory( + sizer=grid.sizer, backend=stencil_factory.backend, ) + self._are_tracers_setup = False def compute_from_storage(self, inputs): wsd_2d = utils.make_storage_from_shape( @@ -108,21 +111,22 @@ def compute_from_storage(self, inputs): ) wsd_2d[:, :] = inputs["wsd"][:, :, 0] inputs["wsd"] = wsd_2d - tracers = Tracers.make_from_4D_array( - quantity_factory=self._quantity_factory, - tracer_mapping=[ - "vapor", - "liquid", - "rain", - "ice", - "snow", - "graupel", - "qo3mr", - "qsgs_tke", - "cloud", - ], - tracer_data=inputs["tracers"], - ) + + if not self._are_tracers_setup: + self._are_tracers_setup = True + tracers = setup_tracers( + number_of_tracers=inputs["tracers"].shape[3], + quantity_factory=self.quantity_factory, + mappings={ + "vapor": 0, + "liquid": 1, + "rain": 3, + "snow": 4, + "ice": 2, + "graupel": 5, + "cloud": 6, + }, + ) inputs["last_step"] = bool(inputs["last_step"]) pfull = self.grid.quantity_factory.zeros([K_DIM], units="Pa") pfull.data[:] = pfull.np.asarray(inputs.pop("pfull")) @@ -138,5 +142,5 @@ def compute_from_storage(self, inputs): exclude_tracers=["cloud"], ) l_to_e_obj(**inputs) - inputs["tracers"] = tracers.as_4D_array() + inputs["tracers"] = tracers.quantity.data[:-1, :-1, :-1.0:-1] return inputs diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 982ba391..10ff3bdd 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,7 +1,10 @@ from types import SimpleNamespace -from ndsl import StencilFactory, Quantity from f90nml import Namelist +from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyFV3.tracers import TracersType, setup_tracers + +from ndsl import Quantity, StencilFactory from ndsl.constants import ( X_DIM, X_INTERFACE_DIM, @@ -13,8 +16,6 @@ from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyFV3 import DynamicalCoreConfig -from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS -from pyFV3.tracers import TracersType, setup_tracers class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): @@ -393,7 +394,7 @@ def outputs_from_state(self, state: dict): return {} outputs = {} storages = {} - for name, properties in self.outputs.items(): + for name, _properties in self.outputs.items(): if name in ["mfx_R4", "mfy_R4", "cx_R4", "cy_R4"]: storages[name] = state[name[:-3]] elif isinstance(state[name], Quantity): diff --git a/tests/savepoint/translate/translate_riem_solver_c.py b/tests/savepoint/translate/translate_riem_solver_c.py index 20d7ea6b..80ca5b01 100644 --- a/tests/savepoint/translate/translate_riem_solver_c.py +++ b/tests/savepoint/translate/translate_riem_solver_c.py @@ -1,8 +1,11 @@ +import numpy as np from f90nml import Namelist from ndsl import StencilFactory +from ndsl.constants import I_DIM, J_DIM, K_DIM from pyfv3.stencils import NonhydrostaticVerticalSolverCGrid from pyfv3.testing import TranslateDycoreFortranData2Py +from pyfv3.utils.functional_validation import get_subset_func class TranslateRiem_Solver_C(TranslateDycoreFortranData2Py): @@ -35,7 +38,7 @@ def __init__( self.stencil_factory = stencil_factory self._subset = get_subset_func( self.grid.grid_indexing, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], n_halo=((3, 3), (3, 3)), ) diff --git a/tests/savepoint/translate/translate_scalar_profile.py b/tests/savepoint/translate/translate_scalar_profile.py index 9a9c8150..8ef7f25d 100644 --- a/tests/savepoint/translate/translate_scalar_profile.py +++ b/tests/savepoint/translate/translate_scalar_profile.py @@ -1,9 +1,10 @@ -from ndsl import StencilFactory from f90nml import Namelist +from pyFV3.stencils.remap_profile import RemapProfile + +from ndsl import StencilFactory from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.remap_profile import RemapProfile class TranslateScalar_Profile(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_tracer2d1l.py b/tests/savepoint/translate/translate_tracer2d1l.py index a0c5b723..6e368de0 100644 --- a/tests/savepoint/translate/translate_tracer2d1l.py +++ b/tests/savepoint/translate/translate_tracer2d1l.py @@ -1,12 +1,12 @@ import pytest from f90nml import Namelist -import ndsl.dsl.gt4py_utils as utils -from ndsl import StencilFactory +from ndsl import QuantityFactory, StencilFactory from ndsl.constants import I_DIM, J_DIM, K_DIM -from ndsl.stencils.testing import ParallelTranslate +from ndsl.stencils.testing import Grid, ParallelTranslate from pyfv3 import DynamicalCoreConfig from pyfv3.stencils import FiniteVolumeTransport, TracerAdvection +from pyfv3.tracers import setup_tracers from pyfv3.utils.functional_validation import get_subset_func @@ -20,7 +20,7 @@ class TranslateTracer2D1L(ParallelTranslate): def __init__( self, - grid, + grid: Grid, namelist: Namelist, stencil_factory: StencilFactory, ): @@ -37,7 +37,7 @@ def __init__( self._base.out_vars = self._base.in_vars["data_vars"] self.stencil_factory = stencil_factory self._quantity_factory = QuantityFactory( - sizer=stencil_factory.grid_indexing._sizer, + sizer=grid.sizer, backend=stencil_factory.backend, ) self._subset = get_subset_func( diff --git a/tests/savepoint/translate/translate_tracer2d1l_cmax.py b/tests/savepoint/translate/translate_tracer2d1l_cmax.py index 0cde88e1..a67ddcac 100644 --- a/tests/savepoint/translate/translate_tracer2d1l_cmax.py +++ b/tests/savepoint/translate/translate_tracer2d1l_cmax.py @@ -1,8 +1,9 @@ -from ndsl import StencilFactory, QuantityFactory, Quantity from f90nml import Namelist -from ndsl.stencils.testing import ParallelTranslate2Py from pyFV3.stencils.tracer_2d_1l import TracerCMax + +from ndsl import Quantity, QuantityFactory, StencilFactory from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.stencils.testing import ParallelTranslate2Py def _quantity_wrap(storage, dims, grid_indexing): diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py index 6a642bc2..bdaee31e 100644 --- a/tests/savepoint/translate/translate_w_fix_consrv_moment.py +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -1,7 +1,8 @@ +from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment + from ndsl.dsl.typing import Float from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid -from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment class TranslateW_fix_consrv_moment(TranslateFortranData2Py): From 303e24ce92bce83f33983f61efe6f592c1af43b2 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Mar 2026 11:20:37 -0400 Subject: [PATCH 246/252] [WIP] Non-tracers lint & fixes --- pyfv3/_config.py | 2 ++ pyfv3/dycore_state.py | 12 +++++--- .../test_cases/initialize_baroclinic.py | 2 +- pyfv3/mpi/mpp_sum.py | 30 ++++++++++++++----- pyfv3/stencils/fxadv.py | 5 ---- pyfv3/stencils/map_single.py | 13 ++++---- pyfv3/stencils/remapping.py | 1 - pyfv3/stencils/tracer_2d_1l.py | 2 +- pyfv3/testing/translate_fvdynamics.py | 4 +-- pyfv3/utils/functional_validation.py | 24 --------------- pyfv3/wrappers/geos_wrapper.py | 5 ++-- 11 files changed, 45 insertions(+), 55 deletions(-) diff --git a/pyfv3/_config.py b/pyfv3/_config.py index 0f226979..5e0ec525 100644 --- a/pyfv3/_config.py +++ b/pyfv3/_config.py @@ -121,6 +121,8 @@ class AcousticDynamicsConfig: """ riemann: RiemannConfig d_grid_shallow_water: DGridShallowWaterLagrangianDynamicsConfig + dz_min: float = 2.0 + """Controls minimum thickness in NH solver""" @property def nord(self) -> int: diff --git a/pyfv3/dycore_state.py b/pyfv3/dycore_state.py index 0f6a771a..1557b717 100644 --- a/pyfv3/dycore_state.py +++ b/pyfv3/dycore_state.py @@ -1,7 +1,9 @@ from collections.abc import Mapping from dataclasses import asdict, dataclass, field, fields +from types import MappingProxyType from typing import Any, Self +import numpy.typing as npt import xarray as xr import ndsl.dsl.gt4py_utils as gt_utils @@ -315,7 +317,7 @@ def init_zeros( tracer_count: int, backend: Backend, dtype_dict: dict[str, type] | None = None, - allow_mismatch_float_precision=False, + allow_mismatch_float_precision: bool = False, ) -> Self: initial_storages = {} for _field in fields(cls): @@ -493,7 +495,9 @@ def from_fortran_restart( return new - def _xr_dataarray_from_quantity(self, name: str, metadata: dict[str, Any], data): + def _xr_dataarray_from_array( + self, name: str, metadata: MappingProxyType[Any, Any], data: npt.ArrayLike + ): dims = [f"{dim_name}_{name}" for dim_name in metadata["dims"]] return xr.DataArray( gt_utils.asarray(data), @@ -509,13 +513,13 @@ def xr_dataset(self) -> xr.Dataset: data_vars = {} for name, field_info in self.__dataclass_fields__.items(): if issubclass(field_info.type, Quantity): - data_vars[name] = self._xr_dataarray_from_quantity( + data_vars[name] = self._xr_dataarray_from_array( name=name, metadata=field_info.metadata, data=getattr(self, name).data, ) if isinstance(field_info.type, FieldBundle): - data_vars[name] = self._xr_dataarray_from_quantity( + data_vars[name] = self._xr_dataarray_from_array( name=name, metadata=field_info.metadata, data=getattr(self, name).quantity.data, diff --git a/pyfv3/initialization/test_cases/initialize_baroclinic.py b/pyfv3/initialization/test_cases/initialize_baroclinic.py index ef3203de..90e69ad1 100644 --- a/pyfv3/initialization/test_cases/initialize_baroclinic.py +++ b/pyfv3/initialization/test_cases/initialize_baroclinic.py @@ -369,7 +369,7 @@ def init_baroclinic_state( state = DycoreState.init_from_numpy_arrays( numpy_state.__dict__, quantity_factory=quantity_factory, - backend=sample_quantity.metadata.backend, + backend=sample_quantity.backend, tracer_list=["vapor", "liquid", "rain", "snow", "ice", "graupel", "cloud"], ) state.tracers["vapor"].view[:] = numpy_state.qvapor[slice_3d] diff --git a/pyfv3/mpi/mpp_sum.py b/pyfv3/mpi/mpp_sum.py index ab84bd26..a37cae03 100644 --- a/pyfv3/mpi/mpp_sum.py +++ b/pyfv3/mpi/mpp_sum.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np from ndsl import Quantity, StencilFactory @@ -5,7 +7,13 @@ from ndsl.dsl.typing import Float -def _increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): +def _increment_ints_faster( + int_sum: np.ndarray, + pr: list[float], + I_pr: list[float], + r: float, + max_mag_term: float, +) -> None: if (r >= 1e30) == r < 1e30: print("NaN_error") return @@ -23,7 +31,13 @@ def _increment_ints_faster(int_sum, pr, I_pr, r, max_mag_term): int_sum[i] = int_sum[i] + sgn * ival -def _carry_overflow(int_sum, prec, I_prec, prec_error): +def _carry_overflow( + int_sum: np.ndarray, + prec: int, + I_prec: float, + prec_error: float, +) -> bool: + overflow_error = False for i in range(len(int_sum) - 1, 0, -1): if abs(int_sum[i]) > prec: num_carry = int(int_sum[i] * I_prec) @@ -31,9 +45,10 @@ def _carry_overflow(int_sum, prec, I_prec, prec_error): int_sum[i - 1] = int_sum[i - 1] + num_carry if abs(int_sum[0]) > prec_error: overflow_error = True + return overflow_error -def _regularize_ints(int_sum, prec, I_prec): +def _regularize_ints(int_sum: np.ndarray, prec: int, I_prec: float) -> None: for i in range(len(int_sum) - 1, 0, -1): if abs(int_sum[i]) > prec: num_carry = int(int_sum[i] * I_prec) @@ -61,7 +76,7 @@ def _regularize_ints(int_sum, prec, I_prec): int_sum[i - 1] = int_sum[i - 1] + 1 -def _ints_to_real(ints, pr): +def _ints_to_real(ints: np.ndarray, pr: list[float]) -> float: r = 0.0 for i in range(len(ints)): @@ -80,14 +95,14 @@ def __init__( data=np.zeros((NUMINT), dtype=Float), dims=["K"], units="dunno", - gt4py_backend=stencil_factory.backend, + backend=stencil_factory.backend, ) self._ints_sum_reduce = Quantity( data=np.zeros((NUMINT), dtype=Float), dims=["K"], units="dunno", - gt4py_backend=stencil_factory.backend, + backend=stencil_factory.backend, ) def __call__(self, qty_to_sum: Quantity) -> Float: @@ -115,7 +130,8 @@ def __call__(self, qty_to_sum: Quantity) -> Float: self._ints_sum.data, pr, I_pr, qty_to_sum.field[i, j], mag_max_term ) - _carry_overflow(self._ints_sum.data, prec, I_prec, prec_error) + if not _carry_overflow(self._ints_sum.data, prec, I_prec, prec_error): + warnings.warn("Overflow in MPP sum", category=UserWarning, stacklevel=2) self._comm.all_reduce( self._ints_sum, diff --git a/pyfv3/stencils/fxadv.py b/pyfv3/stencils/fxadv.py index 5fff9df3..4341af4e 100644 --- a/pyfv3/stencils/fxadv.py +++ b/pyfv3/stencils/fxadv.py @@ -610,11 +610,6 @@ def __init__( origin=origin, domain=domain, ) - # self._set_nans = get_set_nan_func( - # grid_indexing, - # dims=[I_DIM, J_DIM, K_DIM], - # n_halo=((2, 2), (2, 2)), - # ) def __call__( self, diff --git a/pyfv3/stencils/map_single.py b/pyfv3/stencils/map_single.py index e7420a54..a715ae23 100644 --- a/pyfv3/stencils/map_single.py +++ b/pyfv3/stencils/map_single.py @@ -372,14 +372,11 @@ def make_quantity(): dims=dims, ) - if interpolate_contribution: - self._lagrangian_contributions = LagrangianContributionInterpolated( - stencil_factory, quantity_factory, dims - ) - else: - self._lagrangian_contributions = LagrangianContribution( - stencil_factory, dims - ) + self._lagrangian_contributions = ( + LagrangianContributionInterpolated(stencil_factory, quantity_factory, dims) + if interpolate_contribution + else LagrangianContribution(stencil_factory, dims) + ) def __call__( self, diff --git a/pyfv3/stencils/remapping.py b/pyfv3/stencils/remapping.py index 8f6032d5..c44fb3b6 100644 --- a/pyfv3/stencils/remapping.py +++ b/pyfv3/stencils/remapping.py @@ -484,7 +484,6 @@ def __init__( abs(config.kord_tr), fill=config.fill, tracers=tracers, - exclude_tracers=exclude_tracers, ) self._map_single_w = MapSingle( diff --git a/pyfv3/stencils/tracer_2d_1l.py b/pyfv3/stencils/tracer_2d_1l.py index 045ea162..37e1bb7f 100644 --- a/pyfv3/stencils/tracer_2d_1l.py +++ b/pyfv3/stencils/tracer_2d_1l.py @@ -327,7 +327,7 @@ def __init__( def __call__( self, - tracers: TracersType, # type: ignore + tracers, dp1, x_mass_flux, y_mass_flux, diff --git a/pyfv3/testing/translate_fvdynamics.py b/pyfv3/testing/translate_fvdynamics.py index 27da22b1..dd00894f 100644 --- a/pyfv3/testing/translate_fvdynamics.py +++ b/pyfv3/testing/translate_fvdynamics.py @@ -232,10 +232,10 @@ def __init__( self.dycore: fv_dynamics.DynamicalCore | None = None self.stencil_factory = stencil_factory self._quantity_factory = QuantityFactory( - sizer=stencil_factory.grid_indexing._sizer, + sizer=grid.sizer, backend=stencil_factory.backend, ) - self.namelist: DynamicalCoreConfig = DynamicalCoreConfig.from_namelist(namelist) + self.namelist: DynamicalCoreConfig = DynamicalCoreConfig.from_f90nml(namelist) def state_from_inputs(self, inputs): tracers = self._quantity_factory._numpy.empty( diff --git a/pyfv3/utils/functional_validation.py b/pyfv3/utils/functional_validation.py index 3cd3fe29..3a279a8c 100644 --- a/pyfv3/utils/functional_validation.py +++ b/pyfv3/utils/functional_validation.py @@ -1,4 +1,3 @@ -import copy from collections.abc import Callable, Sequence import numpy as np @@ -38,26 +37,3 @@ def subset(data: np.ndarray) -> np.ndarray: ) return subset - - -def get_set_nan_func( - grid_indexing: GridIndexing, - dims: Sequence[str], - n_halo: tuple[tuple[int, int], tuple[int, int]] = ((0, 0), (0, 0)), -) -> Callable[[np.ndarray], np.ndarray]: - subset = get_subset_func(grid_indexing=grid_indexing, dims=dims, n_halo=n_halo) - - def set_nans(data: np.ndarray) -> np.ndarray: - try: - safe = copy.deepcopy(data) - data[:] = np.nan - # data_subset is a view of data, so modifying data_subset modifies data - data_subset = subset(data) - data_subset[:] = subset(safe) - except TypeError: - safe = copy.deepcopy(data.data) - data.data[:] = np.nan - data_subset = subset(data.data) - data_subset[:] = subset(safe) - - return set_nans diff --git a/pyfv3/wrappers/geos_wrapper.py b/pyfv3/wrappers/geos_wrapper.py index 0a5127e3..19fd36ed 100644 --- a/pyfv3/wrappers/geos_wrapper.py +++ b/pyfv3/wrappers/geos_wrapper.py @@ -176,7 +176,7 @@ def __init__( metric_terms = MetricTerms( quantity_factory=quantity_factory, communicator=self.communicator, - eta_file=namelist["grid_config"]["config"]["eta_file"], # type: ignore + eta_file=namelist["grid_config"]["config"]["eta_file"], ) grid_data = GridData.new_from_metric_terms(metric_terms) @@ -213,7 +213,8 @@ def __init__( self.dycore_state = pyfv3.DycoreState.init_zeros( quantity_factory=quantity_factory, - tracer_list=self._tracers_mapping, + tracer_count=len(self._tracers_mapping), + backend=quantity_factory.backend, ) self.dycore_state.bdt = self.dycore_config.dt_atmos From 120950b850586d83e811313b04032932408d948a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Mar 2026 12:10:24 -0400 Subject: [PATCH 247/252] [WIP] Remove the DIM constants --- pyfv3/mpi/sum.py | 6 +- pyfv3/stencils/compute_total_energy.py | 6 +- pyfv3/stencils/delnflux.py | 2 +- pyfv3/stencils/remapping_GEOS.py | 94 +++++++++--------- pyfv3/tracers.py | 4 +- .../translate/translate_Pressures_mapU.py | 6 +- .../translate/translate_Pressures_mapV.py | 6 +- .../translate/translate_cs_profile.py | 4 +- ...ranslate_lagrangian_contribution_interp.py | 14 ++- .../translate/translate_map1_ppm_W.py | 6 +- .../translate/translate_map1_ppm_delz.py | 6 +- .../translate/translate_map_scalar.py | 6 +- .../translate/translate_moistcvpluspt_2d.py | 4 +- .../translate/translate_remapping_GEOS.py | 98 +++++++++---------- .../translate/translate_scalar_profile.py | 4 +- .../translate/translate_tracer2d1l_cmax.py | 8 +- 16 files changed, 136 insertions(+), 138 deletions(-) diff --git a/pyfv3/mpi/sum.py b/pyfv3/mpi/sum.py index e5697326..52aaf701 100644 --- a/pyfv3/mpi/sum.py +++ b/pyfv3/mpi/sum.py @@ -2,7 +2,7 @@ from ndsl import Quantity, QuantityFactory from ndsl.comm.communicator import Communicator, ReductionOperator -from ndsl.constants import X_DIM, Y_DIM +from ndsl.constants import I_DIM, J_DIM from ndsl.dsl.dace.orchestration import dace_inhibitor from ndsl.dsl.stencil import GridIndexing from ndsl.dsl.typing import Float @@ -17,8 +17,8 @@ def __init__( grid_indexing: GridIndexing = None, ) -> None: self._comm = communicator - # self._tmp_reduce = quantity_factory.empty(dims=[X_DIM, Y_DIM], units="n/a") - self._tmp_reduce = quantity_factory.zeros(dims=[X_DIM, Y_DIM], units="n/a") + # self._tmp_reduce = quantity_factory.empty(dims=[I_DIM, J_DIM], units="n/a") + self._tmp_reduce = quantity_factory.zeros(dims=[I_DIM, J_DIM], units="n/a") self._isc = grid_indexing.isc self._iec = grid_indexing.iec self._jsc = grid_indexing.jsc diff --git a/pyfv3/stencils/compute_total_energy.py b/pyfv3/stencils/compute_total_energy.py index 4ca3fc13..73f7d3a2 100644 --- a/pyfv3/stencils/compute_total_energy.py +++ b/pyfv3/stencils/compute_total_energy.py @@ -1,7 +1,7 @@ from gt4py.cartesian.gtscript import BACKWARD, FORWARD, K, computation, interval from ndsl import QuantityFactory, StencilFactory, orchestrate -from ndsl.constants import GRAV, X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM +from ndsl.constants import GRAV, I_DIM, J_DIM, K_DIM, K_INTERFACE_DIM from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from ndsl.grid import GridData from pyfv3._config import DynamicalCoreConfig @@ -110,14 +110,14 @@ def __init__( ) self._phyz = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="Unknown", dtype=Float, ) self._compute_total_energy = stencil_factory.from_dims_halo( func=_compute_total_energy__stencil, - compute_dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], + compute_dims=[I_DIM, J_DIM, K_INTERFACE_DIM], ) self._rsin2 = grid_data.rsin2 self._cosa_s = grid_data.cosa_s diff --git a/pyfv3/stencils/delnflux.py b/pyfv3/stencils/delnflux.py index cfbdfc54..ee44aada 100644 --- a/pyfv3/stencils/delnflux.py +++ b/pyfv3/stencils/delnflux.py @@ -1,7 +1,7 @@ from typing import Optional import dace -import np +import numpy as np from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate from ndsl.constants import I_DIM, I_INTERFACE_DIM, J_DIM, J_INTERFACE_DIM, K_DIM diff --git a/pyfv3/stencils/remapping_GEOS.py b/pyfv3/stencils/remapping_GEOS.py index 920da858..0a557a28 100644 --- a/pyfv3/stencils/remapping_GEOS.py +++ b/pyfv3/stencils/remapping_GEOS.py @@ -1,40 +1,40 @@ from gt4py.cartesian.gtscript import FORWARD, computation, interval -from pyFV3._config import RemappingConfig -from pyFV3.mpi.sum import GlobalSum -from pyFV3.stencils.map_single import MapSingle -from pyFV3.stencils.mapn_tracer import MapNTracer -from pyFV3.stencils.moist_cv import moist_pt_last_step -from pyFV3.stencils.remapping import ( - CONSV_MIN, - init_pe, - moist_cv_pt_pressure, - pe0_ptop_xmax, - pe_pk_delp_peln, - pn2_pk_delp, - pressures_mapu, - pressures_mapv, -) -from pyFV3.stencils.saturation_adjustment import SatAdjust3d -from pyFV3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 -from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment -from pyFV3.tracers import TracersType from ndsl import QuantityFactory, StencilFactory, orchestrate from ndsl.comm.communicator import Communicator from ndsl.constants import ( CV_AIR, GRAV, - X_DIM, - X_INTERFACE_DIM, - Y_DIM, - Y_INTERFACE_DIM, - Z_DIM, - Z_INTERFACE_DIM, + I_DIM, + I_INTERFACE_DIM, + J_DIM, + J_INTERFACE_DIM, + K_DIM, + K_INTERFACE_DIM, ) from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldIJ64, FloatFieldK from ndsl.grid import GridData from ndsl.stencils.basic_operations import adjust_divide_stencil -from pyFV3.stencils import moist_cv +from pyfv3._config import RemappingConfig +from pyfv3.mpi.sum import GlobalSum +from pyfv3.stencils import moist_cv +from pyfv3.stencils.map_single import MapSingle +from pyfv3.stencils.mapn_tracer import MapNTracer +from pyfv3.stencils.moist_cv import moist_pt_last_step +from pyfv3.stencils.remapping import ( + CONSV_MIN, + init_pe, + moist_cv_pt_pressure, + pe0_ptop_xmax, + pe_pk_delp_peln, + pn2_pk_delp, + pressures_mapu, + pressures_mapv, +) +from pyfv3.stencils.saturation_adjustment import SatAdjust3d +from pyfv3.stencils.scale_delz import rescale_delz_1, rescale_delz_2 +from pyfv3.stencils.w_fix_consrv_moment import W_fix_consrv_moment +from pyfv3.tracers import TracersType def _normalize_to_grid_stencil( @@ -102,85 +102,85 @@ def __init__( # Quantities self._pe1 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], + [I_DIM, J_DIM, K_INTERFACE_DIM], units="Pa", dtype=Float, ) self._pe2 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], + [I_DIM, J_DIM, K_INTERFACE_DIM], units="Pa", dtype=Float, ) self._pe3 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], + [I_DIM, J_DIM, K_INTERFACE_DIM], units="Pa", dtype=Float, ) self._dp2 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="Pa", dtype=Float, ) self._pn1 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="Pa", dtype=Float, ) self._pn2 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="Pa", dtype=Float, ) self._pe0 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], + [I_DIM, J_DIM, K_INTERFACE_DIM], units="Pa", dtype=Float, ) self._pe3 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], + [I_DIM, J_DIM, K_INTERFACE_DIM], units="Pa", dtype=Float, ) self._gz = quantity_factory.zeros( - [X_DIM, Y_DIM], + [I_DIM, J_DIM], units="m^2 s^-2", dtype=Float, ) self._cvm = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="unknown", dtype=Float, ) self._compute_performed = quantity_factory.zeros( - [X_DIM, Y_DIM], + [I_DIM, J_DIM], units="mask", dtype=bool, ) self._w2 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="temp W", dtype=Float, ) self._pk2 = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="Pa", dtype=Float, ) self._te_2d = quantity_factory.zeros( - [X_DIM, Y_DIM], + [I_DIM, J_DIM], units="Pa", dtype=Float, ) self._zsum1 = quantity_factory.zeros( - [X_DIM, Y_DIM], + [I_DIM, J_DIM], units="Pa", dtype=Float, ) self._phis = quantity_factory.zeros( - [X_DIM, Y_DIM, Z_INTERFACE_DIM], + [I_DIM, J_DIM, K_INTERFACE_DIM], units="n/a", dtype=Float, ) @@ -217,7 +217,7 @@ def __init__( quantity_factory, self._kord_tm, mode=1, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], interpolate_contribution=True, ) @@ -234,7 +234,7 @@ def __init__( quantity_factory, self._kord_wz, mode=-2, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], ) self._map_single_delz = MapSingle( @@ -242,7 +242,7 @@ def __init__( quantity_factory, self._kord_wz, mode=1, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], ) self._moist_cv_pkz = stencil_factory.from_origin_domain( @@ -262,7 +262,7 @@ def __init__( quantity_factory, self._kord_mt, mode=-1, - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + dims=[I_DIM, J_INTERFACE_DIM, K_DIM], ) self._pressures_mapv = stencil_factory.from_origin_domain( @@ -276,7 +276,7 @@ def __init__( quantity_factory, self._kord_mt, mode=-1, - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + dims=[I_INTERFACE_DIM, J_DIM, K_DIM], ) self._saturation_adjustment = SatAdjust3d( diff --git a/pyfv3/tracers.py b/pyfv3/tracers.py index 6faefe74..c4c0d84f 100644 --- a/pyfv3/tracers.py +++ b/pyfv3/tracers.py @@ -1,7 +1,7 @@ from typing import TypeAlias from ndsl import QuantityFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType from pyfv3.version import IS_GEOS @@ -41,7 +41,7 @@ def setup_tracers( FieldBundleType.register("Tracers", (number_of_tracers,)) _unit = "g/kg" - _dims = [X_DIM, Y_DIM, Z_DIM, "tracers"] + _dims = [I_DIM, J_DIM, K_DIM, "tracers"] tracers_qty_factory = FieldBundle.extend_3D_quantity_factory( quantity_factory, {"tracers": number_of_tracers} diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index 07350db5..5029f5ae 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -3,7 +3,7 @@ from pyFV3.stencils.remapping import pe0_ptop_xmax, pressures_mapu from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, J_INTERFACE_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -106,7 +106,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto grid_indexing = stencil_factory.grid_indexing - self.dims = [X_DIM, Y_DIM, Z_DIM] + self.dims = [I_DIM, J_DIM, K_DIM] self._pressures_mapu = stencil_factory.from_origin_domain( pressures_mapu, @@ -126,7 +126,7 @@ def compute_from_storage(self, inputs): self.quantity_factory, inputs["kord_mt"], -1, - dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM], + dims=[I_DIM, J_INTERFACE_DIM, K_DIM], ) self._pressures_mapu( diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index 6427b268..b098b410 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -3,7 +3,7 @@ from pyFV3.stencils.remapping import pressures_mapv from ndsl import StencilFactory -from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, I_INTERFACE_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -105,7 +105,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto grid_indexing = stencil_factory.grid_indexing - self.dims = [X_DIM, Y_DIM, Z_DIM] + self.dims = [I_DIM, J_DIM, K_DIM] self._pressures_mapv = stencil_factory.from_origin_domain( pressures_mapv, @@ -123,7 +123,7 @@ def compute_from_storage(self, inputs): self.quantity_factory, inputs["kord_mt"], -1, - dims=[X_INTERFACE_DIM, Y_DIM, Z_DIM], + dims=[I_INTERFACE_DIM, J_DIM, K_DIM], ) self._pressures_mapv( diff --git a/tests/savepoint/translate/translate_cs_profile.py b/tests/savepoint/translate/translate_cs_profile.py index 12e90ded..89d1e2cc 100644 --- a/tests/savepoint/translate/translate_cs_profile.py +++ b/tests/savepoint/translate/translate_cs_profile.py @@ -2,7 +2,7 @@ from pyFV3.stencils.remap_profile import RemapProfile from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -100,7 +100,7 @@ def compute_from_storage(self, inputs): self.quantity_factory, inputs["kord_"], inputs["iv_"], - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], ) self._compute_func( diff --git a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py index daf29630..7c260171 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py @@ -2,7 +2,7 @@ from pyFV3.stencils.map_single import lagrangian_contributions_interp from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.dsl.typing import Bool, BoolFieldIJ, FloatField, Int, IntField, IntFieldIJ from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -62,9 +62,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto super().__init__(grid, stencil_factory) self.stencil_factory = stencil_factory self.grid = grid - self.compute_func = test_Lagragian_Contribution_Interp( - self.stencil_factory, self.grid - ) # type: ignore + self.compute_func = test_Lagragian_Contribution_Interp(self.stencil_factory, self.grid) # type: ignore self.quantity_factory = grid.quantity_factory self.in_vars["data_vars"] = { @@ -130,25 +128,25 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto def compute_from_storage(self, inputs): self._not_exit_loop = self.quantity_factory.zeros( - [X_DIM, Y_DIM], + [I_DIM, J_DIM], units="", dtype=Bool, ) self._INDEX_LM1 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="", dtype=Int, ) self._INDEX_LP0 = self.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="", dtype=Int, ) self._lev = self.quantity_factory.zeros( - [X_DIM, Y_DIM], + [I_DIM, J_DIM], units="", dtype=Int, ) diff --git a/tests/savepoint/translate/translate_map1_ppm_W.py b/tests/savepoint/translate/translate_map1_ppm_W.py index 6c399474..a72180e7 100644 --- a/tests/savepoint/translate/translate_map1_ppm_W.py +++ b/tests/savepoint/translate/translate_map1_ppm_W.py @@ -2,7 +2,7 @@ from pyFV3.stencils.map_single import MapSingle from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -52,7 +52,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto # mode / iv set to -2 from GEOS self.mode = -2 - self.dims = [X_DIM, Y_DIM, Z_DIM] + self.dims = [I_DIM, J_DIM, K_DIM] def compute_from_storage(self, inputs): self._compute_func = MapSingle( @@ -60,7 +60,7 @@ def compute_from_storage(self, inputs): self.quantity_factory, inputs["kord_wz"], self.mode, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], ) self._compute_func( diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py index e85d41bc..2e78d5fe 100644 --- a/tests/savepoint/translate/translate_map1_ppm_delz.py +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -3,7 +3,7 @@ from pyFV3.stencils.map_single import MapSingle from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.dsl.typing import FloatField from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -78,7 +78,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto # mode / iv set to 1 from GEOS self.mode = 1 - self.dims = [X_DIM, Y_DIM, Z_DIM] + self.dims = [I_DIM, J_DIM, K_DIM] self._rescale_delz_1 = stencil_factory.from_origin_domain( rescale_delz_1, @@ -98,7 +98,7 @@ def compute_from_storage(self, inputs): self.quantity_factory, inputs["kord_wz"], self.mode, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], ) self._rescale_delz_1( diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py index decfb0af..39a15ed1 100644 --- a/tests/savepoint/translate/translate_map_scalar.py +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -2,7 +2,7 @@ from pyFV3.stencils.map_single import MapSingle from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -49,14 +49,14 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto # mode / iv set to 1 from GEOS self.mode = 1 - self.dims = [X_DIM, Y_DIM, Z_DIM] + self.dims = [I_DIM, J_DIM, K_DIM] self._compute_func = MapSingle( self.stencil_factory, self.quantity_factory, self._kord_tm, self.mode, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], interpolate_contribution=True, ) diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d.py b/tests/savepoint/translate/translate_moistcvpluspt_2d.py index 33815138..a54740e3 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d.py @@ -1,7 +1,7 @@ from gt4py.cartesian.gtscript import PARALLEL, computation, interval from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.dsl.typing import Float, FloatField from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j from pyfv3.stencils import moist_cv @@ -55,7 +55,7 @@ def __init__( ) self._q_con = grid.quantity_factory.zeros( - [X_DIM, Y_DIM, Z_DIM], + [I_DIM, J_DIM, K_DIM], units="unknown", dtype=Float, ) diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index 10ff3bdd..ad02bba6 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,235 +1,235 @@ from types import SimpleNamespace from f90nml import Namelist +from pyFV3 import DynamicalCoreConfig from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS from pyFV3.tracers import TracersType, setup_tracers from ndsl import Quantity, StencilFactory from ndsl.constants import ( - X_DIM, - X_INTERFACE_DIM, - Y_DIM, - Y_INTERFACE_DIM, - Z_DIM, - Z_INTERFACE_DIM, + I_DIM, + I_INTERFACE_DIM, + J_DIM, + J_INTERFACE_DIM, + K_DIM, + K_INTERFACE_DIM, ) from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing -from pyFV3 import DynamicalCoreConfig class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): inputs = { "pe": { "name": "pe", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [I_DIM, J_DIM, K_INTERFACE_DIM], "units": "No Units", }, "delp": { "name": "delp", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "delz": { "name": "delz", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "q_con": { "name": "q_con", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "pt": { "name": "pt", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "cappa": { "name": "cappa", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "ps": { "name": "ps", - "dims": [X_DIM, Y_DIM], + "dims": [I_DIM, J_DIM], "units": "No Units", }, "peln": { "name": "peln", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [I_DIM, J_DIM, K_INTERFACE_DIM], "units": "No Units", }, "ak": { "name": "ak", - "dims": [Z_INTERFACE_DIM], + "dims": [K_INTERFACE_DIM], "units": "No Units", }, "bk": { "name": "bk", - "dims": [Z_INTERFACE_DIM], + "dims": [K_INTERFACE_DIM], "units": "No Units", }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "pkz": { "name": "pkz", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "w": { "name": "w", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "u": { "name": "u", - "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], + "dims": [I_DIM, J_INTERFACE_DIM, K_DIM], "units": "No Units", }, "v": { "name": "v", - "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], + "dims": [I_INTERFACE_DIM, J_DIM, K_DIM], "units": "No Units", }, "mfy_R4": { "name": "mfy", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "cy_R4": { "name": "cy", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "mfx_R4": { "name": "mfx", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "cx_R4": { "name": "cx", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "phis": { "name": "phis", - "dims": [X_DIM, Y_DIM], + "dims": [I_DIM, J_DIM], "units": "No Units", }, "te_2d": { "name": "te_2d", - "dims": [X_DIM, Y_DIM], + "dims": [I_DIM, J_DIM], "units": "No Units", }, "wsd": { "name": "wsd", - "dims": [X_DIM, Y_DIM], + "dims": [I_DIM, J_DIM], "units": "No Units", }, "dp1": { "name": "dp1", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "pfull": { "name": "pfull", - "dims": [Z_DIM], + "dims": [K_DIM], "units": "No Units", }, } outputs = { "pt": { "name": "pt", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "cappa": { "name": "cappa", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "delp": { "name": "delp", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "delz": { "name": "delz", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "w": { "name": "w", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "u": { "name": "u", - "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], + "dims": [I_DIM, J_INTERFACE_DIM, K_DIM], "units": "No Units", }, "v": { "name": "v", - "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], + "dims": [I_INTERFACE_DIM, J_DIM, K_DIM], "units": "No Units", }, "mfy_R4": { "name": "mfy", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "cy_R4": { "name": "cy", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "mfx_R4": { "name": "mfx", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "cx_R4": { "name": "cx", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "peln": { "name": "peln", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [I_DIM, J_DIM, K_INTERFACE_DIM], "units": "No Units", }, "pe": { "name": "pe", - "dims": [X_DIM, Y_DIM, Z_INTERFACE_DIM], + "dims": [I_DIM, J_DIM, K_INTERFACE_DIM], "units": "No Units", }, "pk": { "name": "pk", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "pkz": { "name": "pkz", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "q_con": { "name": "q_con", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "dp1": { "name": "dp1", - "dims": [X_DIM, Y_DIM, Z_DIM], + "dims": [I_DIM, J_DIM, K_DIM], "units": "No Units", }, "ps": { "name": "ps", - "dims": [X_DIM, Y_DIM], + "dims": [I_DIM, J_DIM], "units": "No Units", }, } diff --git a/tests/savepoint/translate/translate_scalar_profile.py b/tests/savepoint/translate/translate_scalar_profile.py index 8ef7f25d..4fdad09a 100644 --- a/tests/savepoint/translate/translate_scalar_profile.py +++ b/tests/savepoint/translate/translate_scalar_profile.py @@ -2,7 +2,7 @@ from pyFV3.stencils.remap_profile import RemapProfile from ndsl import StencilFactory -from ndsl.constants import X_DIM, Y_DIM, Z_DIM +from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid @@ -104,7 +104,7 @@ def __init__(self, grid: Grid, namelist: Namelist, stencil_factory: StencilFacto self.quantity_factory, self.kord, self.mode, - dims=[X_DIM, Y_DIM, Z_DIM], + dims=[I_DIM, J_DIM, K_DIM], ) def compute_from_storage(self, inputs): diff --git a/tests/savepoint/translate/translate_tracer2d1l_cmax.py b/tests/savepoint/translate/translate_tracer2d1l_cmax.py index a67ddcac..68594686 100644 --- a/tests/savepoint/translate/translate_tracer2d1l_cmax.py +++ b/tests/savepoint/translate/translate_tracer2d1l_cmax.py @@ -2,7 +2,7 @@ from pyFV3.stencils.tracer_2d_1l import TracerCMax from ndsl import Quantity, QuantityFactory, StencilFactory -from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM +from ndsl.constants import I_DIM, I_INTERFACE_DIM, J_DIM, J_INTERFACE_DIM, K_DIM from ndsl.stencils.testing import ParallelTranslate2Py @@ -21,17 +21,17 @@ class TranslateTracerCMax(ParallelTranslate2Py): inputs = { "cx_R4": { "name": "cx_R4", - "dims": [X_INTERFACE_DIM, Y_DIM, Z_DIM], + "dims": [I_INTERFACE_DIM, J_DIM, K_DIM], "units": "unitless", }, "cy_R4": { "name": "cy_R4", - "dims": [X_DIM, Y_INTERFACE_DIM, Z_DIM], + "dims": [I_DIM, J_INTERFACE_DIM, K_DIM], "units": "unitless", }, "cmax": { "name": "cmaxgrid", - "dims": [Z_DIM], + "dims": [K_DIM], "units": "unitless", }, } From bba011e7ce810653d182d83f5d4dde5ef696212d Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Mar 2026 12:19:40 -0400 Subject: [PATCH 248/252] pyFV3 -> pyfv3 --- pyfv3/dycore_state.py | 2 +- pyfv3/stencils/del2cubed.py | 2 +- pyfv3/stencils/pk3_halo.py | 2 +- pyfv3/stencils/tracer_2d_1l.py | 2 +- pyfv3/wrappers/geos_wrapper.py | 10 +++++----- tests/savepoint/translate/translate_MapN_Tracer_2d.py | 4 ++-- tests/savepoint/translate/translate_Pressures_mapU.py | 4 ++-- tests/savepoint/translate/translate_Pressures_mapV.py | 4 ++-- tests/savepoint/translate/translate_cond_output.py | 2 +- tests/savepoint/translate/translate_cs_profile.py | 2 +- .../translate_lagrangian_contribution_interp.py | 2 +- tests/savepoint/translate/translate_map1_ppm_W.py | 2 +- tests/savepoint/translate/translate_map1_ppm_delz.py | 2 +- tests/savepoint/translate/translate_map_scalar.py | 2 +- .../translate/translate_moistcvpluspt_2d_last_step.py | 2 +- .../savepoint/translate/translate_moistcvpluste_2d.py | 2 +- tests/savepoint/translate/translate_mpp_global_sum.py | 2 +- tests/savepoint/translate/translate_pe_pk_delp_peln.py | 2 +- tests/savepoint/translate/translate_pn2_pk_delp.py | 3 +-- tests/savepoint/translate/translate_remapping_GEOS.py | 6 +++--- tests/savepoint/translate/translate_scalar_profile.py | 2 +- tests/savepoint/translate/translate_te_zsum.py | 2 +- tests/savepoint/translate/translate_tracer2d1l_cmax.py | 2 +- .../translate/translate_w_fix_consrv_moment.py | 3 +-- 24 files changed, 33 insertions(+), 35 deletions(-) diff --git a/pyfv3/dycore_state.py b/pyfv3/dycore_state.py index 1557b717..ff13e2e3 100644 --- a/pyfv3/dycore_state.py +++ b/pyfv3/dycore_state.py @@ -377,7 +377,7 @@ def init_from_numpy_arrays( elif issubclass(_field.type, TracersType): if len(dict_of_numpy_arrays[_field.name]) != len(tracer_list): raise ValueError( - "[pyFV3] DycoreState init:" + "[pyfv3] DycoreState init:" f" tracer list size ({len(tracer_list)})" " doesn't match the inputs size" f" ({len(dict_of_numpy_arrays[_field.name])})" diff --git a/pyfv3/stencils/del2cubed.py b/pyfv3/stencils/del2cubed.py index 139e93d8..816a8405 100644 --- a/pyfv3/stencils/del2cubed.py +++ b/pyfv3/stencils/del2cubed.py @@ -90,7 +90,7 @@ def __init__( ): """ Args: - grid: pyFV3 grid object + grid: pyfv3 grid object """ orchestrate(obj=self, config=stencil_factory.config.dace_config) grid_indexing = stencil_factory.grid_indexing diff --git a/pyfv3/stencils/pk3_halo.py b/pyfv3/stencils/pk3_halo.py index f078bb12..c57c0b26 100644 --- a/pyfv3/stencils/pk3_halo.py +++ b/pyfv3/stencils/pk3_halo.py @@ -5,7 +5,7 @@ # TODO merge with pe_halo? reuse partials? -# NOTE: This is different from pyFV3.stencils.pe_halo.edge_pe +# NOTE: This is different from pyfv3.stencils.pe_halo.edge_pe def edge_pe_update( pe: FloatFieldIJ, delp: FloatField, pk3: FloatField, ptop: Float, akap: Float ): diff --git a/pyfv3/stencils/tracer_2d_1l.py b/pyfv3/stencils/tracer_2d_1l.py index 37e1bb7f..9a70f2ef 100644 --- a/pyfv3/stencils/tracer_2d_1l.py +++ b/pyfv3/stencils/tracer_2d_1l.py @@ -526,7 +526,7 @@ def _reduce(self, cmax: Quantity): if __debug__: if not isinstance(cmax, Quantity): raise TypeError( - f"[pyFV3][Tracer]: cmax must be a quantity, got {type(cmax)}" + f"[pyfv3][Tracer]: cmax must be a quantity, got {type(cmax)}" ) cmax.data[:] = self._tmp_cmax.data.max(axis=0).max(axis=0)[:] self._comm.all_reduce_per_element_in_place(cmax, ReductionOperator.MAX) diff --git a/pyfv3/wrappers/geos_wrapper.py b/pyfv3/wrappers/geos_wrapper.py index 19fd36ed..f9f5eb40 100644 --- a/pyfv3/wrappers/geos_wrapper.py +++ b/pyfv3/wrappers/geos_wrapper.py @@ -124,7 +124,7 @@ def __init__( # Check for water species configuration not handled by the interface if water_tracers_count != 6: raise NotImplementedError( - f"[pyFV3 Bridge] Bridge expect 6 water species, got {water_tracers_count}." + f"[pyfv3 Bridge] Bridge expect 6 water species, got {water_tracers_count}." ) # Build the full tracer mapping by appending None to the expected tracer list @@ -339,11 +339,11 @@ def __call__( # Collect performance of the timestep and write a json file for rank 0 self.perf_collector.collect_performance() - for k, v in self.perf_collector.times_per_step[0].items(): + for k, v in self.perf_collector.times_per_step[0].items(): # type: ignore if k not in timings.keys(): - timings[k] = [v] + timings[k] = [v] # type: ignore else: - timings[k].append(v) + timings[k].append(v) # type: ignore self.perf_collector.clear() return self.output_dict, timings @@ -528,7 +528,7 @@ def _prep_outputs_for_geos(self) -> dict[str, np.ndarray]: return output_dict - def _allocate_output_dir(self): + def _allocate_output_dir(self) -> None: if len(self.output_dict) != 0: return if self._fortran_mem_space != self._pace_mem_space: diff --git a/tests/savepoint/translate/translate_MapN_Tracer_2d.py b/tests/savepoint/translate/translate_MapN_Tracer_2d.py index 48826b98..03c5ad0b 100644 --- a/tests/savepoint/translate/translate_MapN_Tracer_2d.py +++ b/tests/savepoint/translate/translate_MapN_Tracer_2d.py @@ -1,10 +1,10 @@ from f90nml import Namelist -from pyFV3.stencils.mapn_tracer import MapNTracer -from pyFV3.tracers import setup_tracers from ndsl import StencilFactory from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.mapn_tracer import MapNTracer +from pyfv3.tracers import setup_tracers class TranslateMapN_Tracer_2d(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_Pressures_mapU.py b/tests/savepoint/translate/translate_Pressures_mapU.py index 5029f5ae..bfa8d93f 100644 --- a/tests/savepoint/translate/translate_Pressures_mapU.py +++ b/tests/savepoint/translate/translate_Pressures_mapU.py @@ -1,11 +1,11 @@ from f90nml import Namelist -from pyFV3.stencils.map_single import MapSingle -from pyFV3.stencils.remapping import pe0_ptop_xmax, pressures_mapu from ndsl import StencilFactory from ndsl.constants import I_DIM, J_DIM, J_INTERFACE_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.map_single import MapSingle +from pyfv3.stencils.remapping import pe0_ptop_xmax, pressures_mapu class TranslatePressures_mapU(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_Pressures_mapV.py b/tests/savepoint/translate/translate_Pressures_mapV.py index b098b410..1e473d06 100644 --- a/tests/savepoint/translate/translate_Pressures_mapV.py +++ b/tests/savepoint/translate/translate_Pressures_mapV.py @@ -1,11 +1,11 @@ from f90nml import Namelist -from pyFV3.stencils.map_single import MapSingle -from pyFV3.stencils.remapping import pressures_mapv from ndsl import StencilFactory from ndsl.constants import I_DIM, I_INTERFACE_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.map_single import MapSingle +from pyfv3.stencils.remapping import pressures_mapv class TranslatePressures_mapV(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_cond_output.py b/tests/savepoint/translate/translate_cond_output.py index 6bd94b2e..0fecf45c 100644 --- a/tests/savepoint/translate/translate_cond_output.py +++ b/tests/savepoint/translate/translate_cond_output.py @@ -1,5 +1,5 @@ from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils import moist_cv +from pyfv3.stencils import moist_cv class TranslateCond_output(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_cs_profile.py b/tests/savepoint/translate/translate_cs_profile.py index 89d1e2cc..0e8d3487 100644 --- a/tests/savepoint/translate/translate_cs_profile.py +++ b/tests/savepoint/translate/translate_cs_profile.py @@ -1,10 +1,10 @@ from f90nml import Namelist -from pyFV3.stencils.remap_profile import RemapProfile from ndsl import StencilFactory from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.remap_profile import RemapProfile class TranslateCS_Profile(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py index 7c260171..046d9766 100644 --- a/tests/savepoint/translate/translate_lagrangian_contribution_interp.py +++ b/tests/savepoint/translate/translate_lagrangian_contribution_interp.py @@ -1,11 +1,11 @@ from f90nml import Namelist -from pyFV3.stencils.map_single import lagrangian_contributions_interp from ndsl import StencilFactory from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.dsl.typing import Bool, BoolFieldIJ, FloatField, Int, IntField, IntFieldIJ from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.map_single import lagrangian_contributions_interp class test_Lagragian_Contribution_Interp: diff --git a/tests/savepoint/translate/translate_map1_ppm_W.py b/tests/savepoint/translate/translate_map1_ppm_W.py index a72180e7..5c3ee95b 100644 --- a/tests/savepoint/translate/translate_map1_ppm_W.py +++ b/tests/savepoint/translate/translate_map1_ppm_W.py @@ -1,10 +1,10 @@ from f90nml import Namelist -from pyFV3.stencils.map_single import MapSingle from ndsl import StencilFactory from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.map_single import MapSingle class TranslateMap1_PPM_W(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_map1_ppm_delz.py b/tests/savepoint/translate/translate_map1_ppm_delz.py index 2e78d5fe..f0afac1f 100644 --- a/tests/savepoint/translate/translate_map1_ppm_delz.py +++ b/tests/savepoint/translate/translate_map1_ppm_delz.py @@ -1,12 +1,12 @@ from f90nml import Namelist from gt4py.cartesian.gtscript import PARALLEL, computation, interval -from pyFV3.stencils.map_single import MapSingle from ndsl import StencilFactory from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.dsl.typing import FloatField from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.map_single import MapSingle def rescale_delz_1( diff --git a/tests/savepoint/translate/translate_map_scalar.py b/tests/savepoint/translate/translate_map_scalar.py index 39a15ed1..e8fdcda3 100644 --- a/tests/savepoint/translate/translate_map_scalar.py +++ b/tests/savepoint/translate/translate_map_scalar.py @@ -1,10 +1,10 @@ from f90nml import Namelist -from pyFV3.stencils.map_single import MapSingle from ndsl import StencilFactory from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.map_single import MapSingle class TranslateMap_Scalar(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py index f947444a..7c9b2aa5 100644 --- a/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py +++ b/tests/savepoint/translate/translate_moistcvpluspt_2d_last_step.py @@ -1,6 +1,6 @@ from ndsl.dsl.typing import Float from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils import moist_cv +from pyfv3.stencils import moist_cv class TranslateMoistCVPlusPt_2d_last_step(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_moistcvpluste_2d.py b/tests/savepoint/translate/translate_moistcvpluste_2d.py index 8723bf3a..23880cc9 100644 --- a/tests/savepoint/translate/translate_moistcvpluste_2d.py +++ b/tests/savepoint/translate/translate_moistcvpluste_2d.py @@ -1,5 +1,5 @@ from ndsl.stencils.testing import TranslateFortranData2Py, pad_field_in_j -from pyFV3.stencils import moist_cv +from pyfv3.stencils import moist_cv class TranslateMoistCVPlusTe_2d(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_mpp_global_sum.py b/tests/savepoint/translate/translate_mpp_global_sum.py index 5471d033..53b1e4ca 100644 --- a/tests/savepoint/translate/translate_mpp_global_sum.py +++ b/tests/savepoint/translate/translate_mpp_global_sum.py @@ -1,10 +1,10 @@ from f90nml import Namelist -from pyFV3.mpi.mpp_sum import MPPGlobalSum from ndsl import StencilFactory from ndsl.stencils.testing import ParallelTranslate from ndsl.stencils.testing.grid import Grid from ndsl.typing import Communicator +from pyfv3.mpi.mpp_sum import MPPGlobalSum class TranslateMpp_global_sum(ParallelTranslate): diff --git a/tests/savepoint/translate/translate_pe_pk_delp_peln.py b/tests/savepoint/translate/translate_pe_pk_delp_peln.py index c7526edb..e9dc935f 100644 --- a/tests/savepoint/translate/translate_pe_pk_delp_peln.py +++ b/tests/savepoint/translate/translate_pe_pk_delp_peln.py @@ -1,9 +1,9 @@ from f90nml import Namelist -from pyFV3.stencils.remapping import pe_pk_delp_peln from ndsl import StencilFactory from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.remapping import pe_pk_delp_peln class TranslatePE_pk_delp_peln(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_pn2_pk_delp.py b/tests/savepoint/translate/translate_pn2_pk_delp.py index 4016e91c..fe8eb041 100644 --- a/tests/savepoint/translate/translate_pn2_pk_delp.py +++ b/tests/savepoint/translate/translate_pn2_pk_delp.py @@ -1,8 +1,7 @@ -from pyFV3.stencils.remapping import pn2_pk_delp - from ndsl import StencilFactory from ndsl.dsl.typing import Float, FloatField from ndsl.stencils.testing import TranslateFortranData2Py +from pyfv3.stencils.remapping import pn2_pk_delp class testClass: diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index ad02bba6..d83a12a7 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -1,9 +1,6 @@ from types import SimpleNamespace from f90nml import Namelist -from pyFV3 import DynamicalCoreConfig -from pyFV3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS -from pyFV3.tracers import TracersType, setup_tracers from ndsl import Quantity, StencilFactory from ndsl.constants import ( @@ -16,6 +13,9 @@ ) from ndsl.dsl.typing import Float from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing +from pyfv3 import DynamicalCoreConfig +from pyfv3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS +from pyfv3.tracers import TracersType, setup_tracers class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): diff --git a/tests/savepoint/translate/translate_scalar_profile.py b/tests/savepoint/translate/translate_scalar_profile.py index 4fdad09a..269aadf4 100644 --- a/tests/savepoint/translate/translate_scalar_profile.py +++ b/tests/savepoint/translate/translate_scalar_profile.py @@ -1,10 +1,10 @@ from f90nml import Namelist -from pyFV3.stencils.remap_profile import RemapProfile from ndsl import StencilFactory from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.remap_profile import RemapProfile class TranslateScalar_Profile(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_te_zsum.py b/tests/savepoint/translate/translate_te_zsum.py index 2b39e1a5..9b7a3191 100644 --- a/tests/savepoint/translate/translate_te_zsum.py +++ b/tests/savepoint/translate/translate_te_zsum.py @@ -1,5 +1,5 @@ from ndsl.stencils.testing import TranslateFortranData2Py -from pyFV3.stencils import moist_cv +from pyfv3.stencils import moist_cv class TranslateTe_Zsum(TranslateFortranData2Py): diff --git a/tests/savepoint/translate/translate_tracer2d1l_cmax.py b/tests/savepoint/translate/translate_tracer2d1l_cmax.py index 68594686..8932196b 100644 --- a/tests/savepoint/translate/translate_tracer2d1l_cmax.py +++ b/tests/savepoint/translate/translate_tracer2d1l_cmax.py @@ -1,9 +1,9 @@ from f90nml import Namelist -from pyFV3.stencils.tracer_2d_1l import TracerCMax from ndsl import Quantity, QuantityFactory, StencilFactory from ndsl.constants import I_DIM, I_INTERFACE_DIM, J_DIM, J_INTERFACE_DIM, K_DIM from ndsl.stencils.testing import ParallelTranslate2Py +from pyfv3.stencils.tracer_2d_1l import TracerCMax def _quantity_wrap(storage, dims, grid_indexing): diff --git a/tests/savepoint/translate/translate_w_fix_consrv_moment.py b/tests/savepoint/translate/translate_w_fix_consrv_moment.py index bdaee31e..2784d366 100644 --- a/tests/savepoint/translate/translate_w_fix_consrv_moment.py +++ b/tests/savepoint/translate/translate_w_fix_consrv_moment.py @@ -1,8 +1,7 @@ -from pyFV3.stencils.w_fix_consrv_moment import W_fix_consrv_moment - from ndsl.dsl.typing import Float from ndsl.stencils.testing import TranslateFortranData2Py from ndsl.stencils.testing.grid import Grid +from pyfv3.stencils.w_fix_consrv_moment import W_fix_consrv_moment class TranslateW_fix_consrv_moment(TranslateFortranData2Py): From b00d2d0b362184ae33bb97ff69918ecbe620b923 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Mar 2026 14:22:44 -0400 Subject: [PATCH 249/252] [WIPP] Fix acoustics translate tests --- tests/savepoint/translate/translate_d_sw.py | 6 +++++- tests/savepoint/translate/translate_updatedzc.py | 2 +- tests/savepoint/translate/translate_updatedzd.py | 7 ++++--- 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/savepoint/translate/translate_d_sw.py b/tests/savepoint/translate/translate_d_sw.py index 27f8c4db..b4cbc671 100644 --- a/tests/savepoint/translate/translate_d_sw.py +++ b/tests/savepoint/translate/translate_d_sw.py @@ -242,7 +242,8 @@ def __init__( def compute_from_storage(self, inputs): column_namelist = d_sw.get_column_namelist( - config=self.config, quantity_factory=self.grid.quantity_factory + config=self.config.d_grid_shallow_water, + quantity_factory=self.grid.quantity_factory, ) # TODO add these to the serialized data or remove the test inputs["damp_w"] = column_namelist["damp_w"] @@ -255,6 +256,9 @@ def compute_from_storage(self, inputs): d_sw.heat_diss, origin=self.grid.compute_origin(), domain=self.grid.domain_shape_compute(), + externals={ + "do_stochastic_ke_backscatter": self.config.do_skeb, + }, ) heat_diss_stencil(**inputs) return inputs diff --git a/tests/savepoint/translate/translate_updatedzc.py b/tests/savepoint/translate/translate_updatedzc.py index 6b54bea5..b81516df 100644 --- a/tests/savepoint/translate/translate_updatedzc.py +++ b/tests/savepoint/translate/translate_updatedzc.py @@ -24,7 +24,7 @@ def compute(**kwargs): quantity_factory=self.grid.quantity_factory, area=grid.grid_data.area, dp_ref=grid.grid_data.dp_ref, - grid_type=namelist.grid_type, + grid_type=self.config.grid_type, dz_min=kwargs.pop("dz_min"), ) diff --git a/tests/savepoint/translate/translate_updatedzd.py b/tests/savepoint/translate/translate_updatedzd.py index 4e623b07..92d74766 100644 --- a/tests/savepoint/translate/translate_updatedzd.py +++ b/tests/savepoint/translate/translate_updatedzd.py @@ -33,7 +33,7 @@ def __init__( self.in_vars["data_vars"]["height"]["serialname"] = "zh" self.in_vars["data_vars"]["ws"]["serialname"] = "wsd" - self.in_vars["parameters"] = ["dt", "dz_min"] + self.in_vars["parameters"] = ["dt"] out_vars = [ "height", "courant_number_x", @@ -65,9 +65,10 @@ def compute(self, inputs): self.grid.grid_data, self.grid.grid_type, self.config.hord_tm, - dz_min=inputs.pop("dz_min"), + dz_min=self.config.acoustic_dynamics.dz_min, column_namelist=d_sw.get_column_namelist( - self.config, quantity_factory=self.grid.quantity_factory + self.config.d_grid_shallow_water, + quantity_factory=self.grid.quantity_factory, ), ) self.updatedzd(**inputs) From d2a372f12c9cd31a3e9753eecf02499781164d92 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 12 Mar 2026 15:49:42 -0400 Subject: [PATCH 250/252] [WIP] Fix dyncore translate test (and some lint) --- pyfv3/dycore_state.py | 62 ++++++++++--------- .../test_cases/initialize_baroclinic.py | 14 ++++- .../test_cases/initialize_rossby.py | 4 +- .../test_cases/initialize_tc.py | 15 ++++- pyfv3/testing/translate_dyncore.py | 5 +- pyfv3/testing/translate_fvdynamics.py | 7 ++- pyfv3/tracers.py | 37 +++++------ pyfv3/wrappers/geos_wrapper.py | 6 +- .../translate/overrides/standard.yaml | 6 +- 9 files changed, 92 insertions(+), 64 deletions(-) diff --git a/pyfv3/dycore_state.py b/pyfv3/dycore_state.py index ff13e2e3..dd41a62f 100644 --- a/pyfv3/dycore_state.py +++ b/pyfv3/dycore_state.py @@ -20,7 +20,7 @@ from ndsl.quantity.field_bundle import FieldBundle from ndsl.restart._legacy_restart import open_restart from ndsl.typing import Communicator -from pyfv3.tracers import TracersType, setup_tracers +from pyfv3.tracers import TracersType, make_tracers DEFAULT_TRACER_PROPERTIES = { @@ -314,11 +314,12 @@ def __post_init__(self) -> None: def init_zeros( cls, quantity_factory: QuantityFactory, - tracer_count: int, - backend: Backend, dtype_dict: dict[str, type] | None = None, - allow_mismatch_float_precision: bool = False, + allow_mismatch_float_precision: bool = True, ) -> Self: + """Initialize the dynamics memory state to zero. Default to allow for + mixed precision as 32-bit dynamics requires it.""" + initial_storages = {} for _field in fields(cls): if "dims" in _field.metadata.keys(): @@ -333,10 +334,12 @@ def init_zeros( allow_mismatch_float_precision=allow_mismatch_float_precision, ).data elif _field.name == "tracers": - qty_factory_tracers = FieldBundle.extend_3D_quantity_factory( - quantity_factory, {"tracers": tracer_count} - ) - initial_storages[_field.name] = qty_factory_tracers.zeros( + if "tracers" not in quantity_factory.sizer.data_dimensions: + raise ValueError( + "[pyFV3] Tracers data dimensions undefined in Quantity Factory." + "You need to register the tracers!" + ) + initial_storages[_field.name] = quantity_factory.zeros( [I_DIM, J_DIM, K_DIM, "tracers"], _field.metadata["units"], dtype=Float, @@ -344,7 +347,7 @@ def init_zeros( return cls.init_from_storages( storages=initial_storages, quantity_factory=quantity_factory, - backend=backend, + allow_mismatch_float_precision=allow_mismatch_float_precision, ) @classmethod @@ -352,8 +355,6 @@ def init_from_numpy_arrays( cls, dict_of_numpy_arrays: dict, sizer: GridSizer, - backend: Backend, - tracer_list: list[str], quantity_factory: QuantityFactory, ) -> Self: field_names = [_field.name for _field in fields(cls)] @@ -372,20 +373,26 @@ def init_from_numpy_arrays( _field.metadata["units"], origin=quantity_factory.sizer.get_origin(dims), extent=quantity_factory.sizer.get_extent(dims), - backend=backend, + backend=quantity_factory.backend, ) - elif issubclass(_field.type, TracersType): - if len(dict_of_numpy_arrays[_field.name]) != len(tracer_list): + elif isinstance(_field.type, type) and issubclass(_field.type, TracersType): + if "tracers" not in quantity_factory.sizer.data_dimensions: + raise ValueError( + "[pyFV3] Tracers data dimensions undefined in Quantity Factory." + "You need to register the tracers!" + ) + quantity_factory.sizer.data_dimensions["tracers"] + if ( + len(dict_of_numpy_arrays[_field.name]) + != quantity_factory.sizer.data_dimensions["tracers"] + ): raise ValueError( "[pyfv3] DycoreState init:" - f" tracer list size ({len(tracer_list)})" + f" tracer list size ({quantity_factory.sizer.data_dimensions['tracers']})" " doesn't match the inputs size" f" ({len(dict_of_numpy_arrays[_field.name])})" ) - dict_state[_field.name] = Tracers.make( - quantity_factory=quantity_factory, - tracer_mapping=tracer_list, - ) + dict_state[_field.name] = make_tracers(quantity_factory) state = cls(**dict_state) return state @@ -396,10 +403,8 @@ def init_from_storages( quantity_factory: QuantityFactory, bdt: float = 0.0, mdt: float = 0.0, - backend: Backend | None = None, + allow_mismatch_float_precision: bool = False, ) -> Self: - if not backend: - backend = Backend.python() inputs = {} for _field in fields(cls): if "dims" in _field.metadata.keys(): @@ -410,11 +415,12 @@ def init_from_storages( _field.metadata["units"], origin=quantity_factory.sizer.get_origin(dims), extent=quantity_factory.sizer.get_extent(dims), - backend=backend, + backend=quantity_factory.backend, + allow_mismatch_float_precision=allow_mismatch_float_precision, ) inputs[_field.name] = quantity elif "tracers" == _field.name: - tracers = setup_tracers(storages["tracers"].shape[3], quantity_factory) + tracers = make_tracers(quantity_factory) tracers.quantity.data[:] = storages["tracers"][:] inputs[_field.name] = tracers @@ -436,8 +442,6 @@ def from_fortran_restart( ) new = cls.init_zeros( quantity_factory=quantity_factory, - tracer_count=len(DEFAULT_TRACER_PROPERTIES), - backend=backend, ) new.pt.view[:] = new.pt.np.asarray( state_dict["air_temperature"].transpose(new.pt.dims).view[:] @@ -497,7 +501,7 @@ def from_fortran_restart( def _xr_dataarray_from_array( self, name: str, metadata: MappingProxyType[Any, Any], data: npt.ArrayLike - ): + ) -> xr.DataArray: dims = [f"{dim_name}_{name}" for dim_name in metadata["dims"]] return xr.DataArray( gt_utils.asarray(data), @@ -512,7 +516,9 @@ def _xr_dataarray_from_array( def xr_dataset(self) -> xr.Dataset: data_vars = {} for name, field_info in self.__dataclass_fields__.items(): - if issubclass(field_info.type, Quantity): + if isinstance(field_info.type, type) and issubclass( + field_info.type, Quantity + ): data_vars[name] = self._xr_dataarray_from_array( name=name, metadata=field_info.metadata, diff --git a/pyfv3/initialization/test_cases/initialize_baroclinic.py b/pyfv3/initialization/test_cases/initialize_baroclinic.py index 90e69ad1..9cc4ae26 100644 --- a/pyfv3/initialization/test_cases/initialize_baroclinic.py +++ b/pyfv3/initialization/test_cases/initialize_baroclinic.py @@ -8,6 +8,7 @@ from ndsl.grid.gnomonic import great_circle_distance_lon_lat, lon_lat_midpoint from pyfv3.dycore_state import DycoreState from pyfv3.initialization import init_utils +from pyfv3.tracers import setup_tracers # maximum windspeed amplitude - close to windspeed of zonal-mean time-mean @@ -366,11 +367,20 @@ def init_baroclinic_state( moist_phys=moist_phys, make_nh=(not hydrostatic), ) + tracers = { + "vapor": 0, + "liquid": 1, + "rain": 2, + "snow": 3, + "ice": 4, + "graupel": 5, + "cloud": 6, + } + setup_tracers(len(tracers), quantity_factory, tracers) state = DycoreState.init_from_numpy_arrays( numpy_state.__dict__, + sizer=quantity_factory.sizer, quantity_factory=quantity_factory, - backend=sample_quantity.backend, - tracer_list=["vapor", "liquid", "rain", "snow", "ice", "graupel", "cloud"], ) state.tracers["vapor"].view[:] = numpy_state.qvapor[slice_3d] diff --git a/pyfv3/initialization/test_cases/initialize_rossby.py b/pyfv3/initialization/test_cases/initialize_rossby.py index 36a57f1d..72ec9ac4 100644 --- a/pyfv3/initialization/test_cases/initialize_rossby.py +++ b/pyfv3/initialization/test_cases/initialize_rossby.py @@ -13,6 +13,7 @@ from ndsl.grid import GridData from pyfv3.dycore_state import DycoreState from pyfv3.initialization import init_utils +from pyfv3.tracers import setup_tracers NHALO = constants.N_HALO_DEFAULT @@ -200,10 +201,11 @@ def init_rossby_state( _init_for_rossby(numpy_state, grid_data, shape) _postinit_for_all_sw(numpy_state) + setup_tracers(1, quantity_factory) state = DycoreState.init_from_numpy_arrays( numpy_state.__dict__, sizer=quantity_factory.sizer, - backend=sample_quantity.metadata.backend, + quantity_factory=quantity_factory, ) comm.halo_update(state.phis, n_points=NHALO) diff --git a/pyfv3/initialization/test_cases/initialize_tc.py b/pyfv3/initialization/test_cases/initialize_tc.py index 9576ce1b..8f81b73f 100644 --- a/pyfv3/initialization/test_cases/initialize_tc.py +++ b/pyfv3/initialization/test_cases/initialize_tc.py @@ -6,6 +6,7 @@ from ndsl.grid.gnomonic import great_circle_distance_lon_lat from pyfv3.dycore_state import DycoreState from pyfv3.initialization import init_utils +from pyfv3.tracers import setup_tracers def _calculate_distance_from_tc_center(pe_v, ps_v, muv, calc, tc_properties): @@ -566,12 +567,20 @@ def init_tc_state( numpy_state.v[:] = vd numpy_state.va[:] = va numpy_state.w[:] = w + tracers = { + "vapor": 0, + "liquid": 1, + "rain": 2, + "snow": 3, + "ice": 4, + "graupel": 5, + "cloud": 6, + } + setup_tracers(len(tracers), quantity_factory, tracers) state = DycoreState.init_from_numpy_arrays( numpy_state.__dict__, - quantity_factory=quantity_factory, - backend=sample_quantity.metadata.backend, - tracer_list=["vapor", "liquid", "rain", "snow", "ice", "graupel", "cloud"], sizer=quantity_factory.sizer, + quantity_factory=quantity_factory, ) state.tracers["vapor"].view[:] = qvapor diff --git a/pyfv3/testing/translate_dyncore.py b/pyfv3/testing/translate_dyncore.py index 5f2fc76a..b31cc447 100644 --- a/pyfv3/testing/translate_dyncore.py +++ b/pyfv3/testing/translate_dyncore.py @@ -8,6 +8,7 @@ from pyfv3._config import DynamicalCoreConfig from pyfv3.dycore_state import DycoreState from pyfv3.stencils import dyn_core +from pyfv3.tracers import setup_tracers class TranslateDynCore(ParallelTranslate2PyState): @@ -148,10 +149,10 @@ def compute_parallel(self, inputs: dict, communicator: Communicator) -> dict: for k, v in inputs.items(): if hasattr(v, "dtype"): inputs_dtypes[k] = v.dtype + setup_tracers(1, self.grid.quantity_factory) # No tracers used in acoustics state = DycoreState.init_zeros( quantity_factory=self.grid.quantity_factory, dtype_dict=inputs_dtypes, - tracer_count=1, # No tracers used in acoustics allow_mismatch_float_precision=True, ) wsd = self.grid.quantity_factory.zeros( @@ -202,7 +203,7 @@ def compute_parallel(self, inputs: dict, communicator: Communicator) -> dict: cyd=state.cyd, dpx=dpx, timestep=inputs["mdt"], - n_map=state.n_map, + n_map=inputs["n_map"], ) # the "inputs" dict is not used to return, we construct a new dict based # on variables attached to `state` diff --git a/pyfv3/testing/translate_fvdynamics.py b/pyfv3/testing/translate_fvdynamics.py index dd00894f..7ca9741a 100644 --- a/pyfv3/testing/translate_fvdynamics.py +++ b/pyfv3/testing/translate_fvdynamics.py @@ -2,6 +2,7 @@ from datetime import timedelta from typing import Any +import numpy as np import pytest from f90nml import Namelist @@ -228,7 +229,7 @@ def __init__( self.max_error = 1e-5 - self.ignore_near_zero_errors = {} + self.ignore_near_zero_errors: dict[str, float | bool] = {} self.dycore: fv_dynamics.DynamicalCore | None = None self.stencil_factory = stencil_factory self._quantity_factory = QuantityFactory( @@ -237,8 +238,8 @@ def __init__( ) self.namelist: DynamicalCoreConfig = DynamicalCoreConfig.from_f90nml(namelist) - def state_from_inputs(self, inputs): - tracers = self._quantity_factory._numpy.empty( + def state_from_inputs(self, inputs: dict[str, np.ndarray]) -> DycoreState: + tracers = self._quantity_factory.empty( ( inputs["tracers"].shape[0] + 1, inputs["tracers"].shape[1] + 1, diff --git a/pyfv3/tracers.py b/pyfv3/tracers.py index c4c0d84f..f59cddb2 100644 --- a/pyfv3/tracers.py +++ b/pyfv3/tracers.py @@ -1,5 +1,3 @@ -from typing import TypeAlias - from ndsl import QuantityFactory from ndsl.constants import I_DIM, J_DIM, K_DIM from ndsl.quantity.field_bundle import FieldBundle, FieldBundleType @@ -28,26 +26,20 @@ } -TracersType: TypeAlias = FieldBundleType.T("Tracers") # type: ignore # noqa +TracersType = FieldBundleType.T("Tracers") + +_mappings: dict[str, int] = {} def setup_tracers( number_of_tracers: int, quantity_factory: QuantityFactory, mappings: dict[str, int] | None = None, -) -> FieldBundle: - """Setup a FieldBundle for tracers. Should be called only once.""" +) -> None: + global _mappings FieldBundleType.register("Tracers", (number_of_tracers,)) - _unit = "g/kg" - _dims = [I_DIM, J_DIM, K_DIM, "tracers"] - - tracers_qty_factory = FieldBundle.extend_3D_quantity_factory( - quantity_factory, {"tracers": number_of_tracers} - ) - data = tracers_qty_factory.zeros(_dims, units=_unit) - # Some default mappings for ease of use with commonly # run models if mappings is None: @@ -56,8 +48,17 @@ def setup_tracers( else: mappings = _default_mapping_PACE - return FieldBundle( - "Tracers", - quantity=data, - mapping=mappings, - ) + _mappings = mappings + quantity_factory.add_data_dimensions({"tracers": number_of_tracers}) + + +def make_tracers(quantity_factory: QuantityFactory) -> FieldBundle: + """Setup a FieldBundle for tracers. Should be called only once.""" + global _mappings # noqa + + _unit = "g/kg" + _dims = [I_DIM, J_DIM, K_DIM, "tracers"] + + data = quantity_factory.zeros(_dims, units=_unit) + + return FieldBundle("Tracers", quantity=data, mapping=_mappings) diff --git a/pyfv3/wrappers/geos_wrapper.py b/pyfv3/wrappers/geos_wrapper.py index f9f5eb40..766efa1c 100644 --- a/pyfv3/wrappers/geos_wrapper.py +++ b/pyfv3/wrappers/geos_wrapper.py @@ -212,9 +212,7 @@ def __init__( ) self.dycore_state = pyfv3.DycoreState.init_zeros( - quantity_factory=quantity_factory, - tracer_count=len(self._tracers_mapping), - backend=quantity_factory.backend, + quantity_factory=quantity_factory ) self.dycore_state.bdt = self.dycore_config.dt_atmos @@ -339,7 +337,7 @@ def __call__( # Collect performance of the timestep and write a json file for rank 0 self.perf_collector.collect_performance() - for k, v in self.perf_collector.times_per_step[0].items(): # type: ignore + for k, v in self.perf_collector.times_per_step[0].items(): if k not in timings.keys(): timings[k] = [v] # type: ignore else: diff --git a/tests/savepoint/translate/overrides/standard.yaml b/tests/savepoint/translate/overrides/standard.yaml index 38e67a41..7a4333cd 100644 --- a/tests/savepoint/translate/overrides/standard.yaml +++ b/tests/savepoint/translate/overrides/standard.yaml @@ -27,16 +27,16 @@ MapN_Tracer_2d: near_zero: 1e-17 ignore_near_zero_errors: - qtracers - - backend: numpy + - backend: st:numpy:cpu:IJK max_error: 9e-9 # 48_6ranks NH_P_Grad: - max_error: 6e-11 + - max_error: 6e-11 Riem_Solver3: - backend: st:gt:gpu:KJI max_error: 5e-6 - - backend: gt:cpu_ifirst + - backend: st:gt:cpu:KJI max_error: 5e-6 - backend: st:dace:gpu:KJI max_error: 5e-6 From 239de93e1e7bb6431cf166ad531b8beab0d4815a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 13 Mar 2026 15:28:35 -0400 Subject: [PATCH 251/252] Fix `dz_min` type in config Fix tracers in Remap & Tracer2D1L modules --- pyfv3/_config.py | 3 +- .../translate/translate_remapping_GEOS.py | 31 +++++++------------ .../translate/translate_tracer2d1l.py | 5 +-- 3 files changed, 16 insertions(+), 23 deletions(-) diff --git a/pyfv3/_config.py b/pyfv3/_config.py index 5e0ec525..78511377 100644 --- a/pyfv3/_config.py +++ b/pyfv3/_config.py @@ -8,6 +8,7 @@ import yaml from dacite import Config, from_dict +from ndsl.constants import Float from ndsl.utils import f90nml_as_dict @@ -121,7 +122,7 @@ class AcousticDynamicsConfig: """ riemann: RiemannConfig d_grid_shallow_water: DGridShallowWaterLagrangianDynamicsConfig - dz_min: float = 2.0 + dz_min: Float = Float(2.0) """Controls minimum thickness in NH solver""" @property diff --git a/tests/savepoint/translate/translate_remapping_GEOS.py b/tests/savepoint/translate/translate_remapping_GEOS.py index d83a12a7..055e4f80 100644 --- a/tests/savepoint/translate/translate_remapping_GEOS.py +++ b/tests/savepoint/translate/translate_remapping_GEOS.py @@ -15,7 +15,7 @@ from ndsl.stencils.testing import Grid, ParallelTranslateBaseSlicing from pyfv3 import DynamicalCoreConfig from pyfv3.stencils.remapping_GEOS import LagrangianToEulerian_GEOS -from pyfv3.tracers import TracersType, setup_tracers +from pyfv3.tracers import TracersType, make_tracers, setup_tracers class TranslateRemapping_GEOS(ParallelTranslateBaseSlicing): @@ -368,7 +368,7 @@ def __init__( self.quantity_factory = grid.quantity_factory self.stencil_factory = stencil_factory - self.namelist = DynamicalCoreConfig.from_namelist(namelist) + self.config = DynamicalCoreConfig.from_f90nml(namelist) self.grid = grid self._are_tracers_setup = False @@ -378,7 +378,9 @@ def __init__( def compute_sequential(self, inputs_list, communicator_list): print("No serial test available") - def state_from_inputs(self, inputs: dict, tracers: TracersType) -> SimpleNamespace: + def state_from_inputs_and_tracers( + self, inputs: dict, tracers: TracersType + ) -> SimpleNamespace: input_storages = super().state_from_inputs(inputs) # Rename fluxes and courant numbers input_storages["mfx"] = input_storages.pop("mfx_R4") @@ -409,23 +411,9 @@ def outputs_from_state(self, state: dict): return outputs def compute_parallel(self, inputs, communicator): - # tracers_mapping = Tracers.blind_mapping_from_data(inputs["tracers"]) - # tracers_mapping[0] = "vapor" - # tracers_mapping[1] = "liquid" - # tracers_mapping[2] = "rain" - # tracers_mapping[3] = "snow" - # tracers_mapping[4] = "ice" - # tracers_mapping[5] = "graupel" - # tracers_mapping[6] = "cloud" - # tracers = Tracers.make_from_4D_array( - # self.quantity_factory, - # tracers_mapping[0:7], - # inputs["tracers"], - # ) - if not self._are_tracers_setup: self._are_tracers_setup = True - self._tracers = setup_tracers( + setup_tracers( number_of_tracers=inputs["tracers"].shape[3], quantity_factory=self.quantity_factory, mappings={ @@ -439,15 +427,18 @@ def compute_parallel(self, inputs, communicator): }, ) + self._tracers = make_tracers(self.quantity_factory) self._tracers.quantity.data[:-1, :-1, :-1, :] = inputs["tracers"] + inputs.pop("tracers") + self._base.in_vars["data_vars"].pop("tracers") inputs["te_2d"] = inputs["te_2d"].astype(Float) - state = self.state_from_inputs(inputs, self._tracers) + state = self.state_from_inputs_and_tracers(inputs, self._tracers) l_to_e = LagrangianToEulerian_GEOS( self.stencil_factory, self.quantity_factory, - DynamicalCoreConfig.from_namelist(self.namelist).remapping, + self.config.remapping, communicator, self.grid.grid_data, state.nq, diff --git a/tests/savepoint/translate/translate_tracer2d1l.py b/tests/savepoint/translate/translate_tracer2d1l.py index 6e368de0..c7b90619 100644 --- a/tests/savepoint/translate/translate_tracer2d1l.py +++ b/tests/savepoint/translate/translate_tracer2d1l.py @@ -6,7 +6,7 @@ from ndsl.stencils.testing import Grid, ParallelTranslate from pyfv3 import DynamicalCoreConfig from pyfv3.stencils import FiniteVolumeTransport, TracerAdvection -from pyfv3.tracers import setup_tracers +from pyfv3.tracers import make_tracers, setup_tracers from pyfv3.utils.functional_validation import get_subset_func @@ -49,10 +49,11 @@ def __init__( def compute_parallel(self, inputs, communicator): self._base.make_storage_data_input_vars(inputs, dict_4d=False) - tracers = setup_tracers( + setup_tracers( number_of_tracers=inputs["tracers"].shape[3], quantity_factory=self._quantity_factory, ) + tracers = make_tracers(self._quantity_factory) tracers.quantity.data[:] = inputs["tracers"][:] inputs.pop("tracers") inputs.pop("nq") # Fortran NQ is intrinsic to Tracers (e.g Tracers.count) From 556785d7ea67a8baa7225408558cddb1afb40557 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Fri, 13 Mar 2026 15:40:35 -0400 Subject: [PATCH 252/252] Fix copy corners for numpy/quantity usage --- pyfv3/stencils/copy_corners.py | 59 ++++++++++++++++++++++++++++++++-- pyfv3/stencils/fvtp2d.py | 4 +-- 2 files changed, 58 insertions(+), 5 deletions(-) diff --git a/pyfv3/stencils/copy_corners.py b/pyfv3/stencils/copy_corners.py index 266627df..9514f038 100644 --- a/pyfv3/stencils/copy_corners.py +++ b/pyfv3/stencils/copy_corners.py @@ -1,8 +1,35 @@ -from ndsl import StencilFactory, orchestrate +from functools import singledispatch + +import numpy as np + +from ndsl import Quantity, StencilFactory, orchestrate from ndsl.dsl.typing import FloatField +from ndsl.optional_imports import cupy as cp +@singledispatch def corner_copy_x(field_to_copy): + raise NotImplementedError(f"No CopyCorners for type {type(field_to_copy)}") + + +if cp is not None: + + @corner_copy_x.register(cp.ndarray) + def _corner_copy_x_cupy(field_to_copy: cp.ndarray): + _blind_copy_corners_x(field_to_copy) + + +@corner_copy_x.register(np.ndarray) +def _corner_copy_x_numpy(field_to_copy: np.ndarray): + _blind_copy_corners_x(field_to_copy) + + +@corner_copy_x.register(Quantity) +def _corner_copy_x_quantity(field_to_copy: Quantity): + _blind_copy_corners_x(field_to_copy.data) + + +def _blind_copy_corners_x(field_to_copy): """Equivalent to the copy_corners_x functions in fortran. This is written to operate on plain ndarrarys and not use the GT4Py framework. @@ -65,7 +92,29 @@ def corner_copy_x(field_to_copy): field_to_copy[-2, -4] = field_to_copy[-4, -7] +@singledispatch def corner_copy_y(field_to_copy): + raise NotImplementedError(f"No CopyCorners for type {type(field_to_copy)}") + + +if cp is not None: + + @corner_copy_y.register(cp.ndarray) + def _corner_copy_y_cupy(field_to_copy: cp.ndarray): + _blind_copy_corners_y(field_to_copy) + + +@corner_copy_y.register(np.ndarray) +def _corner_copy_y_nupy(field_to_copy: np.ndarray): + _blind_copy_corners_y(field_to_copy) + + +@corner_copy_y.register(Quantity) +def _corner_copy_y_quantity(field_to_copy: Quantity): + _blind_copy_corners_y(field_to_copy.data) + + +def _blind_copy_corners_y(field_to_copy): """Equivalent to the copy_corners_y functions in fortran. This is written to operate on plain ndarrarys and not use the GT4Py framework. @@ -144,8 +193,10 @@ def __init__(self, stencil_factory: StencilFactory) -> None: "Corner-Copy only implemented for exactly 3 Halo-Points" ) + self._is_orch = stencil_factory.backend.is_orchestrated() + def __call__(self, field: FloatField): - corner_copy_x(field) + corner_copy_x(field) if not self._is_orch else _blind_copy_corners_x(field) class CopyCornersY: @@ -165,5 +216,7 @@ def __init__(self, stencil_factory: StencilFactory) -> None: "Corner-Copy only implemented for exactly 3 Halo-Points" ) + self._is_orch = stencil_factory.backend.is_orchestrated() + def __call__(self, field: FloatField): - corner_copy_y(field) + corner_copy_y(field) if not self._is_orch else _blind_copy_corners_y(field) diff --git a/pyfv3/stencils/fvtp2d.py b/pyfv3/stencils/fvtp2d.py index fc205c29..3114cabf 100644 --- a/pyfv3/stencils/fvtp2d.py +++ b/pyfv3/stencils/fvtp2d.py @@ -311,7 +311,7 @@ def __call__( # y_area_flux as an input (flux = area_flux * advected_mean), since a flux is # easier to understand than the current output. This would be like merging # yppm with q_i_stencil and xppm with q_j_stencil. - self._copy_corners_y(q.data) + self._copy_corners_y(q) self.y_piecewise_parabolic_inner(q, cry, self._q_y_advected_mean) # q_y_advected_mean is 1/Delta_area * curly-F, where curly-F is defined in # equation 4.3 of the FV3 documentation and Delta_area is the advected area @@ -328,7 +328,7 @@ def __call__( ) # q_advected_y_x_advected_mean is now rho^n + F(rho^y) in PL07 eq 16 - self._copy_corners_x(q.data) + self._copy_corners_x(q) # similarly below for x<->y self.x_piecewise_parabolic_inner(q, crx, self._q_x_advected_mean) self.q_j_stencil(