Enable patchwise training and prediction

deepsensor/data/loader.py

@@ -1,6 +1,7 @@
 import copy
 import itertools
 import json
+import operator
 import os
 import random
 from typing import List, Optional, Sequence, Tuple, Union
@@ -821,41 +822,53 @@ def _compute_global_coordinate_bounds(self) -> List[float]:
 
         return [x1_min, x1_max, x2_min, x2_max]
 
-    def _compute_x1x2_direction(self) -> str:
+    def _compute_x1x2_direction(self) -> dict:
         """Compute whether the x1 and x2 coords are ascending or descending.
 
         Returns:
-        -------
-        coord_directions: dict(str)
-            Dictionary containing two keys: x1 and x2, with boolean values
-            defining if these coordings increase or decrease from top left corner.
+            dict(bool)
+                Dictionary containing two keys: x1 and x2, with boolean values
+                defining if these coordings increase or decrease from top left corner.
+
+        Raises:
+            ValueError:
+                If all datasets are non-gridded or if direction of ascending
+                coordinates does not match across non-gridded datasets.
 
         """
+        non_gridded = {"x1": None, "x2": None}  # value to use for non-gridded data
+        ascending = []
         for var in itertools.chain(self.context, self.target):
             if isinstance(var, (xr.Dataset, xr.DataArray)):
                 coord_x1_left = var.x1[0]
                 coord_x1_right = var.x1[-1]
                 coord_x2_top = var.x2[0]
                 coord_x2_bottom = var.x2[-1]
 
-                x1_ascend = True if coord_x1_left <= coord_x1_right else False
-                x2_ascend = True if coord_x2_top <= coord_x2_bottom else False
-
-                coord_directions = {
-                    "x1": x1_ascend,
-                    "x2": x2_ascend,
-                }
+                ascending.append(
+                    {
+                        "x1": True if coord_x1_left <= coord_x1_right else False,
+                        "x2": True if coord_x2_top <= coord_x2_bottom else False,
+                    }
+                )
 
-            # TODO- what to input for pd.dataframe
             elif isinstance(var, (pd.DataFrame, pd.Series)):
-                # var_x1_min = var.index.get_level_values("x1").min()
-                # var_x1_max = var.index.get_level_values("x1").max()
-                # var_x2_min = var.index.get_level_values("x2").min()
-                # var_x2_max = var.index.get_level_values("x2").max()
+                ascending.append(non_gridded)
 
-                coord_directions = {"x1": None, "x2": None}
+        if len(list(filter(lambda x: x != non_gridded, ascending))) == 0:
+            raise ValueError(
+                "All data is non gridded, can not proceed with sliding window sampling."
+            )
+
+        # get the directions for only the gridded data
+        gridded = list(filter(lambda x: x != non_gridded, ascending))
+        # raise error if directions don't match across gridded data
+        if gridded.count(gridded[0]) != len(gridded):
+            raise ValueError(
+                "Direction of ascending coordinates does not match across all gridded datasets."
+            )
 
-        return coord_directions
+        return gridded[0]
 
     def sample_random_window(self, patch_size: Tuple[float]) -> Sequence[float]:
         """Sample random window uniformly from global coordinates to slice data.
@@ -983,8 +996,8 @@ def task_generation(  # noqa: D102
         ] = None,
         split_frac: float = 0.5,
         bbox: Sequence[float] = None,
-        patch_size: Union[float, tuple[float]] = None,
-        stride: Union[float, tuple[float]] = None,
+        patch_size: Union[float, Tuple[float]] = None,
+        stride: Union[float, Tuple[float]] = None,
         datewise_deterministic: bool = False,
         seed_override: Optional[int] = None,
     ) -> Task:
@@ -1428,96 +1441,74 @@ def sample_sliding_window(
         patch_size : Tuple[float]
             Tuple of window extent
 
-        Stride : Tuple[float]
+        stride : Tuple[float]
             Tuple of step size between each patch along x1 and x2 axis.
 
         Returns:
         -------
-        bbox: List[float]
+        List[float]
             Sequence of patch spatial extent as [x1_min, x1_max, x2_min, x2_max].
         """
         # define patch size in x1/x2
-        x1_extend, x2_extend = patch_size
+        size = {}
+        size["x1"], size["x2"] = patch_size
 
         # define stride length in x1/x2 or set to patch_size if undefined
         if stride is None:
             stride = patch_size
 
-        dy, dx = stride
-        # Calculate the global bounds of context and target set.
-        x1_min, x1_max, x2_min, x2_max = self.coord_bounds
-        ## start with first patch top left hand corner at x1_min, x2_min
-        patch_list = []
-
-        # Todo: simplify these elif statements
-        if self.coord_directions["x1"] == False and self.coord_directions["x2"] == True:
-            for y in np.arange(x1_max, x1_min, -dy):
-                for x in np.arange(x2_min, x2_max, dx):
-                    if y - x1_extend < x1_min:
-                        y0 = x1_min + x1_extend
-                    else:
-                        y0 = y
-                    if x + x2_extend > x2_max:
-                        x0 = x2_max - x2_extend
-                    else:
-                        x0 = x
+        step = {}
+        step["x1"], step["x2"] = stride
 
