Merge branch 'develop' into feature/stencil_basic_ops

examples/mpi/global_timings.py

@@ -3,7 +3,7 @@
 import numpy as np
 from mpi4py import MPI
 
-from ndsl.performance.timer import Timer
+from ndsl import Timer
 
 
 @contextlib.contextmanager

examples/mpi/zarr_monitor.py

@@ -5,11 +5,14 @@
 import zarr
 from mpi4py import MPI
 
-from ndsl.comm.partitioner import CubedSpherePartitioner, TilePartitioner
+from ndsl import (
+    CubedSpherePartitioner,
+    QuantityFactory,
+    SubtileGridSizer,
+    TilePartitioner,
+    ZarrMonitor,
+)
 from ndsl.constants import X_DIM, Y_DIM, Z_DIM
-from ndsl.initialization.allocator import QuantityFactory
-from ndsl.initialization.sizer import SubtileGridSizer
-from ndsl.monitor import ZarrMonitor
 
 
 OUTPUT_PATH = "output/zarr_monitor.zarr"

ndsl/__init__.py

@@ -1 +1,33 @@
+from .comm.communicator import CubedSphereCommunicator, TileCommunicator
+from .comm.local_comm import LocalComm
+from .comm.mpi import MPIComm
+from .comm.null_comm import NullComm
+from .comm.partitioner import CubedSpherePartitioner, TilePartitioner
+from .constants import ConstantVersions
+from .dsl.caches.codepath import FV3CodePath
+from .dsl.dace.dace_config import DaceConfig, DaCeOrchestration, FrozenCompiledSDFG
+from .dsl.dace.orchestration import orchestrate, orchestrate_function
+from .dsl.dace.utils import (
+    ArrayReport,
+    DaCeProgress,
+    MaxBandwithBenchmarkProgram,
+    StorageReport,
+)
+from .dsl.dace.wrapped_halo_exchange import WrappedHaloUpdater
+from .dsl.stencil import FrozenStencil, GridIndexing, StencilFactory, TimingCollector
+from .dsl.stencil_config import CompilationConfig, RunMode, StencilConfig
+from .exceptions import OutOfBoundsError
+from .halo.data_transformer import HaloExchangeSpec
+from .halo.updater import HaloUpdater, HaloUpdateRequest, VectorInterfaceHaloUpdater
+from .initialization.allocator import QuantityFactory
+from .initialization.sizer import GridSizer, SubtileGridSizer
 from .logging import ndsl_log
+from .monitor.netcdf_monitor import NetCDFMonitor
+from .namelist import Namelist
+from .performance.collector import NullPerformanceCollector, PerformanceCollector
+from .performance.profiler import NullProfiler, Profiler
+from .performance.report import Experiment, Report, TimeReport
+from .quantity import Quantity
+from .testing.dummy_comm import DummyComm
+from .types import Allocator
+from .utils import MetaEnumStr

ndsl/checkpointer/__init__.py

