get lazy loading actually working for strds

Author: lrntct

README.md

@@ -153,7 +153,7 @@ Attributes:
   - [ ] Support `end_time`
   - [ ] Accept writing into a specific mapset (GRASS 8.5)
   - [ ] Accept non homogeneous 3D resolution in NS and EW dimensions (GRASS 8.5)
-- [ ] Lazy loading of all raster types
+- [x] Lazy loading of STDS on the time dimension
 - [ ] Properly test with lat-lon location
 
 ### Stretch goals

src/xarray_grass/__init__.py

@@ -1,7 +1,9 @@
 from xarray_grass.grass_interface import GrassConfig as GrassConfig
 from xarray_grass.grass_interface import GrassInterface as GrassInterface
 from xarray_grass.xarray_grass import GrassBackendEntrypoint as GrassBackendEntrypoint
-from xarray_grass.xarray_grass import GrassBackendArray as GrassBackendArray
+from xarray_grass.grass_backend_array import (
+    GrassSTDSBackendArray as GrassSTDSBackendArray,
+)
 from xarray_grass.to_grass import to_grass as to_grass
 from xarray_grass.coord_utils import RegionData as RegionData
 

src/xarray_grass/grass_backend_array.py

@@ -25,25 +25,24 @@
     from xarray_grass.grass_interface import GrassInterface
 
 
-class GrassBackendArray(BackendArray):
-    """Lazy loading of grass arrays"""
+class GrassSTDSBackendArray(BackendArray):
+    """Lazy loading of grass Space-Time DataSets (multiple maps in time series)"""
 
     def __init__(
         self,
         shape,
         dtype,
-        # lock,
-        map_id: str,
+        map_list: list,  # List of map metadata objects
         map_type: str,
         grass_interface: GrassInterface,
     ):
         self.shape = shape
         self.dtype = dtype
         self._lock = threading.Lock()
-        self.map_id = map_id
+        self.map_list = map_list  # List with .id attribute
         self.map_type = map_type  # "raster" or "raster3d"
         self.grass_interface = grass_interface
-        self._array: np.ndarray = None
+        self._cached_maps = {}  # Cache loaded maps by index
 
     def __getitem__(self, key: xr.core.indexing.ExplicitIndexer) -> np.typing.ArrayLike:
         """takes in input an index and returns a NumPy array"""
@@ -55,13 +54,44 @@ def __getitem__(self, key: xr.core.indexing.ExplicitIndexer) -> np.typing.ArrayL
         )
 
     def _raw_indexing_method(self, key: tuple):
+        """Load only the maps needed for the requested slice"""
         with self._lock:
-            if self._array is None:
-                self._array = self._load_map()
-            return self._array[key]
+            # key is a tuple of slices/indices for each dimension
+            # First dimension is time
+            time_key = key[0] if key else slice(None)
+            spatial_key = key[1:] if len(key) > 1 else ()
 
-    def _load_map(self):
-        if self.map_type == "raster":
-            return self.grass_interface.read_raster_map(self.map_id)
-        else:  # 'raster3d'
-            return self.grass_interface.read_raster3d_map(self.map_id)
+            # Determine which time indices are needed
+            if isinstance(time_key, slice):
+                time_indices = range(*time_key.indices(self.shape[0]))
+            elif isinstance(time_key, int):
+                time_indices = [time_key]
+            else:
+                time_indices = list(time_key)
+
+            # Load only the needed maps
+            result_list = []
+            for t_idx in time_indices:
+                if t_idx not in self._cached_maps:
+                    map_data = self.map_list[t_idx]
+                    if self.map_type == "raster":
+                        self._cached_maps[t_idx] = self.grass_interface.read_raster_map(
+                            map_data.id
+                        )
+                    else:  # 'raster3d'
+                        self._cached_maps[t_idx] = (
+                            self.grass_interface.read_raster3d_map(map_data.id)
+                        )
+
+                # Apply spatial indexing
+                if spatial_key:
+                    result_list.append(self._cached_maps[t_idx][spatial_key])
+                else:
+                    result_list.append(self._cached_maps[t_idx])
+
+            # Stack results along time dimension
+            if len(result_list) == 1 and isinstance(time_key, int):
+                # Single time slice requested as integer index
+                return result_list[0]
+            else:
+                return np.stack(result_list, axis=0)

src/xarray_grass/grass_interface.py

@@ -28,7 +28,7 @@
 from grass.script import array as garray
 import grass.pygrass.utils as gutils
 from grass.pygrass import raster as graster
