Optional and non-optional dependencies

ndsl/checkpointer/snapshots.py

@@ -1,10 +1,10 @@
 import collections
 
 import numpy as np
+import xarray as xr
 
 from ndsl.checkpointer.base import Checkpointer
 from ndsl.optional_imports import cupy as cp
-from ndsl.optional_imports import xarray as xr
 
 
 def make_dims(savepoint_dim, label, data_list):
@@ -39,12 +39,7 @@ def dataset(self) -> "xr.Dataset":
             data_vars[f"{variable_name}"] = make_dims(
                 savepoint_dim, variable_name, self._arrays[variable_name]
             )
-        if xr is None:
-            raise ModuleNotFoundError(
-                "xarray must be installed to use Snapshots.dataset"
-            )
-        else:
-            return xr.Dataset(data_vars=data_vars)
+        return xr.Dataset(data_vars=data_vars)
 
 
 class SnapshotCheckpointer(Checkpointer):
@@ -54,10 +49,6 @@ class SnapshotCheckpointer(Checkpointer):
     """
 
     def __init__(self, rank: int):
-        if xr is None:
-            raise ModuleNotFoundError(
-                "xarray must be installed to use SnapshotCheckpointer"
-            )
         self._rank = rank
         self._snapshots = _Snapshots()
 

ndsl/checkpointer/thresholds.py

@@ -9,12 +9,6 @@
 from ndsl.quantity import Quantity
 
 
-try:
-    import cupy as cp
-except ImportError:
-    cp = None
-
-
 SavepointName = str
 VariableName = str
 ArrayLike = Union[Quantity, np.ndarray]

ndsl/checkpointer/validation.py

@@ -4,6 +4,7 @@
 from typing import MutableMapping, Tuple
 
 import numpy as np
+import xarray as xr
 
 from ndsl.checkpointer.base import Checkpointer
 from ndsl.checkpointer.thresholds import (
@@ -12,7 +13,6 @@
     SavepointThresholds,
     cast_to_ndarray,
 )
-from ndsl.optional_imports import xarray as xr
 
 
 def _clip_pace_array_to_target(
@@ -109,8 +109,6 @@ def __call__(self, savepoint_name: str, **kwargs: ArrayLike) -> None:
         Raises:
             AssertionError: if the thresholds on any variable are not met
         """
-        if xr is None:
-            raise ModuleNotFoundError("xarray is not installed")
         nc_file = os.path.join(self._savepoint_data_path, savepoint_name + ".nc")
         ds = xr.open_dataset(nc_file)
 

ndsl/comm/communicator.py

@@ -10,17 +10,12 @@
 from ndsl.comm.comm_abc import ReductionOperator
 from ndsl.comm.partitioner import CubedSpherePartitioner, Partitioner, TilePartitioner
 from ndsl.halo.updater import HaloUpdater, HaloUpdateRequest, VectorInterfaceHaloUpdater
+from ndsl.optional_imports import cupy
 from ndsl.performance.timer import NullTimer, Timer
 from ndsl.quantity import Quantity, QuantityHaloSpec, QuantityMetadata
 from ndsl.types import NumpyModule
 
 