-                    # bbox of x1_min, x1_max, x2_min, x2_max per patch
-                    bbox = [y0 - x1_extend, y0, x0, x0 + x2_extend]
-                    patch_list.append(bbox)
-
-        elif (
-            self.coord_directions["x1"] == False
-            and self.coord_directions["x2"] == False
-        ):
-            for y in np.arange(x1_max, x1_min, -dy):
-                for x in np.arange(x2_max, x2_min, -dx):
-                    if y - x1_extend < x1_min:
-                        y0 = x1_min + x1_extend
-                    else:
-                        y0 = y
-                    if x - x2_extend < x2_min:
-                        x0 = x2_min + x2_extend
-                    else:
-                        x0 = x
+        # Calculate the global bounds of context and target set.
+        coord_min = {}
+        coord_max = {}
+        coord_min["x1"], coord_max["x1"], coord_min["x2"], coord_max["x2"] = (
+            self.coord_bounds
+        )
 
-                    # bbox of x1_min, x1_max, x2_min, x2_max per patch
-                    bbox = [y0 - x1_extend, y0, x0 - x2_extend, x0]
-                    patch_list.append(bbox)
+        ## start with first patch top left hand corner at coord_min["x1"], coord_min["x2"]
+        patch_list = []
 
-        elif (
-            self.coord_directions["x1"] == True and self.coord_directions["x2"] == False
+        # define some lambda functions for use below
+        # round to 12 figures to avoid floating point error but reduce likelihood of unintentional rounding
+        r = lambda x: round(x, 12)
+        bbox_coords_ascend = lambda a, b: [r(a), r(a + b)]
+        bbox_coords_descend = lambda a, b: bbox_coords_ascend(a, b)[::-1]
+
+        compare = {}
+        bbox_coords = {}
+        # for each coordinate direction specify the correct operations for patching
+        for c in ("x1", "x2"):
+            if self.coord_directions[c]:
+                compare[c] = operator.gt
+                bbox_coords[c] = bbox_coords_ascend
+            else:
+                step[c] = -step[c]
+                coord_min[c], coord_max[c] = coord_max[c], coord_min[c]
+                size[c] = -size[c]
+                compare[c] = operator.lt
+                bbox_coords[c] = bbox_coords_descend
+
+        # Define the bounding boxes for all patches, starting in top left corner of dataArray
+        for y, x in itertools.product(
+            np.arange(coord_min["x1"], coord_max["x1"], step["x1"]),
+            np.arange(coord_min["x2"], coord_max["x2"], step["x2"]),
         ):
-            for y in np.arange(x1_min, x1_max, dy):
-                for x in np.arange(x2_max, x2_min, -dx):
-                    if y + x1_extend > x1_max:
-                        y0 = x1_max - x1_extend
-                    else:
-                        y0 = y
-                    if x - x2_extend < x2_min:
-                        x0 = x2_min + x2_extend
-                    else:
-                        x0 = x
-
-                    # bbox of x1_min, x1_max, x2_min, x2_max per patch
-                    bbox = [y0, y0 + x1_extend, x0 - x2_extend, x0]
-                    patch_list.append(bbox)
-        else:
-            for y in np.arange(x1_min, x1_max, dy):
-                for x in np.arange(x2_min, x2_max, dx):
-                    if y + x1_extend > x1_max:
-                        y0 = x1_max - x1_extend
-                    else:
-                        y0 = y
-                    if x + x2_extend > x2_max:
-                        x0 = x2_max - x2_extend
-                    else:
-                        x0 = x
-
-                    # bbox of x1_min, x1_max, x2_min, x2_max per patch
-                    bbox = [y0, y0 + x1_extend, x0, x0 + x2_extend]
+            y0 = (
+                coord_max["x1"] - size["x1"]
+                if compare["x1"](y + size["x1"], coord_max["x1"])
+                else y
+            )
+            x0 = (
+                coord_max["x2"] - size["x2"]
+                if compare["x2"](x + size["x2"], coord_max["x2"])
+                else x
+            )
 
-                    patch_list.append(bbox)
+            # bbox of x1_min, x1_max, x2_min, x2_max per patch
+            bbox = bbox_coords["x1"](y0, size["x1"]) + bbox_coords["x2"](x0, size["x2"])
+            patch_list.append(bbox)
 
         # Remove duplicate patches while preserving order
         seen = set()
@@ -1551,9 +1542,9 @@ def __call__(
             ]
         ] = None,
         split_frac: float = 0.5,
-        patch_size: Union[float, tuple[float]] = None,
+        patch_size: Union[float, Tuple[float]] = None,
         patch_strategy: Optional[str] = None,
-        stride: Union[float, tuple[float]] = None,
+        stride: Union[float, Tuple[float]] = None,
         num_samples_per_date: int = 1,
         datewise_deterministic: bool = False,
         seed_override: Optional[int] = None,