Merging most recent changes

.pre-commit-config.yaml

@@ -0,0 +1,50 @@
+default_language_version:
+  python: python3
+
+repos:
+-   repo: https://github.com/psf/black
+    rev: 20.8b1
+    hooks:
+    - id: black
+      additional_dependencies: ["click==8.0.4"]
+
+-   repo: https://github.com/pre-commit/mirrors-isort
+    rev: v5.4.2
+    hooks:
+    - id: isort
+      args: ["--profile", "black"]
+
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v0.812
+    hooks:
+    - id: mypy
+      name: mypy-ndsl
+      args: [--config-file, setup.cfg]
+      files: ndsl
+      exclude: |
+        (?x)^(
+        ndsl/ndsl/gt4py_utils.py |
+        )$
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+    - id: check-toml
+    - id: check-yaml
+    - id: end-of-file-fixer
+    - id: trailing-whitespace
+-   repo: https://github.com/pycqa/flake8
+    rev: 3.9.2
+    hooks:
+    - id: flake8
+      name: flake8
+      language_version: python3
+      args: [--config, setup.cfg]
+      exclude: |
+        (?x)^(
+        .*/__init__.py |
+        )$
+    - id: flake8
+      name: flake8 __init__.py files
+      files: "__init__.py"
+      # ignore unused import error in __init__.py files
+      args: ["--ignore=F401,E203", --config, setup.cfg]

README.md

@@ -1 +1,3 @@
-# NDSL
+# NOAA/NASA Domain Specific Language middleware
+
+Use `git clone --recurse-submodule` to pull all vetted versions of the submodules used by `ndsl`

ndsl/comm/_boundary_utils.py

@@ -1,8 +1,8 @@
 import functools
 from typing import Union
 
-from .. import constants
-from ..exceptions import OutOfBoundsError
+import ndsl.constants as constants
+from ndsl.exceptions import OutOfBoundsError
 
 
 def shift_boundary_slice_tuple(dims, origin, extent, boundary_type, slice_tuple):

ndsl/dsl/dace/wrapped_halo_exchange.py

@@ -1,5 +1,5 @@
 import dataclasses
-from typing import List, Optional
+from typing import Any, List, Optional
 
 from ndsl.comm.communicator import Communicator
 from ndsl.dsl.dace.orchestration import dace_inhibitor