-try:
-    import cupy
-except ImportError:
-    cupy = None
-
-
 def to_numpy(array, dtype=None) -> np.ndarray:
     """
     Input array can be a numpy array or a cupy array. Returns numpy array.

ndsl/comm/mpi.py

@@ -1,33 +1,30 @@
-try:
-    import mpi4py
-    from mpi4py import MPI
-except ImportError:
-    MPI = None
 from typing import Dict, List, Optional, TypeVar, cast
 
+from mpi4py import MPI
+
 from ndsl.comm.comm_abc import Comm, ReductionOperator, Request
 
 
 T = TypeVar("T")
 
 
 class MPIComm(Comm):
-    _op_mapping: Dict[ReductionOperator, mpi4py.MPI.Op] = {
-        ReductionOperator.OP_NULL: mpi4py.MPI.OP_NULL,
-        ReductionOperator.MAX: mpi4py.MPI.MAX,
-        ReductionOperator.MIN: mpi4py.MPI.MIN,
-        ReductionOperator.SUM: mpi4py.MPI.SUM,
-        ReductionOperator.PROD: mpi4py.MPI.PROD,
-        ReductionOperator.LAND: mpi4py.MPI.LAND,
-        ReductionOperator.BAND: mpi4py.MPI.BAND,
-        ReductionOperator.LOR: mpi4py.MPI.LOR,
-        ReductionOperator.BOR: mpi4py.MPI.BOR,
-        ReductionOperator.LXOR: mpi4py.MPI.LXOR,
-        ReductionOperator.BXOR: mpi4py.MPI.BXOR,
-        ReductionOperator.MAXLOC: mpi4py.MPI.MAXLOC,
-        ReductionOperator.MINLOC: mpi4py.MPI.MINLOC,
-        ReductionOperator.REPLACE: mpi4py.MPI.REPLACE,
-        ReductionOperator.NO_OP: mpi4py.MPI.NO_OP,
+    _op_mapping: Dict[ReductionOperator, MPI.Op] = {
+        ReductionOperator.OP_NULL: MPI.OP_NULL,
+        ReductionOperator.MAX: MPI.MAX,
+        ReductionOperator.MIN: MPI.MIN,
+        ReductionOperator.SUM: MPI.SUM,
+        ReductionOperator.PROD: MPI.PROD,
+        ReductionOperator.LAND: MPI.LAND,
+        ReductionOperator.BAND: MPI.BAND,
+        ReductionOperator.LOR: MPI.LOR,
+        ReductionOperator.BOR: MPI.BOR,
+        ReductionOperator.LXOR: MPI.LXOR,
+        ReductionOperator.BXOR: MPI.BXOR,
+        ReductionOperator.MAXLOC: MPI.MAXLOC,
+        ReductionOperator.MINLOC: MPI.MINLOC,
+        ReductionOperator.REPLACE: MPI.REPLACE,
+        ReductionOperator.NO_OP: MPI.NO_OP,
     }
 
     def __init__(self):
@@ -84,4 +81,4 @@ def Allreduce(self, sendobj_or_inplace: T, recvobj: T, op: ReductionOperator) ->
         return self._comm.Allreduce(sendobj_or_inplace, recvobj, self._op_mapping[op])
 
     def Allreduce_inplace(self, recvobj: T, op: ReductionOperator) -> T:
-        return self._comm.Allreduce(mpi4py.MPI.IN_PLACE, recvobj, self._op_mapping[op])
+        return self._comm.Allreduce(MPI.IN_PLACE, recvobj, self._op_mapping[op])

ndsl/dsl/dace/orchestration.py

@@ -32,12 +32,7 @@
     report_memory_static_analysis,
 )
 from ndsl.logging import ndsl_log
-
-
-try:
-    import cupy as cp
-except ImportError:
-    cp = None
+from ndsl.optional_imports import cupy as cp
 
 
 def dace_inhibitor(func: Callable) -> Callable:

ndsl/dsl/gt4py_utils.py

@@ -1,19 +1,16 @@
 from functools import wraps
 from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
-import gt4py
 import numpy as np
+from gt4py import storage as gt_storage
+from gt4py.cartesian import backend as gt_backend
 
 from ndsl.constants import N_HALO_DEFAULT
 from ndsl.dsl.typing import DTypes, Field, Float
 from ndsl.logging import ndsl_log
+from ndsl.optional_imports import cupy as cp
 
 
-try:
-    import cupy as cp
-except ImportError:
-    cp = None
-
 # If True, automatically transfers memory between CPU and GPU (see gt4py.storage)
 managed_memory = True
 
@@ -188,7 +185,7 @@ def make_storage_data(
             backend=backend,
         )
 
-    storage = gt4py.storage.from_array(
+    storage = gt_storage.from_array(
         data,
         dtype,
         backend=backend,
@@ -340,7 +337,7 @@ def make_storage_from_shape(
             mask = (False, False, True)  # Assume 1D is a k-field
         else:
             mask = (n_dims * (True,)) + ((3 - n_dims) * (False,))
-    storage = gt4py.storage.zeros(
+    storage = gt_storage.zeros(
         shape,
         dtype,
         backend=backend,
@@ -448,7 +445,7 @@ def asarray(array, to_type=np.ndarray, dtype=None, order=None):
 
 
 def is_gpu_backend(backend: str) -> bool:
-    return gt4py.cartesian.backend.from_name(backend).storage_info["device"] == "gpu"
+    return gt_backend.from_name(backend).storage_info["device"] == "gpu"
 
 
 def zeros(shape, dtype=Float, *, backend: str):

ndsl/dsl/stencil.py

@@ -17,8 +17,9 @@
 )
 
 import dace
-import gt4py
 import numpy as np
+from gt4py.cartesian import config as gt_config
+from gt4py.cartesian import definitions as gt_definitions
 from gt4py.cartesian import gtscript
 from gt4py.cartesian.gtc.passes.oir_pipeline import DefaultPipeline, OirPipeline
 from gt4py.cartesian.stencil_object import StencilObject
@@ -39,12 +40,6 @@
 from ndsl.testing.comparison import LegacyMetric
 
 
-try:
-    import cupy as cp
-except ImportError:
-    cp = np
-
-
 def report_difference(args, kwargs, args_copy, kwargs_copy, function_name, gt_id):
     report_head = f"comparing against numpy for func {function_name}, gt_id {gt_id}:"
     report_segments = []
@@ -309,8 +304,8 @@ def __init__(
             dace.Config.set(
                 "default_build_folder",
                 value="{gt_root}/{gt_cache}/dacecache".format(
-                    gt_root=gt4py.cartesian.config.cache_settings["root_path"],
-                    gt_cache=gt4py.cartesian.config.cache_settings["dir_name"],
+                    gt_root=gt_config.cache_settings["root_path"],
+                    gt_cache=gt_config.cache_settings["dir_name"],
                 ),
             )
 
@@ -499,7 +494,7 @@ def _get_written_fields(cls, field_info) -> List[str]:
             if field_info[field_name]
             and bool(
                 field_info[field_name].access
-                & gt4py.cartesian.definitions.AccessKind.WRITE  # type: ignore
+                & gt_definitions.AccessKind.WRITE  # type: ignore
             )
         ]
         return write_fields

ndsl/dsl/typing.py

@@ -1,8 +1,8 @@
 import os
 from typing import Tuple, TypeAlias, Union, cast
 
-import gt4py.cartesian.gtscript as gtscript
 import numpy as np
+from gt4py.cartesian import gtscript
 
 
 # A Field

ndsl/initialization/allocator.py

@@ -1,11 +1,11 @@
 from typing import Callable, Optional, Sequence
 
 import numpy as np
+from gt4py import storage as gt_storage
 
 from ndsl.constants import SPATIAL_DIMS
 from ndsl.dsl.typing import Float
 from ndsl.initialization.sizer import GridSizer
-from ndsl.optional_imports import gt4py
 from ndsl.quantity import Quantity, QuantityHaloSpec
 
 
@@ -20,13 +20,13 @@ def __init__(self, backend: str):
         self.backend = backend
 
     def empty(self, *args, **kwargs) -> np.ndarray:
-        return gt4py.storage.empty(*args, backend=self.backend, **kwargs)
+        return gt_storage.empty(*args, backend=self.backend, **kwargs)
 
     def ones(self, *args, **kwargs) -> np.ndarray:
-        return gt4py.storage.ones(*args, backend=self.backend, **kwargs)
+        return gt_storage.ones(*args, backend=self.backend, **kwargs)
 
     def zeros(self, *args, **kwargs) -> np.ndarray:
-        return gt4py.storage.zeros(*args, backend=self.backend, **kwargs)
+        return gt_storage.zeros(*args, backend=self.backend, **kwargs)
 
 
 class QuantityFactory:

ndsl/io.py

@@ -1,9 +1,9 @@
 from typing import TextIO
 
 import cftime
+import xarray as xr
 
 import ndsl.filesystem as filesystem
-from ndsl.optional_imports import xarray as xr
 from ndsl.quantity import Quantity
 
 

ndsl/monitor/netcdf_monitor.py

@@ -5,13 +5,13 @@
 
 import fsspec
 import numpy as np
+import xarray as xr
 
 from ndsl.comm.communicator import Communicator
 from ndsl.dsl.typing import Float, get_precision
 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 ndsl.quantity import Quantity
 
 

ndsl/monitor/zarr_monitor.py

@@ -2,14 +2,13 @@
 from typing import List, Tuple, Union
 
 import cftime
+import xarray as xr
 
 import ndsl.constants as constants
 from ndsl.comm.partitioner import Partitioner, subtile_slice
 from ndsl.logging import ndsl_log
 from ndsl.monitor.convert import to_numpy
-from ndsl.optional_imports import cupy
-from ndsl.optional_imports import xarray as xr
-from ndsl.optional_imports import zarr
+from ndsl.optional_imports import cupy, zarr
 from ndsl.utils import list_by_dims
 
 

ndsl/optional_imports.py

@@ -14,11 +14,6 @@ def __call__(self, *args, **kwargs):
 except ModuleNotFoundError as err:
     zarr = RaiseWhenAccessed(err)
 
-try:
-    import xarray
-except ModuleNotFoundError as err:
-    xarray = None
-
 try:
     import cupy
 except ImportError:
@@ -30,14 +25,3 @@ def __call__(self, *args, **kwargs):
         cupy.cuda.runtime.deviceSynchronize()
     except cupy.cuda.runtime.CUDARuntimeError:
         cupy = None
-
-
-try:
-    import gt4py
-except ImportError:
-    gt4py = None
-
-try:
-    import dace
-except ImportError:
-    dace = None

ndsl/quantity/field_bundle.py

@@ -2,14 +2,13 @@
 from dataclasses import dataclass
 from typing import Any
 
+from gt4py.cartesian import gtscript
+
 from ndsl.dsl.typing import Float
 from ndsl.initialization.allocator import QuantityFactory
 from ndsl.quantity.quantity import Quantity
 
 
-import gt4py.cartesian.gtscript as gtscript  # isort: skip
-
-
 # ToDo: This is 4th dimensions restricted. We need a concept
 #       of data dimensions index here to be able to extend to N dimensions
 _DataDimensionIndex = int
@@ -44,7 +43,7 @@ def __init__(
             bundle_name: name of the bundle, accessible via `name`.
             quantity: data inputs as a nD array.
             mapping: sparse dict of [name, index] to be able to call tracers by name.
-            register_type: boolean to register the type as part of initilization.
+            register_type: boolean to register the type as part of initialization.
         """
         if len(quantity.shape) != 4:
             raise NotImplementedError("FieldBundle implementation restricted to 4D")