Reordering dimensions to time,depth,lat,lon in Grid class

Author: erikvansebille

parcels/field.py

@@ -130,12 +130,12 @@ class Field:
         A numpy array containing the timestamps for each of the files in filenames, for loading
         from netCDF files only. Default is None if the netCDF dimensions dictionary includes time.
     grid : parcels.grid.Grid
-        :class:`parcels.grid.Grid` object containing all the lon, lat depth, time
+        :class:`parcels.grid.Grid` object containing all the time, depth, lat, lon,
         mesh and time_origin information. Can be constructed from any of the Grid objects
     fieldtype : str
         Type of Field to be used for UnitConverter (either 'U', 'V', 'Kh_zonal', 'Kh_meridional' or None)
     transpose : bool
-        Transpose data to required (lon, lat) layout
+        Transpose data to required (time, depth, lat, lon) layout
     vmin : float
         Minimum allowed value on the field. Data below this value are set to zero
     vmax : float
@@ -211,7 +211,7 @@ def __init__(
         if grid:
             if grid.defer_load and isinstance(data, np.ndarray):
                 raise ValueError(
-                    "Cannot combine Grid from defer_loaded Field with np.ndarray data. please specify lon, lat, depth and time dimensions separately"
+                    "Cannot combine Grid from defer_loaded Field with np.ndarray data. please specify time, depth, lat and lon dimensions separately"
                 )
             self._grid = grid
         else:
@@ -220,7 +220,7 @@ def __init__(
                 time = np.array([time_origin.reltime(t) for t in time])
             else:
                 time_origin = TimeConverter(0)
-            self._grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
+            self._grid = Grid.create_grid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh)
         self.igrid = -1
         self.fieldtype = self.name if fieldtype is None else fieldtype
         self.to_write = to_write
@@ -687,7 +687,7 @@ def from_netcdf(
             time, time_origin, timeslices, dataFiles = cls._collect_timeslices(
                 timestamps, data_filenames, _grid_fb_class, dimensions, indices, netcdf_engine
             )
-            grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
+            grid = Grid.create_grid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh)
             grid.timeslices = timeslices
             kwargs["dataFiles"] = dataFiles
         elif grid is not None and ("dataFiles" not in kwargs or kwargs["dataFiles"] is None):
@@ -838,7 +838,7 @@ def from_xarray(
         time_origin = TimeConverter(time[0])
         time = time_origin.reltime(time)  # type: ignore[assignment]
 
-        grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
+        grid = Grid.create_grid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh)
         kwargs["time_periodic"] = time_periodic
         return cls(
             name,

parcels/fieldset.py

@@ -155,7 +155,7 @@ def from_data(
                 time = np.array([time_origin.reltime(t) for t in time])
             else:
                 time_origin = kwargs.pop("time_origin", TimeConverter(0))
-            grid = Grid.create_grid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh)
+            grid = Grid.create_grid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh)
             if "creation_log" not in kwargs.keys():
                 kwargs["creation_log"] = "from_data"
 

parcels/grid.py

@@ -44,9 +44,10 @@ class Grid:
 
     def __init__(
         self,
-        lon: npt.NDArray,
+        time: npt.NDArray,
+        depth: npt.NDArray,
         lat: npt.NDArray,
-        time: npt.NDArray | None,
+        lon: npt.NDArray,
         time_origin: TimeConverter | None,
         mesh: Mesh,
     ):
@@ -71,6 +72,7 @@ def __init__(
 
         self._lon = lon
         self._lat = lat
+        self._depth = depth
         self.time = time
         self.time_full = self.time  # needed for deferred_loaded Fields
         self._time_origin = TimeConverter() if time_origin is None else time_origin
@@ -98,7 +100,7 @@ def __repr__(self):
         with np.printoptions(threshold=5, suppress=True, linewidth=120, formatter={"float": "{: 0.2f}".format}):
             return (
                 f"{type(self).__name__}("
-                f"lon={self.lon!r}, lat={self.lat!r}, time={self.time!r}, "
+                f"lon={self.lon!r}, lat={self.lat!r}, time={self.time!r}, depth={self.depth!r}, "
                 f"time_origin={self.time_origin!r}, mesh={self.mesh!r})"
             )
 
@@ -195,10 +197,10 @@ def load_chunk(self):
 
     @staticmethod
     def create_grid(
-        lon: npt.ArrayLike,
+        time: npt.ArrayLike,
+        depth: npt.ArrayLike,
         lat: npt.ArrayLike,
-        depth,
-        time,
+        lon: npt.ArrayLike,
         time_origin,
         mesh: Mesh,
         **kwargs,
@@ -211,14 +213,14 @@ def create_grid(
 
         if len(lon.shape) <= 1:
             if depth is None or len(depth.shape) <= 1:
-                return RectilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return RectilinearZGrid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh, **kwargs)
             else:
-                return RectilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return RectilinearSGrid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh, **kwargs)
         else:
             if depth is None or len(depth.shape) <= 1:
-                return CurvilinearZGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return CurvilinearZGrid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh, **kwargs)
             else:
-                return CurvilinearSGrid(lon, lat, depth, time, time_origin=time_origin, mesh=mesh, **kwargs)
+                return CurvilinearSGrid(time, depth, lat, lon, time_origin=time_origin, mesh=mesh, **kwargs)
 
     @property
     def ctypes_struct(self):
@@ -461,14 +463,14 @@ class RectilinearGrid(Grid):
 
     """
 
-    def __init__(self, lon, lat, time, time_origin, mesh: Mesh):
+    def __init__(self, time, depth, lat, lon, time_origin, mesh: Mesh):
         assert isinstance(lon, np.ndarray) and len(lon.shape) <= 1, "lon is not a numpy vector"
         assert isinstance(lat, np.ndarray) and len(lat.shape) <= 1, "lat is not a numpy vector"
         assert isinstance(time, np.ndarray) or not time, "time is not a numpy array"
         if isinstance(time, np.ndarray):
             assert len(time.shape) == 1, "time is not a vector"
 
-        super().__init__(lon, lat, time, time_origin, mesh)
+        super().__init__(time, depth, lat, lon, time_origin, mesh)
         self.tdim = self.time.size
 
         if self.ydim > 1 and self.lat[-1] < self.lat[0]:
@@ -559,8 +561,8 @@ class RectilinearZGrid(RectilinearGrid):
         2. flat: No conversion, lat/lon are assumed to be in m.
     """
 
-    def __init__(self, lon, lat, depth=None, time=None, time_origin=None, mesh: Mesh = "flat"):
-        super().__init__(lon, lat, time, time_origin, mesh)
+    def __init__(self, time, depth, lat, lon, time_origin=None, mesh: Mesh = "flat"):
+        super().__init__(time, depth, lat, lon, time_origin, mesh)
         if isinstance(depth, np.ndarray):
             assert len(depth.shape) <= 1, "depth is not a vector"
 
@@ -592,7 +594,7 @@ class RectilinearSGrid(RectilinearGrid):
         which are s-coordinates.
         s-coordinates can be terrain-following (sigma) or iso-density (rho) layers,
         or any generalised vertical discretisation.
-        The depth of each node depends then on the horizontal position (lon, lat),
+        The depth of each node depends then on the horizontal position (lat, lon),
         the number of the layer and the time is depth is a 4D array.
         depth array is either a 4D array[xdim][ydim][zdim][tdim] or a 3D array[xdim][ydim[zdim].
     time :
@@ -610,14 +612,14 @@ class RectilinearSGrid(RectilinearGrid):
 
     def __init__(
         self,
-        lon: npt.NDArray,
-        lat: npt.NDArray,
+        time: npt.NDArray,
         depth: npt.NDArray,
-        time: npt.NDArray | None = None,
+        lat: npt.NDArray,
+        lon: npt.NDArray,
         time_origin: TimeConverter | None = None,
         mesh: Mesh = "flat",
     ):
-        super().__init__(lon, lat, time, time_origin, mesh)
+        super().__init__(time, depth, lat, lon, time_origin, mesh)
         assert isinstance(depth, np.ndarray) and len(depth.shape) in [3, 4], "depth is not a 3D or 4D numpy array"
 
         self._gtype = GridType.RectilinearSGrid
@@ -656,9 +658,10 @@ def zdim(self):
 class CurvilinearGrid(Grid):
     def __init__(
         self,
-        lon: npt.NDArray,
+        time: npt.NDArray,
+        depth: npt.NDArray,
         lat: npt.NDArray,
-        time: npt.NDArray | None = None,
+        lon: npt.NDArray,
         time_origin: TimeConverter | None = None,
         mesh: Mesh = "flat",
     ):
@@ -670,7 +673,7 @@ def __init__(
 
         lon = lon.squeeze()
         lat = lat.squeeze()
-        super().__init__(lon, lat, time, time_origin, mesh)
+        super().__init__(time, None, lat, lon, time_origin, mesh)
         self.tdim = self.time.size
 
     @property
@@ -770,14 +773,14 @@ class CurvilinearZGrid(CurvilinearGrid):
 
     def __init__(
         self,
-        lon: npt.NDArray,
+        time: npt.NDArray,
+        depth: npt.NDArray,
         lat: npt.NDArray,
-        depth: npt.NDArray | None = None,
-        time: npt.NDArray | None = None,
+        lon: npt.NDArray,
         time_origin: TimeConverter | None = None,
         mesh: Mesh = "flat",
     ):
-        super().__init__(lon, lat, time, time_origin, mesh)
+        super().__init__(time, depth, lat, lon, time_origin, mesh)
         if isinstance(depth, np.ndarray):
             assert len(depth.shape) == 1, "depth is not a vector"
 
@@ -808,7 +811,7 @@ class CurvilinearSGrid(CurvilinearGrid):
         which are s-coordinates.
         s-coordinates can be terrain-following (sigma) or iso-density (rho) layers,
         or any generalised vertical discretisation.
-        The depth of each node depends then on the horizontal position (lon, lat),
+        The depth of each node depends then on the horizontal position (lat, lon),
         the number of the layer and the time is depth is a 4D array.
         depth array is either a 4D array[xdim][ydim][zdim][tdim] or a 3D array[xdim][ydim[zdim].
     time :
@@ -826,14 +829,14 @@ class CurvilinearSGrid(CurvilinearGrid):
 
     def __init__(
         self,
-        lon: npt.NDArray,
-        lat: npt.NDArray,
+        time: npt.NDArray,
         depth: npt.NDArray,
-        time: npt.NDArray | None = None,
+        lat: npt.NDArray,
+        lon: npt.NDArray,
         time_origin: TimeConverter | None = None,
         mesh: Mesh = "flat",
     ):
-        super().__init__(lon, lat, time, time_origin, mesh)
+        super().__init__(time, depth, lat, lon, time_origin, mesh)
         assert isinstance(depth, np.ndarray) and len(depth.shape) in [3, 4], "depth is not a 4D numpy array"
 
         self._gtype = GridType.CurvilinearSGrid

parcels/gridset.py

@@ -44,7 +44,7 @@ def add_grid(self, field):
         field.igrid = self.grids.index(field.grid)
 
     def dimrange(self, dim):
-        """Returns maximum value of a dimension (lon, lat, depth or time)
+        """Returns maximum value of a dimension (time, depth, lat or lon)
         on 'left' side and minimum value on 'right' side for all grids
         in a gridset. Useful for finding e.g. longitude range that
         overlaps on all grids in a gridset.

parcels/particledata.py

@@ -64,7 +64,7 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, n
         ), "particle's initial depth is None - incompatible with the ParticleData class. Invalid state."
         assert lon.size == lat.size and lon.size == depth.size, "lon, lat, depth don't all have the same lenghts."
 
-        assert lon.size == time.size, "time and positions (lon, lat, depth) don't have the same lengths."
+        assert lon.size == time.size, "time and positions (depth, lat, lon) don't have the same lengths."
 
         # If a partitioning function for MPI runs has been passed into the
         # particle creation with the "partition_function" kwarg, retrieve it here.
@@ -74,7 +74,7 @@ def __init__(self, pclass, lon, lat, depth, time, lonlatdepth_dtype, pid_orig, n
         for kwvar in kwargs:
             assert (
                 lon.size == kwargs[kwvar].size
-            ), f"{kwvar} and positions (lon, lat, depth) don't have the same lengths."
+            ), f"{kwvar} and positions (depth, lat, lon) don't have the same lengths."
 
         offset = np.max(pid) if (pid is not None) and len(pid) > 0 else -1
         if MPI:

tests/test_deprecations.py

@@ -285,7 +285,7 @@ def create_test_data():
     lon_g0 = np.linspace(0, 1000, 11, dtype=np.float32)
     lat_g0 = np.linspace(0, 1000, 11, dtype=np.float32)
     time_g0 = np.linspace(0, 1000, 2, dtype=np.float64)
-    grid = RectilinearZGrid(lon_g0, lat_g0, time=time_g0)
+    grid = RectilinearZGrid(time_g0, None, lat_g0, lon_g0)
 
     pfile = ParticleFile("test.zarr", pset, outputdt=1)
 

tests/test_diffusion.py

@@ -68,7 +68,7 @@ def test_fieldKh_SpatiallyVaryingDiffusion(mesh, mode, kernel):
     for x in range(xdim):
         Kh[:, x] = np.tanh(fieldset.U.lon[x] / fieldset.U.lon[-1] * 10.0) * xdim / 2.0 + xdim / 2.0 + 100.0
 
-    grid = RectilinearZGrid(lon=fieldset.U.lon, lat=fieldset.U.lat, mesh=mesh)
+    grid = RectilinearZGrid(time=None, depth=None, lat=fieldset.U.lat, lon=fieldset.U.lon, mesh=mesh)
     fieldset.add_field(Field("Kh_zonal", Kh, grid=grid))
     fieldset.add_field(Field("Kh_meridional", Kh, grid=grid))
     fieldset.add_constant("dres", fieldset.U.lon[1] - fieldset.U.lon[0])

tests/test_fieldset.py

@@ -944,7 +944,7 @@ def test_fieldset_defer_loading_with_diff_time_origin(tmpdir, fail):
         Wtime_origin = TimeConverter(np.datetime64("2018-04-22"))
     else:
         Wtime_origin = TimeConverter(np.datetime64("2018-04-18"))
-    gridW = RectilinearZGrid(dims1["lon"], dims1["lat"], dims1["depth"], dims1["time"], time_origin=Wtime_origin)
+    gridW = RectilinearZGrid(dims1["time"], dims1["depth"], dims1["lat"], dims1["lon"], time_origin=Wtime_origin)
     fieldW = Field("W", np.zeros(data1["U"].shape), grid=gridW)
     fieldset_out.add_field(fieldW)
     fieldset_out.write(filepath)

tests/test_grids.py

@@ -40,7 +40,7 @@ def temp_func(lon, lat):
     lon_g0 = np.linspace(0, a, xdim_g0, dtype=np.float32)
     lat_g0 = np.linspace(0, b, ydim_g0, dtype=np.float32)
     time_g0 = np.linspace(0.0, 1000.0, 2, dtype=np.float64)
-    grid_0 = RectilinearZGrid(lon_g0, lat_g0, time=time_g0)
+    grid_0 = RectilinearZGrid(time_g0, None, lat_g0, lon_g0)
 
     # Grid 1
     xdim_g1 = 51
@@ -49,7 +49,7 @@ def temp_func(lon, lat):
     lon_g1 = np.linspace(0, a, xdim_g1, dtype=np.float32)
     lat_g1 = np.linspace(0, b, ydim_g1, dtype=np.float32)
     time_g1 = np.linspace(0.0, 1000.0, 2, dtype=np.float64)
-    grid_1 = RectilinearZGrid(lon_g1, lat_g1, time=time_g1)
+    grid_1 = RectilinearZGrid(time_g1, None, lat_g1, lon_g1)
 
     u_data = np.ones((lon_g0.size, lat_g0.size, time_g0.size), dtype=np.float32)
     u_data = 2 * u_data
@@ -109,7 +109,7 @@ def test_time_format_in_grid():
     lat = np.linspace(0, 1, 2, dtype=np.float32)
     time = np.array([np.datetime64("2000-01-01")] * 2)
     with pytest.raises(AssertionError, match="Time vector"):
-        RectilinearZGrid(lon, lat, time=time)
+        RectilinearZGrid(time, None, lat, lon)
 
 
 def test_negate_depth():
@@ -126,12 +126,12 @@ def test_avoid_repeated_grids():
     lon_g0 = np.linspace(0, 1000, 11, dtype=np.float32)
     lat_g0 = np.linspace(0, 1000, 11, dtype=np.float32)
     time_g0 = np.linspace(0, 1000, 2, dtype=np.float64)
-    grid_0 = RectilinearZGrid(lon_g0, lat_g0, time=time_g0)
+    grid_0 = RectilinearZGrid(time_g0, None, lat_g0, lon_g0)
 
     lon_g1 = np.linspace(0, 1000, 21, dtype=np.float32)
     lat_g1 = np.linspace(0, 1000, 21, dtype=np.float32)
     time_g1 = np.linspace(0, 1000, 2, dtype=np.float64)
-    grid_1 = RectilinearZGrid(lon_g1, lat_g1, time=time_g1)
+    grid_1 = RectilinearZGrid(time_g1, None, lat_g1, lon_g1)
 
     u_data = np.zeros((lon_g0.size, lat_g0.size, time_g0.size), dtype=np.float32)
     u_field = Field("U", u_data, grid=grid_0, transpose=True)
@@ -166,8 +166,8 @@ def bath_func(lon):
         for k in range(zdim):
             depth_g0[k, :, i] = bath[i] * k / (zdim - 1)
 
-    grid_0 = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
-    grid_1 = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
+    grid_0 = RectilinearSGrid(None, depth_g0, lat_g0, lon_g0)
+    grid_1 = RectilinearSGrid(None, depth_g0, lat_g0, lon_g0)
 
     u_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
     v_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
@@ -216,7 +216,7 @@ def bath_func(lon):
             else:
                 depth_g0[k, :, i] = bath[i] * k / (zdim - 1)
 
-    grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0, time=time_g0)
+    grid = RectilinearSGrid(time_g0, depth_g0, lat_g0, lon_g0)
 
     u_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32)
     v_data = np.zeros((grid.tdim, grid.zdim, grid.ydim, grid.xdim), dtype=np.float32)
@@ -262,7 +262,7 @@ def bath_func(lon):
             depth_g0[i, :, k] = bath[i] * k / (depth_g0.shape[2] - 1)
     depth_g0 = depth_g0.transpose()  # we don't change it on purpose, to check if the transpose op if fixed in jit
 
-    grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
+    grid = RectilinearSGrid(None, depth_g0, lat_g0, lon_g0)
 
     zdim = depth_g0.shape[0]
     u_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
@@ -306,7 +306,7 @@ def bath_func(lon):
         for k in range(zdim):
             depth_g0[k, :, i] = bath[i] * k / (zdim - 1)
 
-    grid = RectilinearSGrid(lon_g0, lat_g0, depth=depth_g0)
+    grid = RectilinearSGrid(None, depth_g0, lat_g0, lon_g0)
 
     u_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
     v_data = np.zeros((zdim, lat_g0.size, lon_g0.size), dtype=np.float32)
@@ -347,7 +347,7 @@ def test_curvilinear_grids(mode):
     lon = r * np.cos(theta)
     lat = r * np.sin(theta)
     time = np.array([0, 86400], dtype=np.float64)
-    grid = CurvilinearZGrid(lon, lat, time=time)
+    grid = CurvilinearZGrid(time, None, lat, lon)
 
     u_data = np.ones((2, y.size, x.size), dtype=np.float32)
     v_data = np.zeros((2, y.size, x.size), dtype=np.float32)

tests/test_particlesets.py

@@ -139,7 +139,7 @@ def test_pset_create_field_curvi():
 
     x = -1 + r * np.cos(theta)
     y = -1 + r * np.sin(theta)
-    grid = CurvilinearZGrid(x, y)
+    grid = CurvilinearZGrid(time=None, depth=None, lat=y, lon=x)
 
     u = np.ones(x.shape)
     v = np.where(np.logical_and(theta > np.pi / 4, theta < np.pi / 3), 1, 0)