From 8400c83b55fe044fd7285f3b80407c584c127bd7 Mon Sep 17 00:00:00 2001 From: Frank Malatino Date: Wed, 28 Feb 2024 16:00:39 -0500 Subject: [PATCH] Updated import method to mimic other package styles --- ndsl/__init__.py | 121 ++++++++++++++++++------ ndsl/checkpointer/__init__.py | 3 - ndsl/comm/__init__.py | 7 -- ndsl/dsl/__init__.py | 11 --- ndsl/dsl/dace/__init__.py | 3 - ndsl/dsl/stencil.py | 6 +- ndsl/grid/__init__.py | 13 --- ndsl/grid/helper.py | 2 +- ndsl/halo/__init__.py | 2 - ndsl/initialization/__init__.py | 2 - ndsl/monitor/__init__.py | 2 - ndsl/performance/__init__.py | 2 - ndsl/stencils/__init__.py | 13 --- ndsl/stencils/testing/__init__.py | 17 ---- ndsl/stencils/testing/test_translate.py | 5 +- ndsl/testing/__init__.py | 3 - ndsl/testing/dummy_comm.py | 1 - setup.py | 1 + tests/checkpointer/test_thresholds.py | 6 +- tests/checkpointer/test_validation.py | 7 +- tests/dsl/test_caches.py | 2 +- tests/mpi/test_mpi_mock.py | 2 +- tests/test_halo_data_transformer.py | 9 +- 23 files changed, 113 insertions(+), 127 deletions(-) diff --git a/ndsl/__init__.py b/ndsl/__init__.py index 221cf1bb..6967d8dc 100644 --- a/ndsl/__init__.py +++ b/ndsl/__init__.py @@ -1,47 +1,110 @@ -from .checkpointer import Checkpointer, NullCheckpointer, SnapshotCheckpointer -from .comm import ( +from .buffer import Buffer +from .checkpointer.base import Checkpointer +from .checkpointer.null import NullCheckpointer +from .checkpointer.snapshots import SnapshotCheckpointer, _Snapshots +from .checkpointer.thresholds import ( + InsufficientTrialsError, + SavepointThresholds, + Threshold, + ThresholdCalibrationCheckpointer, +) +from .checkpointer.validation import ValidationCheckpointer +from .comm.boundary import Boundary, SimpleBoundary +from .comm.caching_comm import ( + CachingCommData, CachingCommReader, CachingCommWriter, - Comm, - Communicator, - ConcurrencyError, - CubedSphereCommunicator, - CubedSpherePartitioner, - LocalComm, - MPIComm, - NullComm, - TileCommunicator, - TilePartitioner, -) -from .dsl import ( + CachingRequestReader, + CachingRequestWriter, + NullRequest, +) +from .comm.comm_abc import Comm, Request +from .comm.communicator import Communicator, CubedSphereCommunicator, TileCommunicator +from .comm.local_comm import AsyncResult, ConcurrencyError, LocalComm +from .comm.mpi import MPIComm +from .comm.null_comm import NullAsyncResult, NullComm +from .comm.partitioner import CubedSpherePartitioner, Partitioner, 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 ( CompareToNumpyStencil, - CompilationConfig, - DaceConfig, - DaCeOrchestration, FrozenStencil, GridIndexing, - RunMode, - StencilConfig, StencilFactory, - WrappedHaloUpdater, + TimingCollector, ) +from .dsl.stencil_config import CompilationConfig, RunMode, StencilConfig from .exceptions import OutOfBoundsError -from .halo import HaloDataTransformer, HaloExchangeSpec, HaloUpdater -from .initialization import GridSizer, QuantityFactory, SubtileGridSizer +from .grid.eta import HybridPressureCoefficients +from .grid.generation import GridDefinition, GridDefinitions, MetricTerms +from .grid.helper import ( + AngleGridData, + ContravariantGridData, + DampingCoefficients, + DriverGridData, + GridData, + HorizontalGridData, + VerticalGridData, +) +from .halo.data_transformer import ( + HaloDataTransformer, + HaloDataTransformerCPU, + HaloDataTransformerGPU, + HaloExchangeSpec, +) +from .halo.updater import HaloUpdater, HaloUpdateRequest, VectorInterfaceHaloUpdater +from .initialization.allocator import QuantityFactory, StorageNumpy +from .initialization.sizer import GridSizer, SubtileGridSizer from .logging import ndsl_log -from .monitor import NetCDFMonitor, ZarrMonitor -from .performance import NullTimer, PerformanceCollector, Timer -from .quantity import Quantity, QuantityHaloSpec -from .stencils import ( - CubedToLatLon, - Grid, +from .monitor.netcdf_monitor import NetCDFMonitor +from .monitor.protocol import Protocol +from .monitor.zarr_monitor import ZarrMonitor +from .namelist import Namelist +from .optional_imports import RaiseWhenAccessed +from .performance.collector import ( + AbstractPerformanceCollector, + NullPerformanceCollector, + PerformanceCollector, +) +from .performance.config import PerformanceConfig +from .performance.profiler import NullProfiler, Profiler +from .performance.report import Experiment, Report, TimeReport +from .performance.timer import NullTimer, Timer +from .quantity import ( + BoundaryArrayView, + BoundedArrayView, + Quantity, + QuantityHaloSpec, + QuantityMetadata, +) +from .stencils.c2l_ord import CubedToLatLon +from .stencils.corners import CopyCorners, CopyCornersXY, FillCornersBGrid +from .stencils.testing.grid import Grid # type: ignore +from .stencils.testing.parallel_translate import ( ParallelTranslate, ParallelTranslate2Py, ParallelTranslate2PyState, ParallelTranslateBaseSlicing, ParallelTranslateGrid, +) +from .stencils.testing.savepoint import SavepointCase, Translate, dataset_to_dict +from .stencils.testing.temporaries import assert_same_temporaries, copy_temporaries +from .stencils.testing.translate import ( TranslateFortranData2Py, TranslateGrid, + pad_field_in_j, + read_serialized_data, ) -from .testing import DummyComm +from .testing.dummy_comm import DummyComm +from .types import Allocator, AsyncRequest, NumpyModule +from .units import UnitsError from .utils import MetaEnumStr diff --git a/ndsl/checkpointer/__init__.py b/ndsl/checkpointer/__init__.py index d24936cb..e69de29b 100644 --- a/ndsl/checkpointer/__init__.py +++ b/ndsl/checkpointer/__init__.py @@ -1,3 +0,0 @@ -from .base import Checkpointer -from .null import NullCheckpointer -from .snapshots import SnapshotCheckpointer diff --git a/ndsl/comm/__init__.py b/ndsl/comm/__init__.py index 31319c78..e69de29b 100644 --- a/ndsl/comm/__init__.py +++ b/ndsl/comm/__init__.py @@ -1,7 +0,0 @@ -from .caching_comm import CachingCommReader, CachingCommWriter -from .comm_abc import Comm -from .communicator import Communicator, CubedSphereCommunicator, TileCommunicator -from .local_comm import ConcurrencyError, LocalComm -from .mpi import MPIComm -from .null_comm import NullComm -from .partitioner import CubedSpherePartitioner, TilePartitioner diff --git a/ndsl/dsl/__init__.py b/ndsl/dsl/__init__.py index 269ae957..ed44420a 100644 --- a/ndsl/dsl/__init__.py +++ b/ndsl/dsl/__init__.py @@ -2,17 +2,6 @@ from ndsl.comm.mpi import MPI -from . import dace -from .dace import ( - DaceConfig, - DaCeOrchestration, - WrappedHaloUpdater, - orchestrate, - orchestrate_function, -) -from .stencil import CompareToNumpyStencil, FrozenStencil, GridIndexing, StencilFactory -from .stencil_config import CompilationConfig, RunMode, StencilConfig - if MPI is not None: import os diff --git a/ndsl/dsl/dace/__init__.py b/ndsl/dsl/dace/__init__.py index c1386ad9..e69de29b 100644 --- a/ndsl/dsl/dace/__init__.py +++ b/ndsl/dsl/dace/__init__.py @@ -1,3 +0,0 @@ -from .dace_config import DaceConfig, DaCeOrchestration -from .orchestration import orchestrate, orchestrate_function -from .wrapped_halo_exchange import WrappedHaloUpdater diff --git a/ndsl/dsl/stencil.py b/ndsl/dsl/stencil.py index 77efd672..b8316727 100644 --- a/ndsl/dsl/stencil.py +++ b/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 @@ -34,6 +33,9 @@ from ndsl.initialization.sizer import GridSizer, SubtileGridSizer from ndsl.quantity import Quantity +# from ndsl import testing +from ndsl.testing import comparison + try: import cupy as cp @@ -68,7 +70,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( diff --git a/ndsl/grid/__init__.py b/ndsl/grid/__init__.py index a7692a8f..e69de29b 100644 --- a/ndsl/grid/__init__.py +++ b/ndsl/grid/__init__.py @@ -1,13 +0,0 @@ -# flake8: noqa: F401 - -from .eta import set_hybrid_pressure_coefficients -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 .stretch_transformation import direct_transform diff --git a/ndsl/grid/helper.py b/ndsl/grid/helper.py index ee97a6b0..fd62d771 100644 --- a/ndsl/grid/helper.py +++ b/ndsl/grid/helper.py @@ -14,7 +14,7 @@ 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.initialization.allocator import QuantityFactory from ndsl.quantity import Quantity diff --git a/ndsl/halo/__init__.py b/ndsl/halo/__init__.py index 823bd226..e69de29b 100644 --- a/ndsl/halo/__init__.py +++ b/ndsl/halo/__init__.py @@ -1,2 +0,0 @@ -from .data_transformer import HaloDataTransformer, HaloExchangeSpec -from .updater import HaloUpdater diff --git a/ndsl/initialization/__init__.py b/ndsl/initialization/__init__.py index fe15db8b..e69de29b 100644 --- a/ndsl/initialization/__init__.py +++ b/ndsl/initialization/__init__.py @@ -1,2 +0,0 @@ -from .allocator import QuantityFactory -from .sizer import GridSizer, SubtileGridSizer diff --git a/ndsl/monitor/__init__.py b/ndsl/monitor/__init__.py index 26b38cc6..e69de29b 100644 --- a/ndsl/monitor/__init__.py +++ b/ndsl/monitor/__init__.py @@ -1,2 +0,0 @@ -from .netcdf_monitor import NetCDFMonitor -from .zarr_monitor import ZarrMonitor diff --git a/ndsl/performance/__init__.py b/ndsl/performance/__init__.py index 51e9bf84..e69de29b 100644 --- a/ndsl/performance/__init__.py +++ b/ndsl/performance/__init__.py @@ -1,2 +0,0 @@ -from .collector import PerformanceCollector -from .timer import NullTimer, Timer diff --git a/ndsl/stencils/__init__.py b/ndsl/stencils/__init__.py index 6083272b..d3ec452c 100644 --- a/ndsl/stencils/__init__.py +++ b/ndsl/stencils/__init__.py @@ -1,14 +1 @@ -from .c2l_ord import CubedToLatLon -from .testing import ( - Grid, - ParallelTranslate, - ParallelTranslate2Py, - ParallelTranslate2PyState, - ParallelTranslateBaseSlicing, - ParallelTranslateGrid, - TranslateFortranData2Py, - TranslateGrid, -) - - __version__ = "0.2.0" diff --git a/ndsl/stencils/testing/__init__.py b/ndsl/stencils/testing/__init__.py index 3ad9ef9c..e69de29b 100644 --- a/ndsl/stencils/testing/__init__.py +++ b/ndsl/stencils/testing/__init__.py @@ -1,17 +0,0 @@ -from . import parallel_translate, translate -from .grid import Grid # type: ignore -from .parallel_translate import ( - ParallelTranslate, - ParallelTranslate2Py, - ParallelTranslate2PyState, - ParallelTranslateBaseSlicing, - ParallelTranslateGrid, -) -from .savepoint import dataset_to_dict -from .temporaries import assert_same_temporaries, copy_temporaries -from .translate import ( - TranslateFortranData2Py, - TranslateGrid, - pad_field_in_j, - read_serialized_data, -) diff --git a/ndsl/stencils/testing/test_translate.py b/ndsl/stencils/testing/test_translate.py index 2e42e273..29c4ed65 100644 --- a/ndsl/stencils/testing/test_translate.py +++ b/ndsl/stencils/testing/test_translate.py @@ -14,8 +14,9 @@ from ndsl.dsl.stencil import CompilationConfig, StencilConfig from ndsl.quantity import Quantity from ndsl.restart._legacy_restart import RESTART_PROPERTIES -from ndsl.stencils.testing import SavepointCase, dataset_to_dict -from ndsl.testing import compare_scalar, perturb, success, success_array +from ndsl.stencils.testing.savepoint import SavepointCase, dataset_to_dict +from ndsl.testing.comparison import compare_scalar, success, success_array +from ndsl.testing.perturbation import perturb # this only matters for manually-added print statements diff --git a/ndsl/testing/__init__.py b/ndsl/testing/__init__.py index a1c927e9..e69de29b 100644 --- a/ndsl/testing/__init__.py +++ b/ndsl/testing/__init__.py @@ -1,3 +0,0 @@ -from .comparison import compare_arr, compare_scalar, success, success_array -from .dummy_comm import ConcurrencyError, DummyComm -from .perturbation import perturb diff --git a/ndsl/testing/dummy_comm.py b/ndsl/testing/dummy_comm.py index b4df2347..f3e93817 100644 --- a/ndsl/testing/dummy_comm.py +++ b/ndsl/testing/dummy_comm.py @@ -1,2 +1 @@ -from ndsl.comm.local_comm import ConcurrencyError # noqa from ndsl.comm.local_comm import LocalComm as DummyComm # noqa diff --git a/setup.py b/setup.py index 73ec210a..c0d71813 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ def local_pkg(name: str, relative_path: str) -> str: "mpi4py", "cftime", "xarray", + "f90nml>=1.1.0", "fsspec", "netcdf4", "scipy", # restart capacities only diff --git a/tests/checkpointer/test_thresholds.py b/tests/checkpointer/test_thresholds.py index 851a9665..90d1f8fc 100644 --- a/tests/checkpointer/test_thresholds.py +++ b/tests/checkpointer/test_thresholds.py @@ -1,11 +1,7 @@ import numpy as np import pytest -from ndsl.checkpointer.thresholds import ( - InsufficientTrialsError, - Threshold, - ThresholdCalibrationCheckpointer, -) +from ndsl import InsufficientTrialsError, Threshold, ThresholdCalibrationCheckpointer def test_thresholds_no_trials(): diff --git a/tests/checkpointer/test_validation.py b/tests/checkpointer/test_validation.py index b696aca3..091bb7c6 100644 --- a/tests/checkpointer/test_validation.py +++ b/tests/checkpointer/test_validation.py @@ -4,11 +4,8 @@ import numpy as np import pytest -from ndsl.checkpointer.thresholds import SavepointThresholds, Threshold -from ndsl.checkpointer.validation import ( - ValidationCheckpointer, - _clip_pace_array_to_target, -) +from ndsl import SavepointThresholds, Threshold, ValidationCheckpointer +from ndsl.checkpointer.validation import _clip_pace_array_to_target from ndsl.optional_imports import xarray as xr diff --git a/tests/dsl/test_caches.py b/tests/dsl/test_caches.py index f90a1828..893fb89d 100644 --- a/tests/dsl/test_caches.py +++ b/tests/dsl/test_caches.py @@ -11,7 +11,7 @@ StencilFactory, ) from ndsl.comm.mpi import MPI -from ndsl.dsl.dace import orchestrate +from ndsl.dsl.dace.orchestration import orchestrate def _make_storage( diff --git a/tests/mpi/test_mpi_mock.py b/tests/mpi/test_mpi_mock.py index d099d76a..def0d342 100644 --- a/tests/mpi/test_mpi_mock.py +++ b/tests/mpi/test_mpi_mock.py @@ -1,8 +1,8 @@ import numpy as np import pytest +from ndsl import ConcurrencyError, DummyComm from ndsl.comm.communicator import recv_buffer -from ndsl.testing import ConcurrencyError, DummyComm from tests.mpi.mpi_comm import MPI diff --git a/tests/test_halo_data_transformer.py b/tests/test_halo_data_transformer.py index 7e1b9f09..ec986f8c 100644 --- a/tests/test_halo_data_transformer.py +++ b/tests/test_halo_data_transformer.py @@ -4,8 +4,13 @@ import numpy as np import pytest -from ndsl import HaloDataTransformer, HaloExchangeSpec, Quantity, QuantityHaloSpec -from ndsl.buffer import Buffer +from ndsl import ( + Buffer, + HaloDataTransformer, + HaloExchangeSpec, + Quantity, + QuantityHaloSpec, +) from ndsl.comm import _boundary_utils from ndsl.constants import ( EAST,