Adding and updating __init__ files for modules

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,31 @@
+from .checkpointer import SnapshotCheckpointer
+from .comm import (
+    CachingCommReader,
+    CachingCommWriter,
+    ConcurrencyError,
+    CubedSphereCommunicator,
+    CubedSpherePartitioner,
+    LocalComm,
+    NullComm,
+    TileCommunicator,
+    TilePartitioner,
+)
+from .dsl import (
+    CompareToNumpyStencil,
+    CompilationConfig,
+    DaceConfig,
+    DaCeOrchestration,
+    FrozenStencil,
+    GridIndexing,
+    RunMode,
+    StencilConfig,
+    StencilFactory,
+)
+from .exceptions import OutOfBoundsError
+from .halo import HaloDataTransformer, HaloExchangeSpec, HaloUpdater
+from .initialization import QuantityFactory, SubtileGridSizer
 from .logging import ndsl_log
+from .monitor import NetCDFMonitor, ZarrMonitor
+from .performance import NullTimer, Timer
+from .quantity import Quantity, QuantityHaloSpec
+from .testing import DummyComm

ndsl/checkpointer/__init__.py

@@ -1,10 +1 @@
-from .base import Checkpointer
-from .null import NullCheckpointer
 from .snapshots import SnapshotCheckpointer
-from .thresholds import (
-    InsufficientTrialsError,
-    SavepointThresholds,
-    Threshold,
-    ThresholdCalibrationCheckpointer,
-)
-from .validation import ValidationCheckpointer

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,5 @@
+from .caching_comm import CachingCommReader, CachingCommWriter
+from .communicator import CubedSphereCommunicator, TileCommunicator
+from .local_comm import ConcurrencyError, LocalComm
+from .null_comm import NullComm
+from .partitioner import CubedSpherePartitioner, TilePartitioner

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/dsl/__init__.py

@@ -3,15 +3,9 @@
 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,
-)
+from .dace import DaceConfig, DaCeOrchestration, orchestrate, orchestrate_function
+from .stencil import CompareToNumpyStencil, FrozenStencil, GridIndexing, StencilFactory
+from .stencil_config import CompilationConfig, RunMode, StencilConfig
 
 
 if MPI is not None:

ndsl/dsl/caches/__init__.py

@@ -0,0 +1 @@
+from .codepath import FV3CodePath

ndsl/dsl/dace/__init__.py

@@ -1,2 +1,9 @@
-from ndsl.dsl.dace.dace_config import DaceConfig
-from ndsl.dsl.dace.orchestration import orchestrate
+from .dace_config import DaceConfig, DaCeOrchestration, FrozenCompiledSDFG
+from .orchestration import (
+    _LazyComputepathFunction,
+    _LazyComputepathMethod,
+    orchestrate,
+    orchestrate_function,
+)
+from .utils import ArrayReport, DaCeProgress, MaxBandwithBenchmarkProgram, StorageReport
+from .wrapped_halo_exchange import WrappedHaloUpdater

ndsl/dsl/stencil.py

@@ -31,7 +31,7 @@
 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
 
 

ndsl/grid/__init__.py

@@ -1,7 +1,6 @@
 # flake8: noqa: F401
 
 from .eta import set_hybrid_pressure_coefficients
-from .generation import GridDefinitions, MetricTerms
 from .gnomonic import (
     great_circle_distance_along_axis,
     great_circle_distance_lon_lat,
@@ -11,13 +10,4 @@
     xyz_midpoint,
     xyz_to_lon_lat,
 )
-from .helper import (
-    AngleGridData,
-    ContravariantGridData,
-    DampingCoefficients,
-    DriverGridData,
-    GridData,
-    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.grid.generation import MetricTerms
 from ndsl.initialization import QuantityFactory
 from ndsl.quantity import Quantity
 
-from .generation import MetricTerms
-
 
 @dataclasses.dataclass(frozen=True)
 class DampingCoefficients:

ndsl/halo/__init__.py

@@ -0,0 +1,2 @@
+from .data_transformer import HaloDataTransformer, HaloExchangeSpec
+from .updater import HaloUpdater

ndsl/initialization/__init__.py

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

ndsl/initialization/allocator.py

@@ -2,10 +2,10 @@
 
 import numpy as np
 
-from ..constants import SPATIAL_DIMS
-from ..optional_imports import gt4py
-from ..quantity import Quantity, QuantityHaloSpec
-from .sizer import GridSizer
+from ndsl.constants import SPATIAL_DIMS
+from ndsl.initialization.sizer import GridSizer
+from ndsl.optional_imports import gt4py
+from ndsl.quantity import Quantity, QuantityHaloSpec
 
 
 class StorageNumpy:

ndsl/monitor/__init__.py

@@ -1,3 +1,2 @@
 from .netcdf_monitor import NetCDFMonitor
-from .protocol import Monitor
 from .zarr_monitor import ZarrMonitor

ndsl/monitor/convert.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from ..optional_imports import cupy
+from ndsl.optional_imports import cupy
 
 
 def to_numpy(array, dtype=None) -> np.ndarray:

ndsl/monitor/netcdf_monitor.py

@@ -6,12 +6,11 @@
 import numpy as np
 
 from ndsl.comm.communicator import Communicator
+from ndsl.filesystem import get_fs
+from ndsl.logging import ndsl_log
+from ndsl.monitor.convert import to_numpy
 from ndsl.optional_imports import xarray as xr
-
-from ..filesystem import get_fs
-from ..logging import ndsl_log
-from ..quantity import Quantity
-from .convert import to_numpy
+from ndsl.quantity import Quantity
 
 
 class _TimeChunkedVariable:

ndsl/performance/__init__.py

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

ndsl/performance/collector.py

@@ -11,6 +11,7 @@
 from ndsl.performance.report import (
     Report,
     TimeReport,
+    collect_data_and_write_to_file,
     collect_keys_from_data,
     gather_hit_counts,
     get_experiment_info,
@@ -19,8 +20,6 @@
 from ndsl.performance.timer import NullTimer, Timer
 from ndsl.utils import GPU_AVAILABLE
 
-from .report import collect_data_and_write_to_file
-
 
 class AbstractPerformanceCollector(Protocol):
     total_timer: Timer

ndsl/performance/config.py

@@ -1,13 +1,12 @@
 import dataclasses
 
 from ndsl.comm.comm_abc import Comm
-from ndsl.performance.profiler import NullProfiler, Profiler
-
-from .collector import (
+from ndsl.performance.collector import (
     AbstractPerformanceCollector,
     NullPerformanceCollector,
     PerformanceCollector,
 )
+from ndsl.performance.profiler import NullProfiler, Profiler
 
 
 @dataclasses.dataclass

ndsl/stencils/c2l_ord.py

@@ -13,7 +13,7 @@
 from ndsl.dsl.dace.wrapped_halo_exchange import WrappedHaloUpdater
 from ndsl.dsl.stencil import StencilFactory
 from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ
-from ndsl.grid import GridData
+from ndsl.grid.helper import GridData
 from ndsl.initialization.allocator import QuantityFactory