|
1 | 1 | from collections.abc import Hashable, Sequence |
2 | 2 | from dataclasses import dataclass |
3 | 3 | from functools import cached_property |
4 | | -from typing import Any, Literal, cast |
| 4 | +from typing import TYPE_CHECKING, Any, Literal, cast |
5 | 5 |
|
6 | 6 | import numpy as np |
7 | 7 | import numpy.typing as npt |
|
17 | 17 | from parcels._sgrid.accessor import _get_dim_to_axis_mapping |
18 | 18 | from parcels._sgrid.core import SGRID_PADDING_TO_XGCM_POSITION |
19 | 19 |
|
| 20 | +if TYPE_CHECKING: |
| 21 | + import xgcm.axis |
| 22 | + |
20 | 23 | _FIELD_DATA_ORDERING: Sequence[ptyping.XgcmAxisDirection] = "TZYX" |
21 | 24 | _XGRID_AXES_ORDERING: Sequence[ptyping.XgridAxis] = "ZYX" |
22 | 25 |
|
23 | | -_DEFAULT_XGCM_KWARGS: dict[str, Any] = {"periodic": False} |
| 26 | +_DEFAULT_XGCM_KWARGS: dict[str, Any] = {"padding": "fill"} |
24 | 27 |
|
25 | 28 |
|
26 | | -def get_cell_count_along_dim(ds: xr.Dataset, axis: xgcm.Axis) -> int: |
| 29 | +def get_cell_count_along_dim(ds: xr.Dataset, axis: xgcm.axis.Axis) -> int: |
27 | 30 | first_coord = list(axis.coords.items())[0] |
28 | 31 | _, coord_var = first_coord |
29 | 32 |
|
30 | 33 | return ds[coord_var].size - 1 |
31 | 34 |
|
32 | 35 |
|
33 | | -def get_time(ds: xr.Dataset, axis: xgcm.Axis) -> npt.NDArray: |
| 36 | +def get_time(ds: xr.Dataset, axis: xgcm.axis.Axis) -> npt.NDArray: |
34 | 37 | return ds[axis.coords["center"]].values |
35 | 38 |
|
36 | 39 |
|
|
0 commit comments