@@ -29,6 +29,14 @@ def __init__(
         self._qtx_y_names = qty_y_names
         self._comm = comm
 
+    @staticmethod
+    def check_for_attribute(state: Any, attr: str):
+        if dataclasses.is_dataclass(state):
+            return state.__getattribute__(attr)
+        elif isinstance(state, dict):
+            return attr in state.keys()
+        return False
+
     @dace_inhibitor
     def start(self):
         if self._qtx_y_names is None:

ndsl/grid/eta.py

@@ -1,10 +1,15 @@
+import math
 import os
 from dataclasses import dataclass
 
 import numpy as np
 import xarray as xr
 
 
+ETA_0 = 0.252
+SURFACE_PRESSURE = 1.0e5  # units of (Pa), from Table VI of DCMIP2016
+
+
 @dataclass
 class HybridPressureCoefficients:
     """
@@ -75,7 +80,23 @@ def set_hybrid_pressure_coefficients(
     return pressure_data
 
 
-def check_eta(ak, bk):
-    from pace.fv3core.initialization.init_utils import compute_eta
+def vertical_coordinate(eta_value):
+    """
+    Equation (1) JRMS2006
+    computes eta_v, the auxiliary variable vertical coordinate
+    """
+    return (eta_value - ETA_0) * math.pi * 0.5
+
 
+def compute_eta(ak, bk):
+    """
+    Equation (1) JRMS2006
+    eta is the vertical coordinate and eta_v is an auxiliary vertical coordinate
+    """
+    eta = 0.5 * ((ak[:-1] + ak[1:]) / SURFACE_PRESSURE + bk[:-1] + bk[1:])
+    eta_v = vertical_coordinate(eta)
+    return eta, eta_v
+
+
+def check_eta(ak, bk):
     return compute_eta(ak, bk)

ndsl/grid/generation.py

@@ -220,7 +220,7 @@ class GridDefinitions:
 
 # TODO
 # corners use sizer + partitioner rather than GridIndexer,
-# have to refactor fv3core calls to corners to do this as well
+# have to refactor pyFV3 calls to corners to do this as well
 class MetricTerms:
     LON_OR_LAT_DIM = GridDefinitions.LON_OR_LAT_DIM
     TILE_DIM = GridDefinitions.TILE_DIM

ndsl/performance/__init__.py

@@ -1 +1,2 @@
 from .config import PerformanceConfig
+from .timer import NullTimer, Timer

ndsl/stencils/c2l_ord.py

@@ -15,7 +15,6 @@
 from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ
 from ndsl.grid import GridData
 from ndsl.initialization.allocator import QuantityFactory
-from pace import fv3core
 
 
 A1 = 0.5625
@@ -157,7 +156,7 @@ class CubedToLatLon:
 
     def __init__(
         self,
-        state: fv3core.DycoreState,
+        state,  # No type hint on purpose to remove dependency on pyFV3
         stencil_factory: StencilFactory,
         quantity_factory: QuantityFactory,
         grid_data: GridData,
@@ -215,8 +214,6 @@ def __init__(
             compute_halos=halos,
         )
 
-        origin = grid_indexing.origin_compute()
-        shape = grid_indexing.max_shape
         if not self.one_rank:
             full_size_xyiz_halo_spec = quantity_factory.get_quantity_halo_spec(
                 dims=[X_DIM, Y_INTERFACE_DIM, Z_DIM],
@@ -228,6 +225,17 @@ def __init__(
                 n_halo=grid_indexing.n_halo,
                 dtype=Float,
             )
+
+            # TODO:
+            # To break the depedency to pyFV3 we allow ourselves to not have a type
+            # hint around state and we check for u and v to make sure we don't
+            # have bad input.
+            # This entire code should be retired when WrappedHaloUpdater is no longer
+            # required.
+            if not WrappedHaloUpdater.check_for_attribute(
+                state, "u"
+            ) and WrappedHaloUpdater.check_for_attribute(state, "v"):
+                raise RuntimeError("Cube To Lat Lon: state given is not readable.")
             self.u__v = WrappedHaloUpdater(
                 comm.get_vector_halo_updater(
                     [full_size_xyiz_halo_spec], [full_size_xiyz_halo_spec]

ndsl/stencils/testing/temporaries.py

@@ -16,10 +16,7 @@ def copy_temporaries(obj, max_depth: int) -> dict:
             attr = None
         if isinstance(attr, Quantity):
             temporaries[attr_name] = copy.deepcopy(np.asarray(attr.data))
-        elif attr.__class__.__module__.split(".")[0] in (  # type: ignore
-            "fv3core",
-            "pace",
-        ):
+        elif attr.__class__.__module__.split(".")[0] in ("pyFV3"):  # type: ignore
             if max_depth > 0:
                 sub_temporaries = copy_temporaries(attr, max_depth - 1)
                 if len(sub_temporaries) > 0:

setup.cfg

@@ -1,21 +1,25 @@
-[bumpversion]
-current_version = 0.10.0
-commit = True
-
-[bdist_wheel]
-universal = 1
-
 [flake8]
 exclude = docs
-ignore = E203,E501,W293,W503
+ignore = E203,E501,W293,W503,E302,E203,F841
 max-line-length = 88
 
 [aliases]
 
-[bumpversion:file:pace/util/__init__.py]
-search = __version__ = "{current_version}"
-replace = __version__ = "{new_version}"
+[tool:isort]
+line_length = 88
+force_grid_wrap = 0
+include_trailing_comma = true
+multi_line_output = 3
+use_parentheses = true
+lines_after_imports = 2
+default_section = THIRDPARTY
+sections = FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER
+known_third_party = f90nml,pytest,xarray,numpy,mpi4py,gt4py,dace
 
-[bumpversion:file:setup.py]
-search = version="{current_version}"
-replace = version="{new_version}"
+[mypy]
+ignore_missing_imports = True
+follow_imports = normal
+namespace_packages = True
+strict_optional = False
+warn_unreachable = True
+explicit_package_bases = True

setup.py

@@ -1,20 +1,35 @@
-from typing import List
-from pathlib import Path
 import os
+from pathlib import Path
+from typing import List
 
 from setuptools import find_namespace_packages, setup
 
-setup_requirements: List[str] = []
-
-test_requirements: List[str] = []
-
 
 def local_pkg(name: str, relative_path: str) -> str:
     """Returns an absolute path to a local package."""
     path = f"{name} @ file://{Path(os.path.abspath(__file__)).parent / relative_path}"
     return path
 
 
+test_requirements = ["pytest", "pytest-subtests"]
+develop_requirements = test_requirements + ["pre-commit"]
+
+extras_requires = {"test": test_requirements, "develop": develop_requirements}
+
+requirements: List[str] = [
+    local_pkg("gt4py", "external/gt4py"),
+    local_pkg("dace", "external/dace"),
+    "mpi4py",
+    "cftime",
+    "xarray",
+    "fsspec",
+    "netcdf4",
+    "scipy",  # restart capacities only
+    "h5netcdf",  # for xarray
+    "dask",  # for xarray
+]
+
+
 setup(
     author="NOAA/NASA",
     python_requires=">=3.8",
@@ -27,15 +42,11 @@ def local_pkg(name: str, relative_path: str) -> str:
         "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: 3.9",
     ],
-    install_requires=[
-        local_pkg("gt4py", "external/gt4py"),
-        local_pkg("dace", "external/dace"),
-    ],
-    setup_requires=setup_requirements,
-    tests_require=test_requirements,
+    install_requires=requirements,
+    extras_requires=extras_requires,
     name="ndsl",
     license="BSD license",
-    packages=find_namespace_packages(include=["ndsl.*"]),
+    packages=find_namespace_packages(include=["ndsl", "ndsl.*"]),
     include_package_data=True,
     url="https://github.com/NOAA-GFDL/NDSL",
     version="0.0.0",

tests/checkpointer/test_snapshot.py

@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 
-import ndsl.util
+from ndsl.checkpointer import SnapshotCheckpointer
 from ndsl.optional_imports import xarray as xr
 
 
@@ -10,13 +10,13 @@
 
 @requires_xarray
 def test_snapshot_checkpointer_no_data():
-    checkpointer = ndsl.SnapshotCheckpointer(rank=0)
+    checkpointer = SnapshotCheckpointer(rank=0)
     xr.testing.assert_identical(checkpointer.dataset, xr.Dataset())
 
 
 @requires_xarray
 def test_snapshot_checkpointer_one_snapshot():
-    checkpointer = ndsl.SnapshotCheckpointer(rank=0)
+    checkpointer = SnapshotCheckpointer(rank=0)
     val1 = np.random.randn(2, 3, 4)
     checkpointer("savepoint_name", val1=val1)
     xr.testing.assert_identical(
@@ -35,7 +35,7 @@ def test_snapshot_checkpointer_one_snapshot():
 
 @requires_xarray
 def test_snapshot_checkpointer_multiple_snapshots():
-    checkpointer = ndsl.SnapshotCheckpointer(rank=0)
+    checkpointer = SnapshotCheckpointer(rank=0)
     val1 = np.random.randn(2, 2, 3, 4)
     val2 = np.random.randn(1, 3, 2, 4)
     checkpointer("savepoint_name_1", val1=val1[0, :])

tests/checkpointer/test_thresholds.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
-from Checkpointer import (
+
+from ndsl.checkpointer import (
     InsufficientTrialsError,
     Threshold,
     ThresholdCalibrationCheckpointer,

tests/checkpointer/test_validation.py

@@ -3,9 +3,9 @@
 
 import numpy as np
 import pytest
-from Checkpointer import SavepointThresholds, Threshold, ValidationCheckpointer
-from Checkpointer.validation import _clip_pace_array_to_target
 
+from ndsl.checkpointer import SavepointThresholds, Threshold, ValidationCheckpointer
+from ndsl.checkpointer.validation import _clip_pace_array_to_target
 from ndsl.optional_imports import xarray as xr
 
 

tests/conftest.py

@@ -12,7 +12,7 @@
     cupy = None
 
 
-@pytest.fixture(params=["numpy", "cupy", "gt4py_numpy", "gt4py_cupy"])
+@pytest.fixture(params=["numpy", "cupy"])
 def backend(request):
     if cupy is None and request.param.endswith("cupy"):
         if request.config.getoption("--gpu-only"):
@@ -29,9 +29,9 @@ def backend(request):
 
 @pytest.fixture
 def gt4py_backend(backend):
-    if backend in ("numpy", "gt4py_numpy"):
+    if backend in ("numpy"):
         return "numpy"
-    elif backend in ("cupy", "gt4py_cupy"):
+    elif backend in ("cupy"):
         return "gt:gpu"
     else:
         return None
@@ -48,11 +48,6 @@ def numpy(backend):
         return np
     elif backend == "cupy":
         return cupy
-    elif backend == "gt4py_numpy":
-        # TODO: Deprecate these "backends".
-        return np
-    elif backend == "gt4py_cupy":
-        return cupy
     else:
         raise NotImplementedError()
 

tests/dsl/test_caches.py

@@ -140,7 +140,7 @@ def test_relocatability(backend: str):
     import gt4py
     from gt4py.cartesian import config as gt_config
 
-    from ndsl.comm.mpi import MPI
+    from ..mpi.mpi_comm import MPI
 
     # Restore original dir name
     gt4py.cartesian.config.cache_settings["dir_name"] = os.environ.get(