diff --git a/src/parcels/_core/fieldset.py b/src/parcels/_core/fieldset.py index 9f8ecb8b2..fa2dfac8d 100644 --- a/src/parcels/_core/fieldset.py +++ b/src/parcels/_core/fieldset.py @@ -7,7 +7,6 @@ import cf_xarray # noqa: F401 import numpy as np -import uxarray as ux import xarray as xr import parcels._typing as ptyping @@ -29,6 +28,8 @@ ) if TYPE_CHECKING: + import uxarray as ux + from parcels._core.basegrid import BaseGrid from parcels._typing import TimeLike __all__ = ["FieldSet"] diff --git a/src/parcels/_core/model.py b/src/parcels/_core/model.py index 848bf2fd2..45a5375ea 100644 --- a/src/parcels/_core/model.py +++ b/src/parcels/_core/model.py @@ -3,10 +3,9 @@ import warnings from abc import ABC, abstractmethod from collections.abc import Hashable, Sequence -from typing import Any, Self +from typing import TYPE_CHECKING, Any, Self import cf_xarray # noqa: F401 -import uxarray as ux import xarray as xr import zarr from dask import is_dask_collection @@ -37,6 +36,9 @@ ) from parcels.interpolators._base import ScalarInterpolator, VectorInterpolator +if TYPE_CHECKING: + import uxarray as ux + class ModelData(ABC): data: Any @@ -320,6 +322,8 @@ def assert_vector_field_components_in_dataset(ds: xr.Dataset, vector_fields: pty class UnstructuredModelData(ModelData): def __init__(self, data: ux.UxDataset, grid: UxGrid, vector_field_components: ptyping.VectorFields): + import uxarray as ux + if not isinstance(data, ux.UxDataset): raise ValueError(f"Expected `data` to be an uxarray.UxDataset . Got {type(data)}") diff --git a/src/parcels/_core/uxgrid.py b/src/parcels/_core/uxgrid.py index 46cb2da03..3257785db 100644 --- a/src/parcels/_core/uxgrid.py +++ b/src/parcels/_core/uxgrid.py @@ -1,15 +1,17 @@ from __future__ import annotations -from typing import Literal +from typing import TYPE_CHECKING, Literal import numpy as np -import uxarray as ux from dask import is_dask_collection from parcels._core.basegrid import BaseGrid from parcels._core.index_search import GRID_SEARCH_ERROR, _search_1d_array, uxgrid_point_in_cell from parcels._core.mesh import SphericalMesh, get_mesh +if TYPE_CHECKING: + import uxarray as ux + _UXGRID_AXES = Literal["Z", "FACE"] @@ -36,6 +38,8 @@ def __init__( mesh : str The type of mesh used for the grid. Either "flat" or "spherical". """ + import uxarray as ux + if grid.n_max_face_nodes > 3: raise ValueError("Provided ux.grid.Grid must contain only triangular cells (n_max_face_nodes=3)") self.uxgrid = grid