@@ -1,4 +1,3 @@
-from .base import Checkpointer
 from .null import NullCheckpointer
 from .snapshots import SnapshotCheckpointer
 from .thresholds import (

ndsl/checkpointer/null.py

@@ -1,4 +1,4 @@
-from .base import Checkpointer
+from ndsl.checkpointer.base import Checkpointer
 
 
 class NullCheckpointer(Checkpointer):

ndsl/checkpointer/snapshots.py

@@ -2,11 +2,10 @@
 
 import numpy as np
 
+from ndsl.checkpointer.base import Checkpointer
 from ndsl.optional_imports import cupy as cp
 from ndsl.optional_imports import xarray as xr
 
-from .base import Checkpointer
-
 
 def make_dims(savepoint_dim, label, data_list):
     """

ndsl/checkpointer/thresholds.py

@@ -5,8 +5,8 @@
 
 import numpy as np
 
-from ..quantity import Quantity
-from .base import Checkpointer
+from ndsl.checkpointer.base import Checkpointer
+from ndsl.quantity import Quantity
 
 
 try:

ndsl/checkpointer/validation.py

@@ -5,11 +5,15 @@
 
 import numpy as np
 
+from ndsl.checkpointer.base import Checkpointer
+from ndsl.checkpointer.thresholds import (
+    ArrayLike,
+    SavepointName,
+    SavepointThresholds,
+    cast_to_ndarray,
+)
 from ndsl.optional_imports import xarray as xr
 
-from .base import Checkpointer
-from .thresholds import ArrayLike, SavepointName, SavepointThresholds, cast_to_ndarray
-
 
 def _clip_pace_array_to_target(
     array: np.ndarray, target_shape: Tuple[int, ...]

ndsl/comm/__init__.py

@@ -0,0 +1,8 @@
+from .caching_comm import (
+    CachingCommData,
+    CachingCommReader,
+    CachingCommWriter,
+    CachingRequestReader,
+    CachingRequestWriter,
+)
+from .comm_abc import Comm, Request

ndsl/comm/boundary.py

@@ -1,8 +1,8 @@
 import dataclasses
 from typing import Tuple
 
-from ..quantity import Quantity, QuantityHaloSpec
-from ._boundary_utils import get_boundary_slice
+from ndsl.comm._boundary_utils import get_boundary_slice
+from ndsl.quantity import Quantity, QuantityHaloSpec
 
 
 @dataclasses.dataclass

ndsl/comm/communicator.py

@@ -4,14 +4,13 @@
 import numpy as np
 
 import ndsl.constants as constants
-from ndsl.buffer import array_buffer, recv_buffer, send_buffer
+from ndsl.buffer import array_buffer, device_synchronize, recv_buffer, send_buffer
 from ndsl.comm.boundary import Boundary
 from ndsl.comm.partitioner import CubedSpherePartitioner, Partitioner, TilePartitioner
 from ndsl.halo.updater import HaloUpdater, HaloUpdateRequest, VectorInterfaceHaloUpdater
 from ndsl.performance.timer import NullTimer, Timer
 from ndsl.quantity import Quantity, QuantityHaloSpec, QuantityMetadata
 from ndsl.types import NumpyModule
-from ndsl.utils import device_synchronize
 
 
 try:
@@ -43,21 +42,6 @@ def to_numpy(array, dtype=None) -> np.ndarray:
     return output
 
 
-def bcast_metadata_list(comm, quantity_list):
-    is_root = comm.Get_rank() == constants.ROOT_RANK
-    if is_root:
-        metadata_list = []
-        for quantity in quantity_list:
-            metadata_list.append(quantity.metadata)
-    else:
-        metadata_list = None
-    return comm.bcast(metadata_list, root=constants.ROOT_RANK)
-
-
-def bcast_metadata(comm, array):
-    return bcast_metadata_list(comm, [array])[0]
-
-
 class Communicator(abc.ABC):
     def __init__(
         self, comm, partitioner, force_cpu: bool = False, timer: Optional[Timer] = None
@@ -585,6 +569,21 @@ def boundaries(self) -> Mapping[int, Boundary]:
         return self._boundaries
 
 
+def bcast_metadata_list(comm, quantity_list):
+    is_root = comm.Get_rank() == constants.ROOT_RANK
+    if is_root:
+        metadata_list = []
+        for quantity in quantity_list:
+            metadata_list.append(quantity.metadata)
+    else:
+        metadata_list = None
+    return comm.bcast(metadata_list, root=constants.ROOT_RANK)
+
+
+def bcast_metadata(comm, array):
+    return bcast_metadata_list(comm, [array])[0]
+
+
 class TileCommunicator(Communicator):
     """Performs communications within a single tile or region of a tile"""
 

ndsl/dsl/__init__.py

@@ -2,17 +2,6 @@
 
 from ndsl.comm.mpi import MPI
 
-from . import dace
-from .dace.dace_config import DaceConfig, DaCeOrchestration
-from .dace.orchestration import orchestrate, orchestrate_function
-from .stencil import (
-    CompilationConfig,
-    FrozenStencil,
-    GridIndexing,
-    StencilConfig,
-    StencilFactory,
-)
-
 
 if MPI is not None:
     import os

ndsl/dsl/dace/__init__.py

@@ -1,2 +0,0 @@
-from ndsl.dsl.dace.dace_config import DaceConfig
-from ndsl.dsl.dace.orchestration import orchestrate

ndsl/dsl/dace/dace_config.py

@@ -6,7 +6,8 @@
 from dace.codegen.compiled_sdfg import CompiledSDFG
 from dace.frontend.python.parser import DaceProgram
 
-from ndsl.comm.communicator import Communicator, Partitioner
+from ndsl.comm.communicator import Communicator
+from ndsl.comm.partitioner import Partitioner
 from ndsl.dsl.caches.cache_location import identify_code_path
 from ndsl.dsl.caches.codepath import FV3CodePath
 from ndsl.dsl.gt4py_utils import is_gpu_backend

ndsl/dsl/stencil.py

@@ -22,7 +22,6 @@
 from gt4py.cartesian import gtscript
 from gt4py.cartesian.gtc.passes.oir_pipeline import DefaultPipeline, OirPipeline
 
-from ndsl import testing
 from ndsl.comm.comm_abc import Comm
 from ndsl.comm.communicator import Communicator
 from ndsl.comm.decomposition import block_waiting_for_compilation, unblock_waiting_tiles
@@ -31,8 +30,9 @@
 from ndsl.dsl.dace.orchestration import SDFGConvertible
 from ndsl.dsl.stencil_config import CompilationConfig, RunMode, StencilConfig
 from ndsl.dsl.typing import Float, Index3D, cast_to_index3d
-from ndsl.initialization import GridSizer, SubtileGridSizer
+from ndsl.initialization.sizer import GridSizer, SubtileGridSizer
 from ndsl.quantity import Quantity
+from ndsl.testing import comparison
 
 
 try:
@@ -68,7 +68,7 @@ def report_difference(args, kwargs, args_copy, kwargs_copy, function_name, gt_id
 
 
 def report_diff(arg: np.ndarray, numpy_arg: np.ndarray, label) -> str:
-    metric_err = testing.compare_arr(arg, numpy_arg)
+    metric_err = comparison.compare_arr(arg, numpy_arg)
     nans_match = np.logical_and(np.isnan(arg), np.isnan(numpy_arg))
     n_points = np.product(arg.shape)
     failures_14 = n_points - np.sum(

ndsl/exceptions.py

@@ -1,2 +1,7 @@
+# flake8: noqa
+from ndsl.comm.local_comm import ConcurrencyError
+from ndsl.units import UnitsError
+
+
 class OutOfBoundsError(ValueError):
     pass

ndsl/grid/__init__.py

@@ -1,16 +1,5 @@
-# flake8: noqa: F401
-
-from .eta import set_hybrid_pressure_coefficients
+from .eta import HybridPressureCoefficients
 from .generation import GridDefinitions, MetricTerms
-from .gnomonic import (
-    great_circle_distance_along_axis,
-    great_circle_distance_lon_lat,
-    lon_lat_corner_to_cell_center,
-    lon_lat_midpoint,
-    lon_lat_to_xyz,
-    xyz_midpoint,
-    xyz_to_lon_lat,
-)
 from .helper import (
     AngleGridData,
     ContravariantGridData,
@@ -20,4 +9,3 @@
     HorizontalGridData,
     VerticalGridData,
 )
-from .stretch_transformation import direct_transform

ndsl/grid/generation.py

@@ -21,17 +21,7 @@
 from ndsl.dsl.stencil import GridIndexing
 from ndsl.dsl.typing import Float
 from ndsl.grid import eta
-from ndsl.initialization.allocator import QuantityFactory
-from ndsl.initialization.sizer import SubtileGridSizer
-from ndsl.quantity import Quantity
-from ndsl.stencils.corners import (
-    fill_corners_2d,
-    fill_corners_agrid,
-    fill_corners_cgrid,
-    fill_corners_dgrid,
-)
-
-from .geometry import (
+from ndsl.grid.geometry import (
     calc_unit_vector_south,
     calc_unit_vector_west,
     calculate_divg_del6,
@@ -47,7 +37,7 @@
     supergrid_corner_fix,
     unit_vector_lonlat,
 )
-from .gnomonic import (
+from ndsl.grid.gnomonic import (
     get_area,
     great_circle_distance_along_axis,
     local_gnomonic_ed,
@@ -59,7 +49,16 @@
     set_tile_border_dxc,
     set_tile_border_dyc,
 )
-from .mirror import mirror_grid
+from ndsl.grid.mirror import mirror_grid
+from ndsl.initialization.allocator import QuantityFactory
+from ndsl.initialization.sizer import SubtileGridSizer
+from ndsl.quantity import Quantity
+from ndsl.stencils.corners import (
+    fill_corners_2d,
+    fill_corners_agrid,
+    fill_corners_cgrid,
+    fill_corners_dgrid,
+)
 
 
 # TODO: when every environment in python3.8, remove

ndsl/grid/geometry.py

@@ -1,7 +1,5 @@
 from ndsl.comm.partitioner import TilePartitioner
-from ndsl.quantity import Quantity
-
-from .gnomonic import (
+from ndsl.grid.gnomonic import (
     get_lonlat_vect,
     get_unit_vector_direction,
     great_circle_distance_lon_lat,
@@ -10,6 +8,7 @@
     spherical_cos,
     xyz_midpoint,
 )
+from ndsl.quantity import Quantity
 
 
 def get_center_vector(

ndsl/grid/global_setup.py

@@ -1,16 +1,15 @@
 import math
 
 from ndsl.constants import PI, RADIUS
-
-from .generation import MetricTerms
-from .gnomonic import (
+from ndsl.grid.generation import MetricTerms
+from ndsl.grid.gnomonic import (
     _cart_to_latlon,
     _check_shapes,
     _latlon2xyz,
     _mirror_latlon,
     symm_ed,
 )
-from .mirror import _rot_3d
+from ndsl.grid.mirror import _rot_3d
 
 
 def gnomonic_grid(grid_type: int, lon, lat, np):

ndsl/grid/helper.py

@@ -13,11 +13,10 @@
 import ndsl.constants as constants
 from ndsl.constants import Z_DIM, Z_INTERFACE_DIM
 from ndsl.filesystem import get_fs
-from ndsl.initialization import QuantityFactory
+from ndsl.grid.generation import MetricTerms
+from ndsl.initialization.allocator import QuantityFactory
 from ndsl.quantity import Quantity
 
-from .generation import MetricTerms
-
 
 @dataclasses.dataclass(frozen=True)
 class DampingCoefficients:

ndsl/halo/__init__.py

@@ -0,0 +1,5 @@
+from .data_transformer import (
+    HaloDataTransformer,
+    HaloDataTransformerCPU,
+    HaloDataTransformerGPU,
+)

ndsl/initialization/__init__.py

@@ -1,2 +1 @@
-from .allocator import QuantityFactory
-from .sizer import GridSizer, SubtileGridSizer
+from .sizer import GridSizer