-from grass.pygrass.raster.abstract import Info
+from grass.pygrass.raster.abstract import Info, RasterAbstractBase
 import grass.temporal as tgis
 
 from xarray_grass.coord_utils import (
@@ -230,11 +230,11 @@ def name_is_str3ds(self, name: str) -> bool:
         return bool(tgis.SpaceTimeRaster3DDataset(str3ds_id).is_in_db())
 
     def name_is_raster(self, raster_name: str) -> bool:
-        """return True if the given name is a map in the grass database
-        False if not
-        """
+        """return True if the given name is a raster map in the grass database."""
+        # Using pygrass instead of gscript is at least 40x faster
         map_id = self.get_id_from_name(raster_name)
-        return bool(gs.find_file(name=map_id, element="raster").get("file"))
+        map_object = RasterAbstractBase(map_id)
+        return map_object.exist()
 
     def name_is_raster_3d(self, raster3d_name: str) -> bool:
         """return True if the given name is a 3D raster in the grass database."""
@@ -261,11 +261,11 @@ def grass_dtype(self, dtype: str) -> str:
     @staticmethod
     def numpy_dtype(mtype: str) -> np.dtype:
         if mtype == "CELL":
-            dtype = np.int64
+            dtype = np.dtype("int64")
         elif mtype == "FCELL":
-            dtype = np.float32
+            dtype = np.dtype("float32")
         elif mtype == "DCELL":
-            dtype = np.float64
+            dtype = np.dtype("float64")
         else:
             raise ValueError(f"Unknown GRASS data type: {mtype}")
         return dtype

src/xarray_grass/xarray_grass.py

@@ -24,7 +24,7 @@
 
 import xarray_grass
 from xarray_grass.grass_interface import GrassInterface
-from xarray_grass.grass_backend_array import GrassBackendArray
+from xarray_grass.grass_backend_array import GrassSTDSBackendArray
 
 
 class GrassBackendEntrypoint(BackendEntrypoint):
@@ -265,6 +265,8 @@ def open_grass_maps(
                 data_array_list.append(data_array)
         if raise_on_not_found and any(not_found.values()):
             raise ValueError(f"Objects not found: {not_found}")
+
+        crs_wkt = gi.get_crs_wkt_str()
     finally:
         if session is not None:
             session.__exit__(None, None, None)
@@ -277,7 +279,7 @@ def open_grass_maps(
     data_array_dict = {da.name: da for da in data_array_list}
 
     attrs = {
-        "crs_wkt": gi.get_crs_wkt_str(),
+        "crs_wkt": crs_wkt,
         "Conventions": "CF-1.13-draft",
         # "title": "",
         "history": f"{datetime.now(timezone.utc)}: Created with xarray-grass version {xarray_grass.__version__}",
@@ -347,9 +349,7 @@ def open_grass_raster_3d(raster_3d_name: str, grass_i: GrassInterface) -> xr.Dat
 
 
 def open_grass_strds(strds_name: str, grass_i: GrassInterface) -> xr.DataArray:
-    """must be called from within a grass session
-    TODO: lazy loading
-    """
+    """Open a STRDS with lazy loading - data is only loaded when accessed"""
     strds_id = grass_i.get_id_from_name(strds_name)
     strds_name = grass_i.get_name_from_id(strds_id)
     x_coords, y_coords, _ = get_coordinates(grass_i, raster_3d=False).values()
@@ -360,45 +360,46 @@ def open_grass_strds(strds_name: str, grass_i: GrassInterface) -> xr.DataArray:
         time_unit = strds_infos.time_unit
     start_time_dim = f"start_time_{strds_name}"
     end_time_dim = f"end_time_{strds_name}"
-    dims = [start_time_dim, "y", "x"]
-    coordinates = dict.fromkeys(dims)
-    coordinates["x"] = x_coords
-    coordinates["y"] = y_coords
+
     map_list = grass_i.list_maps_in_strds(strds_id)
     region = grass_i.get_region()
-    array_list = []
-    for map_data in map_list:
-        # Lazy load the array
-        backend_array = GrassBackendArray(
-            shape=(region.rows, region.cols),
-            dtype=map_data.dtype,
-            map_id=map_data.id,
-            map_type="raster",
-            grass_interface=grass_i,
-        )
-        lazy_array = xr.core.indexing.LazilyIndexedArray(backend_array)
-        # add time dimension at the beginning
-        lazy_array_with_time = np.expand_dims(lazy_array, axis=0)
-
-        # ndarray = grass_i.read_raster_map(map_data.id)
-        # # add time dimension at the beginning
-        # ndarray = np.expand_dims(ndarray, axis=0)
-
-        coordinates[start_time_dim] = [map_data.start_time]
-        coordinates[end_time_dim] = (start_time_dim, [map_data.end_time])
-
-        data_array = xr.DataArray(
-            lazy_array_with_time,
-            coords=coordinates,
-            dims=dims,
-            name=strds_name,
-        )
-        array_list.append(data_array)
-    da_concat = xr.concat(array_list, dim=start_time_dim)
+
+    # Create a single backend array for the entire STRDS
+    backend_array = GrassSTDSBackendArray(
+        shape=(len(map_list), region.rows, region.cols),
+        dtype=map_list[0].dtype,
+        map_list=map_list,
+        map_type="raster",
+        grass_interface=grass_i,
+    )
+    lazy_array = xr.core.indexing.LazilyIndexedArray(backend_array)
+
+    # Create Variable with lazy array
+    var = xr.Variable(dims=[start_time_dim, "y", "x"], data=lazy_array)
+
+    # Extract time coordinates
+    start_times = [map_data.start_time for map_data in map_list]
+    end_times = [map_data.end_time for map_data in map_list]
+
+    # Create coordinates
+    coordinates = {
+        "x": x_coords,
+        "y": y_coords,
+        start_time_dim: start_times,
+        end_time_dim: (start_time_dim, end_times),
+    }
+
+    # Convert to DataArray
+    data_array = xr.DataArray(
+        var,
+        coords=coordinates,
+        name=strds_name,
+    )
+
     # Add CF attributes
     r_infos = grass_i.get_raster_info(map_list[0].id)
     da_with_attrs = set_cf_coordinates(
-        da_concat,
+        data_array,
         gi=grass_i,
         is_3d=False,
         time_dims=[start_time_dim, end_time_dim],
@@ -414,7 +415,7 @@ def open_grass_strds(strds_name: str, grass_i: GrassInterface) -> xr.DataArray:
 
 
 def open_grass_str3ds(str3ds_name: str, grass_i: GrassInterface) -> xr.DataArray:
-    """Open a series of 3D raster maps.
+    """Open a STR3DS with lazy loading - data is only loaded when accessed
     TODO: Figure out what to do when the z value of the maps is time."""
     str3ds_id = grass_i.get_id_from_name(str3ds_name)
     str3ds_name = grass_i.get_name_from_id(str3ds_id)
@@ -426,43 +427,47 @@ def open_grass_str3ds(str3ds_name: str, grass_i: GrassInterface) -> xr.DataArray
         time_unit = strds_infos.time_unit
     start_time_dim = f"start_time_{str3ds_name}"
     end_time_dim = f"end_time_{str3ds_name}"
-    dims = [start_time_dim, "z", "y_3d", "x_3d"]
-    coordinates = dict.fromkeys(dims)
-    coordinates["x_3d"] = x_coords
-    coordinates["y_3d"] = y_coords
-    coordinates["z"] = z_coords
+
     map_list = grass_i.list_maps_in_str3ds(str3ds_id)
     region = grass_i.get_region()
-    array_list = []
-    for map_data in map_list:
-        # Lazy load the map
-        backend_array = GrassBackendArray(
-            shape=(region.depths, region.rows3, region.cols3),
-            dtype=map_data.dtype,
-            map_id=map_data.id,
-            map_type="raster3d",
-            grass_interface=grass_i,
-        )
-        lazy_array = xr.core.indexing.LazilyIndexedArray(backend_array)
-        # add time dimension at the beginning
-        lazy_array_with_time = np.expand_dims(lazy_array, axis=0)
-
-        coordinates[start_time_dim] = [map_data.start_time]
-        coordinates[end_time_dim] = (start_time_dim, [map_data.end_time])
-
-        data_array = xr.DataArray(
-            lazy_array_with_time,
-            coords=coordinates,
-            dims=dims,
-            name=str3ds_name,
-        )
-        array_list.append(data_array)
 
-    da_concat = xr.concat(array_list, dim=start_time_dim)
+    # Create a single backend array for the entire STR3DS
+    backend_array = GrassSTDSBackendArray(
+        shape=(len(map_list), region.depths, region.rows3, region.cols3),
+        dtype=map_list[0].dtype,
+        map_list=map_list,
+        map_type="raster3d",
+        grass_interface=grass_i,
+    )
+    lazy_array = xr.core.indexing.LazilyIndexedArray(backend_array)
+
+    # Create Variable with lazy array
+    var = xr.Variable(dims=[start_time_dim, "z", "y_3d", "x_3d"], data=lazy_array)
+
+    # Extract time coordinates
+    start_times = [map_data.start_time for map_data in map_list]
+    end_times = [map_data.end_time for map_data in map_list]
+
+    # Create coordinates
+    coordinates = {
+        "x_3d": x_coords,
+        "y_3d": y_coords,
+        "z": z_coords,
+        start_time_dim: start_times,
+        end_time_dim: (start_time_dim, end_times),
+    }
+
+    # Convert to DataArray
+    data_array = xr.DataArray(
+        var,
+        coords=coordinates,
+        name=str3ds_name,
+    )
+
     # Add CF attributes
     r3_infos = grass_i.get_raster3d_info(map_list[0].id)
     da_with_attrs = set_cf_coordinates(
-        da_concat,
+        data_array,
         gi=grass_i,
         is_3d=True,
         z_unit=r3_infos["vertical_units"],