New NDSL basic function API

pyshield/stencils/gfdl_cld_microphysics/gfdl_cld_mp_driver.py

@@ -1,7 +1,6 @@
 import numpy as np
 
 import ndsl.constants as constants
-import ndsl.stencils.basic_operations as basic
 import pyshield.stencils.gfdl_cld_microphysics.constants as mpcons
 import pyshield.stencils.gfdl_cld_microphysics.physical_functions as physfun
 from ndsl import GridIndexing, QuantityFactory, StencilFactory
@@ -10,6 +9,7 @@
 from ndsl.dsl.typing import Bool, Float, FloatField, FloatFieldIJ
 from ndsl.grid import GridData
 from ndsl.performance.timer import NullTimer, Timer
+from ndsl.stencils import copy
 from pyshield.stencils.gfdl_cld_microphysics._config import GFDLCloudMPConfig
 from pyshield.stencils.gfdl_cld_microphysics.cloud_fraction import CloudFraction
 from pyshield.stencils.gfdl_cld_microphysics.gfdl_cld_microphysics_state import (
@@ -899,7 +899,7 @@ def make_quantity2d(**kwargs):
         self._convert_mm_day = 86400.0 * constants.RGRAV / self.config.dt_split
 
         self._copy_stencil = stencil_factory.from_origin_domain(
-            basic.copy,
+            copy,
             origin=self._idx.origin_compute(),
             domain=self._idx.domain_compute(),
         )

pyshield/stencils/gfdl_cld_microphysics/ice_cloud.py

@@ -1,11 +1,11 @@
-import ndsl.stencils.basic_operations as basic
 import pyshield.stencils.gfdl_cld_microphysics.constants as mpcons
 import pyshield.stencils.gfdl_cld_microphysics.physical_functions as physfun
 from ndsl import GridIndexing, StencilFactory
 from ndsl.dsl.gt4py import FORWARD, PARALLEL, computation, exp
 from ndsl.dsl.gt4py import function as gtfunction
 from ndsl.dsl.gt4py import interval, log
 from ndsl.dsl.typing import FloatField, FloatFieldIJ
+from ndsl.stencils.arithmetical_functions import dim
 from pyshield.stencils.gfdl_cld_microphysics._config import GFDLCloudMPConfig
 
 
@@ -39,7 +39,7 @@ def melt_cloud_ice(
     if (tc > 0.0) and (qice > mpcons.QCMIN):
         sink = fac_imlt * tc / icpk
         sink = min(qice, sink)
-        tmp = min(sink, basic.dim(ql_mlt, qliquid))
+        tmp = min(sink, dim(ql_mlt, qliquid))
 
         (
             qvapor,
@@ -114,7 +114,7 @@ def freeze_cloud_water(
         sink = qliquid * tc / mpcons.DT_FR
         sink = min(qliquid, min(sink, tc / icpk))
         qim = qi0_crt / density
-        tmp = min(sink, basic.dim(qim, qice))
+        tmp = min(sink, dim(qim, qice))
 
         (
             qvapor,
@@ -303,7 +303,7 @@ def melt_snow(
             ),
         )
         sink = min(qsnow, min((sink + pracs) * timestep, tc / icpk))
-        tmp = min(sink, basic.dim(qs_mlt, qliquid))
+        tmp = min(sink, dim(qs_mlt, qliquid))
 
         (
             qvapor,
@@ -943,7 +943,7 @@ def accrete_graupel_with_cloud_water_and_rain(
             )
 
         sink = pgacr + pgacw
-        factor = min(sink, basic.dim(mpcons.TICE0, temperature) / icpk) / max(
+        factor = min(sink, dim(mpcons.TICE0, temperature) / icpk) / max(
             sink, mpcons.QCMIN
         )
         pgacr = factor * pgacr

pyshield/stencils/gfdl_cld_microphysics/mp_fast.py

@@ -1,11 +1,11 @@
-import ndsl.stencils.basic_operations as basic
 import pyshield.stencils.gfdl_cld_microphysics.constants as mpcons
 import pyshield.stencils.gfdl_cld_microphysics.physical_functions as physfun
 from ndsl import GridIndexing, StencilFactory
 from ndsl.dsl.gt4py import FORWARD, computation, exp
 from ndsl.dsl.gt4py import function as gtfunction
 from ndsl.dsl.gt4py import interval
 from ndsl.dsl.typing import Bool, FloatField, FloatFieldIJ
+from ndsl.stencils.arithmetical_functions import dim
 from pyshield.stencils.gfdl_cld_microphysics._config import FastMPConfig
 from pyshield.stencils.gfdl_cld_microphysics.ice_cloud import (
     freeze_cloud_water,
@@ -120,7 +120,7 @@ def melt_snow_simple(
     if (tc > 0.0) and (qsnow > mpcons.QCMIN):
         sink = (tc * 0.1) ** 2 * qsnow
         sink = min(qsnow, sink, fac_smlt * tc / icpk)
-        tmp = min(sink, basic.dim(qs_mlt, qliquid))
+        tmp = min(sink, dim(qs_mlt, qliquid))
 
         (
             qvapor,

pyshield/stencils/gfdl_cld_microphysics/physical_functions.py

@@ -1,9 +1,9 @@
 import ndsl.constants as constants
-import ndsl.stencils.basic_operations as basic
 import pyshield.stencils.gfdl_cld_microphysics.constants as mpcons
 from ndsl.dsl.gt4py import exp, floor
 from ndsl.dsl.gt4py import function as gtfunction
 from ndsl.dsl.gt4py import log, max, min, sqrt
+from ndsl.stencils.arithmetical_functions import dim
 
 
 @gtfunction
@@ -129,7 +129,7 @@ def calc_heat_cap_and_latent_heat_coeff(
     icpk = (li00 + d1_ice * temperature) / cvm
     tcpk = (li20 + (d1_vap + d1_ice) * temperature) / cvm
     tcp3 = lcpk + icpk * min(
-        1.0, basic.dim(mpcons.TICE0, temperature) / (mpcons.TICE0 - t_wfr)
+        1.0, dim(mpcons.TICE0, temperature) / (mpcons.TICE0 - t_wfr)
     )
 
     return q_liq, q_solid, cvm, te, lcpk, icpk, tcpk, tcp3
@@ -182,7 +182,7 @@ def update_hydrometeors_and_temperatures(
     lcpk = (lv00 + d1_vap * tk) / cvm
     icpk = (li00 + d1_ice * tk) / cvm
     tcpk = (li20 + (d1_vap + d1_ice) * tk) / cvm
-    tcp3 = lcpk + icpk * min(1.0, basic.dim(mpcons.TICE0, tk) / (mpcons.TICE0 - t_wfr))
+    tcp3 = lcpk + icpk * min(1.0, dim(mpcons.TICE0, tk) / (mpcons.TICE0 - t_wfr))
 
     return (
         qvapor,

pyshield/stencils/gfdl_cld_microphysics/subgrid_z_proc.py

@@ -1,12 +1,12 @@
 import ndsl.constants as constants
-import ndsl.stencils.basic_operations as basic
 import pyshield.stencils.gfdl_cld_microphysics.constants as mpcons
 import pyshield.stencils.gfdl_cld_microphysics.physical_functions as physfun
 from ndsl import GridIndexing, StencilFactory
 from ndsl.dsl.gt4py import FORWARD, computation, exp
 from ndsl.dsl.gt4py import function as gtfunction
 from ndsl.dsl.gt4py import interval, log
 from ndsl.dsl.typing import Bool, FloatField, FloatFieldIJ
+from ndsl.stencils.arithmetical_functions import dim
 from pyshield.stencils.gfdl_cld_microphysics._config import GFDLCloudMPConfig
 
 
@@ -44,7 +44,7 @@ def perform_instant_processes(
 
     # Instant deposit all water vapor to cloud ice when temperature is super low
     if temperature < t_min:
-        sink = basic.dim(qvapor, mpcons.QCMIN)
+        sink = dim(qvapor, mpcons.QCMIN)
         dep += sink * delp
 
         (
@@ -368,7 +368,7 @@ def wegener_bergeron_findeisen(
 
         sink = min(fac_wbf * qliquid, tc / icpk)
         qim = qi0_crt / density
-        tmp = min(sink, basic.dim(qim, qice))
+        tmp = min(sink, dim(qim, qice))
 
         (
             qvapor,
@@ -606,7 +606,7 @@ def deposit_and_sublimate_ice(
             sink = min(tmp, min(max(qi_crt - qice, pidep), tc / tcpk))
             dep += sink * delp
         else:
-            pidep = pidep * min(1, basic.dim(temperature, t_sub) * is_fac)
+            pidep = pidep * min(1, dim(temperature, t_sub) * is_fac)
             sink = max(pidep, max(tmp, -qice))
             sub -= sink * delp
 
@@ -726,7 +726,7 @@ def deposit_and_sublimate_snow(
         dq = dq / (1 + tcpk * dqdt)
 
         if pssub > 0:
-            sink = min(pssub * min(1.0, basic.dim(temperature, t_sub) * ss_fac), qsnow)
+            sink = min(pssub * min(1.0, dim(temperature, t_sub) * ss_fac), qsnow)
             sub += sink * delp
         else:
             sink = 0.0
@@ -851,9 +851,7 @@ def deposit_and_sublimate_graupel(
         dq = dq / (1 + tcpk * dqdt)
 
         if pgsub > 0:
-            sink = min(
-                pgsub * min(1.0, basic.dim(temperature, t_sub) * gs_fac), qgraupel
-            )
+            sink = min(pgsub * min(1.0, dim(temperature, t_sub) * gs_fac), qgraupel)
             sub += sink * delp
         else:
             sink = 0.0

pyshield/stencils/gfdl_cld_microphysics/terminal_fall.py

@@ -7,7 +7,7 @@
 from ndsl.constants import I_DIM, J_DIM, K_DIM
 from ndsl.dsl.gt4py import BACKWARD, FORWARD, PARALLEL, computation, interval
 from ndsl.dsl.typing import FloatField, FloatFieldIJ, IntFieldIJ
-from ndsl.stencils.basic_operations import copy
+from ndsl.stencils import copy
 from pyfv3.stencils.remap_profile import RemapProfile
 from pyshield.stencils.gfdl_cld_microphysics._config import GFDLCloudMPConfig
 

pyshield/stencils/physics.py

@@ -21,7 +21,7 @@
 )
 from ndsl.grid import GridData
 from ndsl.logging import ndsl_log
-from ndsl.stencils.basic_operations import copy
+from ndsl.stencils import copy
 from pyshield._config import (
     PHYSICS_PACKAGES,
     TRACER_DIM,

pyshield/stencils/surface/sfc_diff.py

@@ -7,7 +7,7 @@
 from ndsl.dsl.gt4py import function as gtfunction
 from ndsl.dsl.gt4py import interval, log, sqrt
 from ndsl.dsl.typing import Bool, BoolFieldIJ, Float, FloatFieldIJ, Int, IntFieldIJ
-from ndsl.stencils.basic_operations import sign
+from ndsl.stencils.arithmetical_functions import sign
 from pyshield.functions.physics_functions import fpvsx
 
 

tests/savepoint/translate/translate_icesub.py

@@ -1,11 +1,11 @@
-import ndsl.stencils.basic_operations as basic
 import pyshield.stencils.gfdl_cld_microphysics.constants as mpcons
 import pyshield.stencils.gfdl_cld_microphysics.physical_functions as physfun
 from ndsl import StencilFactory
 from ndsl.dsl.gt4py import interval  # noqa
 from ndsl.dsl.gt4py import FORWARD, PARALLEL, computation  # noqa
 from ndsl.dsl.gt4py import function as gtfunction
 from ndsl.dsl.typing import FloatField, FloatFieldIJ
+from ndsl.stencils.arithmetical_functions import dim
 from pyshield.stencils.gfdl_cld_microphysics import GFDLCloudMPConfig
 from pyshield.stencils.gfdl_cld_microphysics.ice_cloud import (  # noqa
     accrete_graupel_with_cloud_water_and_rain,
@@ -175,7 +175,7 @@ def melt_snow_test(
         )
         sink = max(0.0, sink0)
         sink = min(qsnow, min((sink + pracs) * timestep, tc / icpk))
-        tmp = min(sink, basic.dim(qs_mlt, qliquid))
+        tmp = min(sink, dim(qs_mlt, qliquid))
 
         (
             qvapor,

tests/savepoint/translate/translate_pbl_subtests.py

@@ -11,7 +11,7 @@
     Int,
     IntFieldIJ,
 )
-from ndsl.stencils.basic_operations import copy
+from ndsl.stencils import copy
 from pyshield._config import TRACER_DIM, FloatFieldTracer
 from pyshield.stencils.pbl import PBLConfig
 from pyshield.stencils.pbl import constants as pbl_constants

tests/savepoint/translate/translate_preliminary_mp.py

@@ -1,6 +1,6 @@
-import ndsl.stencils.basic_operations as basic
 from ndsl import QuantityFactory, StencilFactory, SubtileGridSizer
 from ndsl.constants import I_DIM, J_DIM
+from ndsl.stencils import copy
 from pyshield.stencils.gfdl_cld_microphysics import GFDLCloudMPConfig
 from pyshield.stencils.gfdl_cld_microphysics.gfdl_cld_mp_driver import (
     calculate_density_factor,
@@ -39,7 +39,7 @@ def make_quantity2d(**kwargs):
         self._bottom_density = make_quantity2d()
 
         self._copy_stencil = stencil_factory.from_origin_domain(
-            basic.copy,
+            copy,
             origin=self._idx.origin_compute(),
             domain=self._idx.domain_compute(),
         )

tests/savepoint/translate/translate_subsub.py

@@ -1,11 +1,11 @@
-import ndsl.stencils.basic_operations as basic  # noqa
 import pyshield.stencils.gfdl_cld_microphysics.constants as mpcons
 import pyshield.stencils.gfdl_cld_microphysics.physical_functions as physfun  # noqa
 from ndsl import GridIndexing, StencilFactory
 from ndsl.dsl.gt4py import FORWARD, computation, exp  # noqa
 from ndsl.dsl.gt4py import function as gtfunction
 from ndsl.dsl.gt4py import interval, log  # noqa
 from ndsl.dsl.typing import FloatField, FloatFieldIJ
+from ndsl.stencils.arithmetical_functions import dim  # noqa
 from pyshield.stencils.gfdl_cld_microphysics import GFDLCloudMPConfig
 from pyshield.stencils.gfdl_cld_microphysics.subgrid_z_proc import (  # noqa
     cloud_condensation_evaporation,
@@ -52,7 +52,7 @@ def perform_instant_processes_test(
 
     # Instant deposit all water vapor to cloud ice when temperature is super low
     if temperature < t_min:
-        sink = basic.dim(qvapor, mpcons.QCMIN)
+        sink = dim(qvapor, mpcons.QCMIN)
         dep += sink * delp
 
         (

tests/savepoint/translate/translate_tridiag.py

@@ -3,7 +3,7 @@
 from ndsl import QuantityFactory, StencilFactory, SubtileGridSizer
 from ndsl.constants import I_DIM, J_DIM, K_DIM
 from ndsl.dsl.typing import Float
-from ndsl.stencils.basic_operations import copy
+from ndsl.stencils import copy
 from pyshield._config import TRACER_DIM, FloatFieldTracer
 from pyshield.stencils.pbl.tridiag import tridi2, tridin, tridit
 from tests.savepoint.translate.translate_physics import TranslatePhysicsFortranData2Py